Skip to content

Allow DataSourceConfig to represent a plain database table - #73273

Open
pankajastro wants to merge 1 commit into
apache:mainfrom
pankajastro:fix-datafusion-datasource-afl-214
Open

pankajastro wants to merge 1 commit into
apache:mainfrom
pankajastro:fix-datafusion-datasource-afl-214

Conversation

@pankajastro

Copy link
Copy Markdown
Member

DataSourceConfig always inferred storage_type from uri, which defaults to empty. LLMSchemaCompareOperator builds a DataSourceConfig with only conn_id and table_name (no uri, no format) to introspect a plain database connection via DbApiHook, and that construction failed.

Reproduction (local airflow dags test run):

  • Before the fix task plain_db_table would fail
@task
def plain_db_table():
    config = DataSourceConfig(conn_id="postgres_default", table_name="my_table")
  • After the fix: the same task succeeds, storage_type stays None.

Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Sonnet 5)

Generated-by: Claude Code (Sonnet 5) following the guidelines

Comment thread providers/common/sql/src/airflow/providers/common/sql/config.py Outdated
Comment thread providers/common/sql/src/airflow/providers/common/sql/config.py
Comment thread providers/common/sql/src/airflow/providers/common/sql/config.py
@kaxil

kaxil commented Sep 17, 2026

Copy link
Copy Markdown
Member

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 DataSourceConfig.

The feature itself works. LLMSchemaCompareOperator now introspects a plain table through DbApiHook and returns real schema text:

Source: pg_e2e (postgresql)
Table: e2e_orders_v1
Columns: order_id INTEGER, customer VARCHAR(64), amount NUMERIC(10, 2), created_at TIMESTAMP
Index: e2e_orders_v1_customer_idx (customer)

That path was unreachable before, since the config could not be constructed, so the DbApiHook branch in _introspect_datasource_schema was dead code. Worth correcting in the description though: the operator does not build the config, the caller passes it in data_sources, and nothing in the repo constructs a plain-database one yet.

Where it needs another look is table_name validation. Same script, before and after moving the check above the new early return:

case before after
plain DB, valid table_name OK, storage_type=None OK, storage_type=None
plain DB, blank table_name accepted raises
plain DB, whitespace table_name accepted raises
explicit storage_type=S3, blank table_name, no uri accepted raises
uri set, blank table_name raises raises

Row three is a behaviour change against current main, which raises there today. Row two costs a worse error at task runtime: hook.get_table_schema("") surfaces a bare NoSuchTableError with an empty message, where the check would have said Table name must be provided at construction. I left a committable suggestion on the diff for it.

With the check hoisted, the full common.sql and common.ai unit suites pass (2084 tests), and both existing assertions on that message, in test_config.py and test_format_handlers.py, keep working.

On docs: llm_schema_compare.rst still scopes data_sources to object-storage and catalog-managed sources, so the new shape has no example anywhere. I have a short section written that frames it as the way to compare differently named tables, since db_conn_ids and table_names form a cross-product and cannot pair orders with orders_v2. Happy to hand it over for this PR or push it as a follow-up, whichever you prefer.

Separately, and for a follow-up rather than this PR: three of the four register_datasource callers do not gate on connection kind, so a plain-database config reaching LLMSQLOperator, DataFusionToolset or AnalyticsOperator fails with ValueError: Unknown connection type postgres from _get_credentials. Only LLMSchemaCompareOperator checks _is_dbapi_connection first.

Comment thread providers/common/sql/src/airflow/providers/common/sql/config.py Outdated
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>
@pankajastro
pankajastro force-pushed the fix-datafusion-datasource-afl-214 branch from 03b26fa to 6dd30df Compare September 17, 2026 12:58
@pankajastro

Copy link
Copy Markdown
Member Author

Pushed an update: hoisted the table_name check (also catches the iceberg-blank-table_name case), added a check requiring uri when storage_type is set, and updated the docs/example DAG for the plain-DB shape. Left the register_datasource guard for a separate follow-up PR. Thanks for the thorough verification!


Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting

Comment on lines +102 to +110
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")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Suggested change
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``

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants