kaxil commented on code in PR #74228:
URL: https://github.com/apache/airflow/pull/74228#discussion_r4243880529
##########
providers/common/sql/docs/operators.rst:
##########
@@ -365,8 +365,12 @@ resolved in this order:
Azure Storage
-------------
-Use an ``az://`` URI with a ``conn_id`` pointing to a ``wasb`` connection.
-``abfs://`` and ``abfss://`` URIs are not recognized yet. The account name
+Use an ``az://``, ``abfs://``, or ``abfss://`` URI with a ``conn_id`` pointing
+to a ``wasb`` connection. ``az://container/path`` takes the storage account
Review Comment:
Small one: `az://` doesn't take the account only from the connection. When
the connection has no host or login, the binding falls back to
`AZURE_STORAGE_ACCOUNT_NAME` on the worker. It would also help to say that an
abfs(s) URI's account is used when the connection names none, since that's new
here.
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -102,19 +103,52 @@ def create_object_store(self, path: str,
connection_config: ConnectionConfig | N
if temp_key_path is not None:
Path(temp_key_path).unlink(missing_ok=True)
- def get_scheme(self) -> str:
- """Return the scheme for GCS."""
- return "gs://"
+
+# Host suffix matched loosely (not pinned to dfs.core.windows.net) to also
cover sovereign
+# clouds, consistent with DataFusionEngine._resolve_wasb_account's
AZURE_STORAGE_ENDPOINT tolerance.
+_ABFS_URI_RE =
re.compile(r"^abfss?://(?P<container>[^@/]+)@(?P<account>[^./]+)\.[^/]+(?P<path>/.*)?$")
Review Comment:
The host suffix is matched here and then dropped, because `MicrosoftAzure`
has no endpoint kwarg. So `abfss://[email protected]/...`
with an account-less connection injects `acct` and sends the request to
`acct.blob.core.windows.net` in the public cloud. `_resolve_wasb_account`
raises for a non-public host unless `AZURE_STORAGE_ENDPOINT`/`AZURE_ENDPOINT`
is set, but nothing checks the URI host, so the "consistent with" comment above
doesn't hold. Could the URI host get the same check (a `.core.windows.net`
suffix, or the env var set), with a test for a sovereign-cloud URI?
##########
providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py:
##########
@@ -230,8 +238,127 @@ def test_execute_query_with_local_csv(self,
mock_get_conn):
result = engine.execute_query("SELECT * FROM test_csv ORDER BY
name")
- expected = {"name": ["Alice", "Bob"], "age": [30, 25]}
- assert result == expected
+ assert result == {"name": ["Alice", "Bob"], "age": [30, 25]}
+ finally:
+ os.unlink(csv_path)
+
+ @patch.object(DataFusionEngine, "_get_connection_config")
+ def test_execute_query_with_abfs_uri_matches_real_registry(self,
mock_get_conn):
+ """Only object-store construction is mocked; normalize_uri and
DataFusion's real
+ registry run unmocked, to catch "No suitable object store found" at
query time."""
+ mock_get_conn.return_value = None
+
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv",
delete=False) as f:
+ f.write("name,age\nAlice,30\nBob,25\n")
+ csv_path = f.name
+
+ try:
+ engine = DataFusionEngine()
+ datasource_config = DataSourceConfig(
+ table_name="abfs_csv",
+
uri=f"abfss://[email protected]{csv_path}",
+ format="csv",
+ conn_id="wasb_default",
+ )
+
+ with patch(
+
"airflow.providers.common.sql.datafusion.object_storage_provider"
+ ".AzureObjectStorageProvider.create_object_store",
+ autospec=True,
+ return_value=LocalFileSystem(),
+ ):
+ engine.register_datasource(datasource_config)
+
+ result = engine.execute_query("SELECT * FROM abfs_csv ORDER BY
name")
+
+ assert result == {"name": ["Alice", "Bob"], "age": [30, 25]}
+ # The caller's own config is untouched; the canonical URI only
shows up internally.
+ assert datasource_config.uri ==
f"abfss://[email protected]{csv_path}"
+ assert engine.registered_tables["abfs_csv"] ==
f"az://myaccount.demo-container{csv_path}"
+ finally:
+ os.unlink(csv_path)
+
+ @pytest.mark.parametrize(
+ ("uri_a", "uri_b"),
+ [
+ pytest.param(
+ "abfss://[email protected]/data.csv",
+ "abfss://[email protected]/data.csv",
+ id="same-account-different-containers",
+ ),
+ pytest.param(
+ "abfss://[email protected]/data.csv",
+ "abfss://[email protected]/data.csv",
+ id="same-container-different-accounts",
+ ),
+ ],
+ )
+ @patch.object(DataFusionEngine, "_get_connection_config")
+ def test_execute_query_with_colliding_abfs_identifiers_does_not_collide(
+ self, mock_get_conn, uri_a, uri_b
+ ):
+ """Two ways an abfs(s) URI's identifying info could collide in
DataFusion's (schema,
+ host) registry: two containers on one account, or one container name
on two accounts."""
+ mock_get_conn.return_value = None
+
+ with (
+ tempfile.TemporaryDirectory() as dir_a,
+ tempfile.TemporaryDirectory() as dir_b,
+ ):
+ (Path(dir_a) / "data.csv").write_text("name,age\nAlice,30\n")
+ (Path(dir_b) / "data.csv").write_text("name,age\nBob,25\n")
+
+ engine = DataFusionEngine()
+ config_a = DataSourceConfig(table_name="table_a", uri=uri_a,
format="csv", conn_id="wasb_a")
+ config_b = DataSourceConfig(table_name="table_b", uri=uri_b,
format="csv", conn_id="wasb_b")
+
+ with patch(
+
"airflow.providers.common.sql.datafusion.object_storage_provider"
+ ".AzureObjectStorageProvider.create_object_store",
+ autospec=True,
+ side_effect=[LocalFileSystem(prefix=dir_a),
LocalFileSystem(prefix=dir_b)],
+ ):
+ engine.register_datasource(config_a)
+ engine.register_datasource(config_b)
+
+ assert engine.execute_query("SELECT * FROM table_a") == {"name":
["Alice"], "age": [30]}
+ assert engine.execute_query("SELECT * FROM table_b") == {"name":
["Bob"], "age": [25]}
+
+ @patch.object(DataFusionEngine, "_get_connection_config")
+ def test_register_datasource_does_not_mutate_caller_config(self,
mock_get_conn):
+ """
+ Regression test: _register_object_store used to overwrite
datasource_config.uri in
Review Comment:
The in-place overwrite this docstring describes isn't in any pushed commit,
so a future reader can't find the history it refers to. With
`create_object_store` patched, `get_bucket` never runs, so the second-engine
queries can't catch the container mis-extraction either. The
`datasource_config.uri == original_uri` check is the part that matters, and
line 276 already asserts it. I'd drop this test, or patch `MicrosoftAzure`
instead and assert `container_name="demo-container"` on both calls.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
