Allow DataSourceConfig to represent a plain database table - #73273
pankajastro wants to merge 1 commit into
Conversation
|
Ran this end to end against a real Postgres backend in breeze: two real tables, a real Airflow connection, and each consumer of a plain-database The feature itself works. That path was unreachable before, since the config could not be constructed, so the Where it needs another look is
Row three is a behaviour change against current main, which raises there today. Row two costs a worse error at task runtime: With the check hoisted, the full On docs: Separately, and for a follow-up rather than this PR: three of the four |
LLMSchemaCompareOperator accepts a DataSourceConfig with only conn_id and table_name, introspected via DbApiHook instead of DataFusion, but construction failed for that shape. Also close a related gap in the same validation: an object-store DataSourceConfig with storage_type set but no uri. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
03b26fa to
6dd30df
Compare
|
Pushed an update: hoisted the Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting |
| if not self.format and not self.uri and self.storage_type is None: | ||
| # Plain database table: no object store involved, so storage_type stays unset. | ||
| return | ||
|
|
||
| if self.storage_type is None: | ||
| self.storage_type = self._extract_storage_type | ||
|
|
||
| if self.storage_type is not None and (not self.table_name or not self.table_name.strip()): | ||
| raise ValueError("Table name must be provided for storage type") | ||
| if not self.uri: | ||
| raise ValueError("URI must be provided when storage_type is set") |
There was a problem hiding this comment.
This is the one part that went past the round-1 suggestion, and it rejects a shape the released provider accepts. Measured in breeze against providers-common-sql/2.1.1, whose __post_init__ is byte-identical to this PR's base:
case | 2.1.1 (released) | HEAD 6dd30dfa20
plain + explicit storage_type | ok (storage_type=local) | raise: URI must be provided when storage_type is set
format, no uri | raise: Unsupported storage ... | raise: Unsupported storage type for URI:
DataSourceConfig(conn_id="postgres_default", table_name="customers", storage_type=StorageType.LOCAL) was the only way to get a plain database table past __post_init__ before this fix, because an explicit storage_type skips _extract_storage_type, and it works end to end since the DbApiHook branch never reads uri, format or storage_type (llm_schema_compare.py L242-L250). So the workaround for the bug this PR fixes now fails at Dag import, asking for a URI the user never wanted.
Row two is the other half: the check cannot fire for the case its message describes. format="parquet" with the URI forgotten still raises Unsupported storage type for URI: because _extract_storage_type runs first, which test_format_handlers.py L143-L146 already pins. Both new tests pass storage_type explicitly, so nothing covers the inferred path.
Keying the guard on format and moving it above the inference keeps the old shape constructible and makes the message both reachable and accurate:
| if not self.format and not self.uri and self.storage_type is None: | |
| # Plain database table: no object store involved, so storage_type stays unset. | |
| return | |
| if self.storage_type is None: | |
| self.storage_type = self._extract_storage_type | |
| if self.storage_type is not None and (not self.table_name or not self.table_name.strip()): | |
| raise ValueError("Table name must be provided for storage type") | |
| if not self.uri: | |
| raise ValueError("URI must be provided when storage_type is set") | |
| if not self.format and not self.uri: | |
| # Plain database table: no object store involved, so storage_type stays unset. | |
| return | |
| if not self.uri: | |
| raise ValueError("URI must be provided when format is set") | |
| if self.storage_type is None: | |
| self.storage_type = self._extract_storage_type |
I grepped the call sites before proposing this: 57 DataSourceConfig(...) constructions in providers/, 9 outside tests, none outside providers/, and none of them passes an explicit storage_type or sets format without a uri. So the only fallout is three test edits. test_explicit_storage_type_without_uri_raises_error becomes an accepted case, and test_explicit_storage_type_without_uri_raises_error_with_format plus test_format_handlers.py L145 move to the new message.
One more delta from the same hoist, not covered above. Putting the table_name check above the is_table_provider branch also rejects DataSourceConfig(conn_id="c", table_name="", format="iceberg", db_name="default"), which main accepts today. Registering an iceberg table under an empty name is broken downstream regardless, so tightening it looks right, but it is a third behaviour change against main and reads as incidental rather than chosen. Measured the same way as the table above, main returns ok(storage_type=None) and this head raises Table name must be provided for storage type.
| object-storage sources (S3 Parquet, CSV, Iceberg, etc.) in the comparison. | ||
| These can be freely combined with ``db_conn_ids``: | ||
| These can be freely combined with ``db_conn_ids``. A ``DataSourceConfig`` | ||
| with neither ``uri`` nor ``format`` set is introspected via ``DbApiHook`` |
There was a problem hiding this comment.
The branch is picked by the connection, not by the config's fields: _introspect_datasource_schema L242-L252 reads _is_dbapi_connection(ds_config.conn_id) and never looks at uri or format, which is what the operator's own class docstring says at L87-L93. Two things then go wrong for a reader following this sentence. A config that does set uri and format, on a conn that resolves to a DbApiHook, also takes the hook path and silently ignores both fields, which "instead of DataFusion" says cannot happen. And a config with neither, on a conn that is not a DbApiHook, does not take the hook path at all: _is_dbapi_connection swallows the failure at DEBUG level (L183-L185), so a missing provider or a typo'd conn id ends up failing inside DataFusion. Restating the docstring's rule here and in the data_sources bullet at L191 covers both, and the With Object Storage heading at L60 could use a word too now that it introduces a non-object-storage shape.
DataSourceConfigalways inferredstorage_typefromuri, which defaults to empty.LLMSchemaCompareOperatorbuilds aDataSourceConfigwith onlyconn_idandtable_name(nouri, noformat) to introspect a plain database connection viaDbApiHook, and that construction failed.Reproduction (local
airflow dags testrun):plain_db_tablewould failstorage_typestaysNone.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Sonnet 5) following the guidelines