Skip to content

Allow DataSourceConfig to represent a plain database table - #73273

Merged
pankajastro merged 1 commit into
apache:mainfrom
pankajastro:fix-datafusion-datasource-afl-214
Sep 19, 2026
Merged

pankajastro merged 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
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
@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 thread providers/common/sql/src/airflow/providers/common/sql/config.py Outdated
Comment thread providers/common/ai/docs/operators/llm_schema_compare.rst Outdated
@pankajastro
pankajastro force-pushed the fix-datafusion-datasource-afl-214 branch from 6dd30df to c2a634a Compare September 18, 2026 13:15

@kaxil kaxil left a comment

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.

Verified the last push independently rather than off the diff alone. Extracted config.py at apache/main and at this head by SHA and exec'd both side by side: the storage_type=LOCAL/S3-with-no-uri case that regressed in the previous push now matches main again (ok, not a raise), and the format-with-no-uri case gets the new, reachable message instead of the old Unsupported storage type for URI: on an empty string. Every other delta from main is an intended part of this PR: the bug it fixes, and the justified iceberg blank-table_name tightening backed by its own test.

Full common.sql + common.ai unit suites pass (2086 passed, 7 skipped, mysql-only), and mypy is clean on config.py and the example DAG. I could not get a clean build-docs common.ai in my session, but that's an environment artifact, not this PR's: this branch's merge-base predates main's DuckDB provider, so the container's provider-discovery metadata references a module this checkout doesn't have. Round 2's docs build was clean on nearly the same RST content, and this round's diff to it is wording only.

Also checked #73287 directly: it guards DataFusionEngine.register_datasource itself, which covers the three previously-ungated call sites at their one choke point rather than patching each individually, with a test for the plain-DB case.

LLMSchemaCompareOperator accepts a DataSourceConfig with only conn_id
and table_name, introspected via DbApiHook instead of DataFusion when
the connection resolves to one, but construction failed for that
shape. Also hoist the table_name check so it can no longer be skipped
by an explicit storage_type, and require a uri whenever format is set,
since format alone can't be registered without one.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@kaxil
kaxil force-pushed the fix-datafusion-datasource-afl-214 branch from c2a634a to 5faa8f6 Compare September 18, 2026 20:38
@pankajastro
pankajastro merged commit 3c19778 into apache:main Sep 19, 2026
106 checks passed
@pankajastro
pankajastro deleted the fix-datafusion-datasource-afl-214 branch September 19, 2026 05:00
potiuk pushed a commit that referenced this pull request Sep 20, 2026
…73287)

Follow-up per [review feedback](#73273 (review)) on #73273.

`LLMSQLOperator`, `DataFusionToolset`, and `AnalyticsOperator` hand `DataFusionEngine.register_datasource` whatever `DataSourceConfig` the caller supplies, unlike `LLMSchemaCompareOperator` which checks the connection kind first. A plain-database-table config reaching any of them failed deep inside credential resolution with a confusing `Unknown connection type` error instead of a clear message. `register_datasource` now rejects that shape immediately.
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