ashb commented on code in PR #25972:
URL: https://github.com/apache/airflow/pull/25972#discussion_r958616778


##########
tests/models/test_dataset.py:
##########
@@ -32,14 +35,41 @@ def test_uri_with_scheme(self, dag_maker, session):
         with dag_maker(dag_id="example_dataset"):
             EmptyOperator(task_id="task1", outlets=[dataset])
 
-    def test_uri_with_airflow_scheme_restricted(self, dag_maker, session):
-        dataset = Dataset(uri="airflow://example_dataset")
-        with pytest.raises(ValueError, match='Scheme `airflow` is reserved'):
-            with dag_maker(dag_id="example_dataset"):
-                EmptyOperator(task_id="task1", outlets=[dataset])
-
     def test_uri_with_invalid_characters(self, dag_maker, session):
         dataset = Dataset(uri="èxample_datašet")
         with pytest.raises(ValueError, match='URI must be ascii'):
             with dag_maker(dag_id="example_dataset"):
                 EmptyOperator(task_id="task1", outlets=[dataset])
+
+
+class TestDatasetModel:
+    @pytest.mark.parametrize(
+        "conn_uri, dataset_uri, expected_canonical_uri",
+        [
+            ("postgres://somehost/", "airflow://testconn@/", 
"postgres://somehost/"),
+            ("postgres://somehost:111/base", "airflow://testconn@", 
"postgres://somehost:111/base"),
+            ("postgres://somehost:111/base", "airflow+foo://testconn@", 
"foo://somehost:111/base"),
+            ("postgres://somehost:111", "airflow://testconn@foo:222", 
"postgres://foo:222"),
+            (
+                "postgres://somehost:111/base",
+                "airflow://testconn@/extra",
+                "postgres://somehost:111/base/extra",
+            ),
+            ("postgres://somehost:111", "airflow://testconn@/?foo=bar", 
"postgres://somehost:111/?foo=bar"),

Review Comment:
   Oh, this might be a risk. I don't think we should include query params from 
the connection -- there's much to big a risk of including the (old 
style/deprecated) of `aws_secret_access_key` 
https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/connections/aws.html#configuring-the-connection



-- 
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]

Reply via email to