Sql driver

SqlDriver

Bases: BaseSqlDriver

Source code in griptape/griptape/drivers/sql/sql_driver.py

@define
class SqlDriver(BaseSqlDriver):
    engine_url: str = field(kw_only=True)
    create_engine_params: dict = field(factory=dict, kw_only=True)
    engine: Engine = field(init=False)

    def __attrs_post_init__(self) -> None:
        sqlalchemy = import_optional_dependency("sqlalchemy")

        self.engine = sqlalchemy.create_engine(self.engine_url, **self.create_engine_params)

    def execute_query(self, query: str) -> Optional[list[BaseSqlDriver.RowResult]]:
        rows = self.execute_query_raw(query)

        if rows:
            return [BaseSqlDriver.RowResult(row) for row in rows]
        else:
            return None

    def execute_query_raw(self, query: str) -> Optional[list[dict[str, Any]]]:
        sqlalchemy = import_optional_dependency("sqlalchemy")

        with self.engine.begin() as con:
            results = con.execute(sqlalchemy.text(query))

            if results.returns_rows:
                return [{column: value for column, value in result.items()} for result in results]
            else:
                return None

    def get_table_schema(self, table: str, schema: Optional[str] = None) -> Optional[str]:
        sqlalchemy = import_optional_dependency("sqlalchemy")

        try:
            table = sqlalchemy.Table(
                table, sqlalchemy.MetaData(bind=self.engine), schema=schema, autoload=True, autoload_with=self.engine
            )
            return str([(c.name, c.type) for c in table.columns])
        except sqlalchemy.exc.NoSuchTableError:
            return None

create_engine_params: dict = field(factory=dict, kw_only=True) class-attribute instance-attribute

engine: Engine = field(init=False) class-attribute instance-attribute

engine_url: str = field(kw_only=True) class-attribute instance-attribute

__attrs_post_init__()

Source code in griptape/griptape/drivers/sql/sql_driver.py

def __attrs_post_init__(self) -> None:
    sqlalchemy = import_optional_dependency("sqlalchemy")

    self.engine = sqlalchemy.create_engine(self.engine_url, **self.create_engine_params)

execute_query(query)

Source code in griptape/griptape/drivers/sql/sql_driver.py

def execute_query(self, query: str) -> Optional[list[BaseSqlDriver.RowResult]]:
    rows = self.execute_query_raw(query)

    if rows:
        return [BaseSqlDriver.RowResult(row) for row in rows]
    else:
        return None

execute_query_raw(query)

Source code in griptape/griptape/drivers/sql/sql_driver.py

def execute_query_raw(self, query: str) -> Optional[list[dict[str, Any]]]:
    sqlalchemy = import_optional_dependency("sqlalchemy")

    with self.engine.begin() as con:
        results = con.execute(sqlalchemy.text(query))

        if results.returns_rows:
            return [{column: value for column, value in result.items()} for result in results]
        else:
            return None

get_table_schema(table, schema=None)

Source code in griptape/griptape/drivers/sql/sql_driver.py

def get_table_schema(self, table: str, schema: Optional[str] = None) -> Optional[str]:
    sqlalchemy = import_optional_dependency("sqlalchemy")

    try:
        table = sqlalchemy.Table(
            table, sqlalchemy.MetaData(bind=self.engine), schema=schema, autoload=True, autoload_with=self.engine
        )
        return str([(c.name, c.type) for c in table.columns])
    except sqlalchemy.exc.NoSuchTableError:
        return None