Skip to content

Commit

Permalink
fix sqlalchemy
Browse files Browse the repository at this point in the history
  • Loading branch information
sh-rp committed Sep 20, 2024
1 parent 9a1752d commit a77192f
Show file tree
Hide file tree
Showing 7 changed files with 12 additions and 8 deletions.
6 changes: 4 additions & 2 deletions dlt/destinations/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,15 @@ def __init__(
self.client = client
self.query = query
self.table = table
self.columns = columns
self.schema_columns = columns

@contextmanager
def cursor(self) -> Generator[SupportsReadableRelation, Any, Any]:
"""Gets a DBApiCursor for the current relation"""
if self.table:
with self.client.table_relation(table=self.table, columns=self.columns) as cursor:
with self.client.table_relation(
table=self.table, columns=self.schema_columns
) as cursor:
yield cursor
elif self.query:
with self.client.query_relation(query=self.query) as cursor:
Expand Down
2 changes: 1 addition & 1 deletion dlt/destinations/impl/filesystem/filesystem.py
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,7 @@ def table_relation(
) -> Generator[DBApiCursor, Any, Any]:
table = self.sql_client.make_qualified_table_name(table)
with self.sql_client.execute_query(f"SELECT * FROM {table}") as cursor:
cursor.columns = columns
cursor.schema_columns = columns
yield cursor

@contextmanager
Expand Down
2 changes: 1 addition & 1 deletion dlt/destinations/impl/mssql/mssql.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ def table_relation(
query = f"SELECT * FROM {table}"
sql_client._conn.autocommit = False
with sql_client.execute_query(query) as cursor:
cursor.columns = columns
cursor.schema_columns = columns
yield cursor

@contextmanager
Expand Down
2 changes: 2 additions & 0 deletions dlt/destinations/impl/sqlalchemy/db_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,8 @@ def _wrap(self: "SqlalchemyClient", *args: Any, **kwargs: Any) -> Any:

class SqlaDbApiCursor(DBApiCursorImpl):
def __init__(self, curr: sa.engine.CursorResult) -> None:
self.schema_columns = None

# Sqlalchemy CursorResult is *mostly* compatible with DB-API cursor
self.native_cursor = curr # type: ignore[assignment]
curr.columns
Expand Down
2 changes: 1 addition & 1 deletion dlt/destinations/job_client_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -721,7 +721,7 @@ def table_relation(
table = sql_client.make_qualified_table_name(table)
query = f"SELECT * FROM {table}"
with sql_client.execute_query(query) as cursor:
cursor.columns = columns
cursor.schema_columns = columns
yield cursor

@contextmanager
Expand Down
4 changes: 2 additions & 2 deletions dlt/destinations/sql_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ class DBApiCursorImpl(DBApiCursor):

def __init__(self, curr: DBApiCursor) -> None:
self.native_cursor = curr
self.columns = None
self.schema_columns = None

# wire protocol methods
self.execute = curr.execute # type: ignore
Expand Down Expand Up @@ -381,7 +381,7 @@ def iter_arrow(self, chunk_size: int) -> Generator[ArrowTable, None, None]:
caps = DestinationCapabilitiesContext.generic_capabilities()

# provide default columns in case not known
columns = self.columns or cast(
columns = self.schema_columns or cast(
TTableSchemaColumns, {c: {"name": c, "nullable": True} for c in self._get_columns()}
)

Expand Down
2 changes: 1 addition & 1 deletion tests/load/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,7 @@ def destinations_configs(
supports_merge=False,
supports_dbt=False,
destination_name="sqlalchemy_sqlite",
credentials="sqlite:///_storage/dl_data.sqlite"
credentials="sqlite:///_storage/dl_data.sqlite",
),
]

Expand Down

0 comments on commit a77192f

Please sign in to comment.