Fixes #30173: fix(ingestion): security/metadata/mcp sources own their BaseConnection; dedupe close-on-failure#30174
Fixes #30173: fix(ingestion): security/metadata/mcp sources own their BaseConnection; dedupe close-on-failure#30174IceS2 wants to merge 7 commits into
Conversation
…ection SecurityServiceSource built two clients on start-up: one through the owner-discarding `get_connection` seam and a second one through the `get_client` hook. It also discarded the test-connection result, so a failing connection never raised and ingestion carried on against a broken client. The base now owns a single `BaseConnection`, reuses it for the test step through `run_test_connection` (which raises), and disposes it on `close()`. The metadata (atlas, amundsen, alationsink) and mcp sources have no service base but carry the same double-connection pattern in their own `__init__`: they built a client through the seam and then had `test_connection_common` build a second one. Each now owns its connection and tests through it. Pre-existing bugs fixed: - AmundsenConnection never registered the Neo4j driver teardown, so every temporary owner (test connection, automations) leaked a driver. - SecurityServiceSource.get_client was undeclared on the base.
Every service base re-derived `try: self.test_connection() except: <dispose>` in its constructor, and the shape had drifted: most called `self.close()`, pipeline disposed `self._connection` directly. `self.close()` is the wrong teardown for a constructor failure. It is successful-run teardown (it runs `compute_percentile`, `self.metadata.close()`) and it is polymorphic, so it dispatches to subclass overrides that touch attributes set *after* `super().__init__()` returns - airflow's `self.session`, databrickspipeline's `self._table_lookup_cache` - and the resulting AttributeError masks the connection error the user needs to see. saperp's override goes further and closes the OpenMetadata client. `close_on_failure` releases only the owner, which is the only resource in flight when the test step runs. Applied to the 10 service bases and to the leaf sources that own a connection directly. CommonDbSourceService keeps an explicit failure branch: its SSL temp files belong to the source, not the connection owner, until they move into a BaseConnection hook. Also fixed while here: - airflow and openlineage `close()` never called `super().close()`, so the owner was never disposed. airflow closes its session before the engine dispose that `super().close()` triggers. - OpenLineageConnection never registered the Kafka consumer teardown and SftpConnection never registered the SFTP transport teardown, so every temporary owner (test connection, automations) leaked one.
✅ PR checks passedThe linked issue has a description and all required Shipping project fields set. Thanks! |
…ailure - CommonDbSourceService: a connector with no connection_class gets its engine from the legacy seam, so nobody owns it. close_on_failure only releases the owner, which left that engine and its session leaked on a failing test step. The failure branch now releases it (gitar-bot). - close_on_failure: a teardown that raises replaced the verification error that triggered it. The teardown error is logged instead; the caller gets the error that matters (greptile). - OpenLineageConnection: register the Kinesis client teardown too. Only the Kafka branch had one, so Kinesis runs dropped the client without releasing it (greptile). - SecurityServiceSource: drop the `get_client` hook. It duplicated the seam -- for the only implementation, `get_client()` and `get_connection()` build the same client the same way -- and the seam already covers connectors without a connection_class. - atlas, amundsen, alationsink, mcp: these ship a connection_class, so create_connection never returns None. Drop the unreachable legacy fallback and take the client from the owner, as the other migrated leaf sources do.
|
Review addressed in gitar: legacy-path engine leaked on test-connection failure — correct, and it was a regression I introduced. greptile: cleanup masks connection failure — valid. Not a regression (the old greptile: Kinesis client remains unowned — valid. Registering the teardown for the Kafka branch and not its sibling was an inconsistency on my side. Both branches now register it. Two more changes from self-review:
Copilot could not review (quota). |
🟡 Playwright Results — all passed (24 flaky)✅ 4547 passed · ❌ 0 failed · 🟡 24 flaky · ⏭️ 95 skipped
🟡 24 flaky test(s) (passed on retry)
How to debug locally# Download playwright-test-results-<shard> artifact and unzip
npx playwright show-trace path/to/trace.zip # view trace |
Code Review ✅ Approved 1 resolved / 1 findingsCentralizes connection ownership and teardown logic within ✅ 1 resolved✅ Bug: Legacy-path engine leaked on test-connection failure in CommonDbSource
OptionsDisplay: compact → Showing less information. Comment with these commands to change the behavior for this request:
Was this helpful? React with 👍 / 👎 | Gitar |
|



Fixes #30173
Closes out the OSS side of
source-owns-BaseConnection. Two commits, splittable.1. Security, metadata and mcp sources own their BaseConnection (
62577c2381)SecurityServiceSourcebuilt two clients — one through the owner-discardingget_connectionseam, a second through itsget_clienthook — and then discarded the test-connection result, so a failing connection never raised. The base now owns oneBaseConnection, takes its client from the owner, tests throughrun_test_connection(which raises), and disposes it onclose().get_clientis demoted to the legacy path for connectors without aconnection_class, and is now declared on the base instead of being an implicit contract.Key finding: for the only consumer of this base, the client's constructor performs a real login. Routing just the unused
self.connectionattribute through the owner would have left the second login in place — the fix only works becauseself.clientcomes from the owner too.The
metadata(atlas, amundsen, alationsink) andmcpsources have no service base but carry the same pattern in their own__init__. Each now owns its connection and tests through it.drives/(plural) is a dead package containing only__init__.py— nothing to migrate.2. Dedupe close-on-failure behind a context manager (
c4fcc7319b)close_on_failureinsource/connections.pyreplaces the driftedtry: self.test_connection() except: <dispose>; raisein 24 constructors.self.close()was the wrong teardown for a constructor failure: it is successful-run teardown, and it is polymorphic, so it dispatched to subclass overrides touching attributes set aftersuper().__init__()returned — airflow'sself.session, databrickspipeline'sself._table_lookup_cache— masking the real connection error behind anAttributeError.saperp's override closes the OpenMetadata client outright. The context manager releases only the owner, which is the only resource in flight when the test step runs.CommonDbSourceServicekeeps an explicit failure branch: its SSL temp files are owned by the source, not the connection, until they move into aBaseConnectionhook.Pre-existing bugs fixed while here
airflow,openlineageclose()never calledsuper().close(), so the owner was never disposed. airflow closes its session before the engine disposesuper().close()triggers.OpenLineageConnectionSftpConnectionAmundsenConnectionThe last three leaked one client per temporary owner — i.e. on every Automation test-connection.
Compatibility
Additive.
connection_objis still set,test_connection_commonkeeps its 3-arg signature, and theget_connection/import_connection_fnseam is untouched — Collate and custom connectors are unaffected.Greptile Summary
This PR moves ingestion sources to owned connections and centralizes failure cleanup. The main changes are:
close_on_failurefor constructor verification failures.Confidence Score: 5/5
This looks safe to merge.
Important Files Changed
Reviews (6): Last reviewed commit: "Merge branch 'main' into fix/security-an..." | Re-trigger Greptile
Context used (3)