Fixed #4769: Trino Connection Argument Fix (#4794)

* Fixed #4769: Trino Connection Argument Fix

* Added tests

* Imporved test with proxy and without conn args

ingestion/src/metadata/utils/source_connections.py

@@ -219,11 +219,13 @@ def _(connection: TrinoConnection):
         session = Session()
         session.proxies = connection.proxies
         if connection.connectionArguments:
-            return {**connection.connectionArguments, "http_session": session}
+            connection_args = connection.connectionArguments.dict()
+            connection_args.update({"http_session": session})
+            return connection_args
         else:
             return {"http_session": session}
     else:
-        return connection.connectionArguments
+        return connection.connectionArguments if connection.connectionArguments else {}
 
 
 @get_connection_url.register

ingestion/tests/unit/test_source_connection.py

@@ -87,6 +87,67 @@ class SouceConnectionTest(TestCase):
         )
         assert expected_url == get_connection_url(trino_conn_obj)
 
+    def test_trino_conn_arguments(self):
+        # connection arguments without connectionArguments and without proxies
+        expected_args = {}
+        trino_conn_obj = TrinoConnection(
+            username="user",
+            password=None,
+            hostPort="localhost:443",
+            catalog="tpcds",
+            database="tiny",
+            connectionArguments=None,
+            scheme=TrinoScheme.trino,
+        )
+        assert expected_args == get_connection_args(trino_conn_obj)
+
+        # connection arguments with connectionArguments and without proxies
+        expected_args = {"user": "user-to-be-impersonated"}
+        trino_conn_obj = TrinoConnection(
+            username="user",
+            password=None,
+            hostPort="localhost:443",
+            catalog="tpcds",
+            database="tiny",
+            connectionArguments={"user": "user-to-be-impersonated"},
+            scheme=TrinoScheme.trino,
+        )
+        assert expected_args == get_connection_args(trino_conn_obj)
+
+        # connection arguments without connectionArguments and with proxies
+        expected_args = {}
+        trino_conn_obj = TrinoConnection(
+            username="user",
+            password=None,
+            hostPort="localhost:443",
+            catalog="tpcds",
+            database="tiny",
+            connectionArguments=None,
+            proxies={"http": "foo.bar:3128", "http://host.name": "foo.bar:4012"},
+            scheme=TrinoScheme.trino,
+        )
+        conn_args = get_connection_args(trino_conn_obj)
+        assert "http_session" in conn_args
+        conn_args.pop("http_session")
+        assert expected_args == conn_args
+
+        # connection arguments with connectionArguments and with proxies
+        expected_args = {"user": "user-to-be-impersonated"}
+        trino_conn_obj = TrinoConnection(
+            username="user",
+            password=None,
+            hostPort="localhost:443",
+            catalog="tpcds",
+            database="tiny",
+            connectionArguments={"user": "user-to-be-impersonated"},
+            proxies={"http": "foo.bar:3128", "http://host.name": "foo.bar:4012"},
+            scheme=TrinoScheme.trino,
+        )
+        conn_args = get_connection_args(trino_conn_obj)
+        assert "http_session" in conn_args
+        conn_args.pop("http_session")
+        assert expected_args == conn_args
+
     def test_trino_url_with_params(self):
         expected_url = "trino://username:pass@localhost:443/catalog?param=value"
         trino_conn_obj = TrinoConnection(