Skip to content

Sql loader

SqlLoader

Bases: BaseLoader

Source code in griptape/loaders/sql_loader.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@define
class SqlLoader(BaseLoader):
    sql_driver: BaseSqlDriver = field(kw_only=True)
    embedding_driver: Optional[BaseEmbeddingDriver] = field(default=None, kw_only=True)

    def load(self, source: str, *args, **kwargs) -> list[CsvRowArtifact]:
        rows = self.sql_driver.execute_query(source)
        artifacts = []

        if rows:
            chunks = [CsvRowArtifact(row.cells) for row in rows]
        else:
            chunks = []

        if self.embedding_driver:
            for chunk in chunks:
                chunk.generate_embedding(self.embedding_driver)

        for chunk in chunks:
            artifacts.append(chunk)

        return artifacts

    def load_collection(self, sources: list[str], *args, **kwargs) -> dict[str, list[CsvRowArtifact]]:
        return cast(dict[str, list[CsvRowArtifact]], super().load_collection(sources, *args, **kwargs))

embedding_driver: Optional[BaseEmbeddingDriver] = field(default=None, kw_only=True) class-attribute instance-attribute

sql_driver: BaseSqlDriver = field(kw_only=True) class-attribute instance-attribute

load(source, *args, **kwargs)

Source code in griptape/loaders/sql_loader.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def load(self, source: str, *args, **kwargs) -> list[CsvRowArtifact]:
    rows = self.sql_driver.execute_query(source)
    artifacts = []

    if rows:
        chunks = [CsvRowArtifact(row.cells) for row in rows]
    else:
        chunks = []

    if self.embedding_driver:
        for chunk in chunks:
            chunk.generate_embedding(self.embedding_driver)

    for chunk in chunks:
        artifacts.append(chunk)

    return artifacts

load_collection(sources, *args, **kwargs)

Source code in griptape/loaders/sql_loader.py 
33
34
def load_collection(self, sources: list[str], *args, **kwargs) -> dict[str, list[CsvRowArtifact]]:
    return cast(dict[str, list[CsvRowArtifact]], super().load_collection(sources, *args, **kwargs))