Fix UnknownTimeZoneError in SnowflakeIOManager with ADBC connector#33922
Fix UnknownTimeZoneError in SnowflakeIOManager with ADBC connector#33922I-am-Uchenna wants to merge 7 commits into
Conversation
Replace SHOW SCHEMAS with information_schema.schemata query to avoid
pytz.exceptions.UnknownTimeZoneError when using ADBC connector.
SHOW SCHEMAS returns TIMESTAMP_LTZ columns that trigger
pytz.timezone('Local') failure when no session timezone is set.
…BLES Updated _table_exists function to query information_schema.tables instead of using SHOW TABLES to prevent UnknownTimeZoneError with ADBC connections.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e79dd01c5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| f"SELECT 1 FROM {table_slice.database}.information_schema.tables" | ||
| f" WHERE LOWER(table_schema) = LOWER('{table_slice.schema}')" | ||
| f" AND LOWER(table_name) = LOWER('{table_slice.table}')" |
There was a problem hiding this comment.
Exclude views when checking for existing tables
When a partitioned asset targets a name that currently exists only as a view, this new information_schema.tables query returns a row because Snowflake's TABLES view includes rows with TABLE_TYPE = 'VIEW'. That makes handle_output take the append/delete path and issue DELETE FROM ... plus an append write against the view, whereas the previous SHOW TABLES check would not match views and would leave the handler in replace mode. Please filter TABLE_TYPE to actual table types so views do not masquerade as existing tables.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good point. Fixed in 527dc6e -- added AND TABLE_TYPE = 'BASE TABLE' filter to the query so views are excluded from the existence check.
| cursor.execute( | ||
| f"show schemas like '{table_slice.schema}' in database {table_slice.database}" | ||
| f"SELECT 1 FROM {table_slice.database}.information_schema.schemata" | ||
| f" WHERE LOWER(schema_name) = LOWER('{table_slice.schema}')" | ||
| ) |
There was a problem hiding this comment.
SQL injection via
table_slice.schema — the value is interpolated directly into the query string without any sanitization. Although schema names typically come from Dagster asset configuration, a schema name containing a single quote (e.g. it's_schema) will break the query with a syntax error, and a maliciously crafted name could manipulate the WHERE clause. Prefer a parameterized query (:1 or %s depending on the cursor API) or at minimum escape/quote the identifier.
| cursor.execute( | |
| f"show schemas like '{table_slice.schema}' in database {table_slice.database}" | |
| f"SELECT 1 FROM {table_slice.database}.information_schema.schemata" | |
| f" WHERE LOWER(schema_name) = LOWER('{table_slice.schema}')" | |
| ) | |
| cursor.execute( | |
| f"SELECT 1 FROM {table_slice.database}.information_schema.schemata" | |
| " WHERE LOWER(schema_name) = LOWER(%s)", | |
| (table_slice.schema,), | |
| ) |
There was a problem hiding this comment.
Good catch. Fixed in 9d61e03 -- both ensure_schema_exists and _table_exists now use %s parameterized queries instead of f-string interpolation.
| cursor.execute( | ||
| f"SHOW TABLES LIKE '{table_slice.table}' IN SCHEMA" | ||
| f" {table_slice.database}.{table_slice.schema}" | ||
| f"SELECT 1 FROM {table_slice.database}.information_schema.tables" | ||
| f" WHERE LOWER(table_schema) = LOWER('{table_slice.schema}')" | ||
| f" AND LOWER(table_name) = LOWER('{table_slice.table}')" | ||
| ) |
There was a problem hiding this comment.
Same SQL injection risk as in
ensure_schema_exists: table_slice.schema and table_slice.table are interpolated directly into the query string. A schema or table name containing a single quote causes a syntax error; a crafted name could manipulate the WHERE clause. Use a parameterized query or escape identifiers properly.
| cursor.execute( | |
| f"SHOW TABLES LIKE '{table_slice.table}' IN SCHEMA" | |
| f" {table_slice.database}.{table_slice.schema}" | |
| f"SELECT 1 FROM {table_slice.database}.information_schema.tables" | |
| f" WHERE LOWER(table_schema) = LOWER('{table_slice.schema}')" | |
| f" AND LOWER(table_name) = LOWER('{table_slice.table}')" | |
| ) | |
| cursor.execute( | |
| f"SELECT 1 FROM {table_slice.database}.information_schema.tables" | |
| " WHERE LOWER(table_schema) = LOWER(%s)" | |
| " AND LOWER(table_name) = LOWER(%s)", | |
| (table_slice.schema, table_slice.table), | |
| ) |
There was a problem hiding this comment.
Fixed in 527dc6e -- switched to %s parameterized queries and also added a TABLE_TYPE = 'BASE TABLE' filter to exclude views from the existence check.
Addresses review feedback: - Replace f-string interpolation with %s parameterized queries to prevent SQL injection in the _table_exists function - Add TABLE_TYPE = 'BASE TABLE' filter to exclude views from the existence check, since information_schema.tables includes views while SHOW TABLES (the previous approach) did not
Replace f-string interpolation with %s parameterized query in ensure_schema_exists to prevent SQL injection, consistent with the _table_exists fix in the polars type handler.
Summary & Motivation
Closes #33750
When using
SnowflakePolarsIOManager(which defaults toconnector = "adbc"), asset writes fail on Windows with:The error originates in
SnowflakeDbClient.ensure_schema_exists, which runsSHOW SCHEMASvia an ADBC cursor. The result set containsTIMESTAMP_LTZcolumns. When no session-levelTIMEZONEhas been set, Snowflake reports the timezone as'Local'in the Arrow schema metadata. PyArrow callspytz.timezone('Local')duringcursor.fetchall(), which raisesUnknownTimeZoneErrorbecause'Local'is not a valid IANA timezone identifier.The native
snowflake-connector-pythonhandles this internally via atzlocalfallback, but the ADBC path delegates timezone resolution to PyArrow + pytz which have no such fallback.A second occurrence of the same pattern exists in
_table_existsin the Polars type handler, which runsSHOW TABLESthrough the same ADBC cursor path.Fix
Replaced
SHOW SCHEMASandSHOW TABLEScommands with equivalentinformation_schemaqueries:ensure_schema_existsinsnowflake_io_manager.py: queriesinformation_schema.schematainstead ofSHOW SCHEMAS_table_existsinsnowflake_polars_type_handler.py: queriesinformation_schema.tablesinstead ofSHOW TABLESBy selecting a constant (
SELECT 1) rather than all columns, the result set contains noTIMESTAMP_LTZvalues. PyArrow never encounters the'Local'timezone string, which eliminates the error regardless of session timezone configuration.This approach:
ALTER SESSION SET TIMEZONE)LOWER()SHOWcommands sinceinformation_schemais accessible to all roles that have privileges on the referenced objectsTest Plan
Added unit tests for
ensure_schema_exists:test_ensure_schema_exists_creates_missing_schema: verifiesinformation_schema.schematais queried and schema is created when missingtest_ensure_schema_exists_skips_create_when_schema_exists: verifies schema creation is skipped when schema already existsChangelog
Fix
UnknownTimeZoneError: 'Local'when usingSnowflakePolarsIOManagerwith ADBC connector on systems where no session timezone is set. ReplacedSHOW SCHEMASandSHOW TABLEScommands withinformation_schemaqueries to avoidTIMESTAMP_LTZcolumns in result sets.