diff --git a/ingestion/src/metadata/utils/source_connections.py b/ingestion/src/metadata/utils/source_connections.py index f12c6a6b58f..ad78df2cd5a 100644 --- a/ingestion/src/metadata/utils/source_connections.py +++ b/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 diff --git a/ingestion/tests/unit/test_source_connection.py b/ingestion/tests/unit/test_source_connection.py index 5f1587cd2fc..0b0712caaa8 100644 --- a/ingestion/tests/unit/test_source_connection.py +++ b/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(