pankajastro commented on code in PR #74228:
URL: https://github.com/apache/airflow/pull/74228#discussion_r4248326808
##########
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:
Dropped in b63272c, since the uri == original_uri assertion in the abfss
registry test already covers it.
--
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]
