From 00ea858ad84e605100c44b687100e5f9d3b1753f Mon Sep 17 00:00:00 2001 From: Perttu Salonen Date: Mon, 4 Dec 2023 12:17:58 +0200 Subject: [PATCH 1/4] map airflow connection `schema` for vertica `database` Signed-off-by: Perttu Salonen --- cosmos/profiles/vertica/user_pass.py | 3 +-- .../vertica/test_vertica_user_pass.py | 21 ++++++++----------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/cosmos/profiles/vertica/user_pass.py b/cosmos/profiles/vertica/user_pass.py index ccaaf301dc..38fbd5dab4 100644 --- a/cosmos/profiles/vertica/user_pass.py +++ b/cosmos/profiles/vertica/user_pass.py @@ -31,8 +31,7 @@ class VerticaUserPasswordProfileMapping(BaseProfileMapping): "username": "login", "password": "password", "port": "port", - "schema": "schema", - "database": "extra.database", + "database": "schema", "autocommit": "extra.autocommit", "backup_server_node": "extra.backup_server_node", "binary_transfer": "extra.binary_transfer", diff --git a/tests/profiles/vertica/test_vertica_user_pass.py b/tests/profiles/vertica/test_vertica_user_pass.py index 19771c799b..6459dea962 100644 --- a/tests/profiles/vertica/test_vertica_user_pass.py +++ b/tests/profiles/vertica/test_vertica_user_pass.py @@ -23,8 +23,7 @@ def mock_vertica_conn(): # type: ignore login="my_user", password="my_password", port=5433, - schema="my_schema", - extra='{"database": "my_database"}', + schema="my_database", ) with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): @@ -43,8 +42,7 @@ def mock_vertica_conn_custom_port(): # type: ignore login="my_user", password="my_password", port=7472, - schema="my_schema", - extra='{"database": "my_database"}', + schema="my_database", ) with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): @@ -69,8 +67,7 @@ def test_connection_claiming() -> None: "host": "my_host", "login": "my_user", "password": "my_password", - "schema": "my_schema", - "extra": '{"database": "my_database"}', + "schema": "my_database", } # if we're missing any of the values, it shouldn't claim @@ -82,20 +79,20 @@ def test_connection_claiming() -> None: print("testing with", values) with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): - profile_mapping = VerticaUserPasswordProfileMapping(conn) + profile_mapping = VerticaUserPasswordProfileMapping(conn, {"schema": "my_schema"}) assert not profile_mapping.can_claim_connection() - # also test when there's no database + # also test when there's no schema conn = Connection(**potential_values) # type: ignore conn.extra = "" with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): - profile_mapping = VerticaUserPasswordProfileMapping(conn) + profile_mapping = VerticaUserPasswordProfileMapping(conn, {}) assert not profile_mapping.can_claim_connection() # if we have them all, it should claim conn = Connection(**potential_values) # type: ignore with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): - profile_mapping = VerticaUserPasswordProfileMapping(conn) + profile_mapping = VerticaUserPasswordProfileMapping(conn, {"schema": "my_schema"}) assert profile_mapping.can_claim_connection() @@ -107,7 +104,7 @@ def test_profile_mapping_selected( """ profile_mapping = get_automatic_profile_mapping( mock_vertica_conn.conn_id, - {"schema": "my_schema"}, + {"schema": "my_database"}, ) assert isinstance(profile_mapping, VerticaUserPasswordProfileMapping) @@ -145,8 +142,8 @@ def test_profile_args( "username": mock_vertica_conn.login, "password": "{{ env_var('COSMOS_CONN_VERTICA_PASSWORD') }}", "port": mock_vertica_conn.port, + "database": mock_vertica_conn.schema, "schema": "my_schema", - "database": mock_vertica_conn.extra_dejson.get("database"), } From 0a5e4a852e4ea03cd0becc1ea0696c5dfefd2d50 Mon Sep 17 00:00:00 2001 From: Perttu Salonen Date: Thu, 7 Dec 2023 17:28:26 +0200 Subject: [PATCH 2/4] added explanation for the mapping Signed-off-by: Perttu Salonen --- cosmos/profiles/vertica/user_pass.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/cosmos/profiles/vertica/user_pass.py b/cosmos/profiles/vertica/user_pass.py index 38fbd5dab4..0104d947fe 100644 --- a/cosmos/profiles/vertica/user_pass.py +++ b/cosmos/profiles/vertica/user_pass.py @@ -1,4 +1,14 @@ -"Maps Airflow Vertica connections using username + password authentication to dbt profiles." +""" +Maps Airflow Vertica connections using username + password authentication to dbt profiles. + +NOTE: Use Airflow connection `schema` for vertica `database` to keep it consistent with other connection types and profiles. +Also Vertica Airflow provider hook assumes this: +https://github.com/apache/airflow/blob/395ac463494dba1478a05a32900218988495889c/airflow/providers/vertica/hooks/vertica.py#L72 +This seems to be a common approach also for Postgres, Redshift and Exasol since there is no `database` field in Airflow connection +and `schema` is not required for the database connection. +Posgres Hook: +https://github.com/apache/airflow/blob/0953e0f844fa5db81c2b461ec2433de1935260b3/airflow/providers/postgres/hooks/postgres.py#L138 +""" from __future__ import annotations from typing import Any From c9d2565832b9d7519797bf7c0608b3f87c8b28b0 Mon Sep 17 00:00:00 2001 From: Perttu Salonen Date: Fri, 8 Dec 2023 08:10:35 +0200 Subject: [PATCH 3/4] Revert "added explanation for the mapping" This reverts commit 0a5e4a852e4ea03cd0becc1ea0696c5dfefd2d50. --- cosmos/profiles/vertica/user_pass.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/cosmos/profiles/vertica/user_pass.py b/cosmos/profiles/vertica/user_pass.py index 0104d947fe..38fbd5dab4 100644 --- a/cosmos/profiles/vertica/user_pass.py +++ b/cosmos/profiles/vertica/user_pass.py @@ -1,14 +1,4 @@ -""" -Maps Airflow Vertica connections using username + password authentication to dbt profiles. - -NOTE: Use Airflow connection `schema` for vertica `database` to keep it consistent with other connection types and profiles. -Also Vertica Airflow provider hook assumes this: -https://github.com/apache/airflow/blob/395ac463494dba1478a05a32900218988495889c/airflow/providers/vertica/hooks/vertica.py#L72 -This seems to be a common approach also for Postgres, Redshift and Exasol since there is no `database` field in Airflow connection -and `schema` is not required for the database connection. -Posgres Hook: -https://github.com/apache/airflow/blob/0953e0f844fa5db81c2b461ec2433de1935260b3/airflow/providers/postgres/hooks/postgres.py#L138 -""" +"Maps Airflow Vertica connections using username + password authentication to dbt profiles." from __future__ import annotations from typing import Any From f1f82fb64a05d606b41a89d6981e5eb06a730324 Mon Sep 17 00:00:00 2001 From: Perttu Salonen Date: Fri, 8 Dec 2023 08:25:03 +0200 Subject: [PATCH 4/4] move docstring with a sphinx note under the class definition Signed-off-by: Perttu Salonen --- cosmos/profiles/vertica/user_pass.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/cosmos/profiles/vertica/user_pass.py b/cosmos/profiles/vertica/user_pass.py index 38fbd5dab4..e016b612c6 100644 --- a/cosmos/profiles/vertica/user_pass.py +++ b/cosmos/profiles/vertica/user_pass.py @@ -9,8 +9,14 @@ class VerticaUserPasswordProfileMapping(BaseProfileMapping): """ Maps Airflow Vertica connections using username + password authentication to dbt profiles. - https://docs.getdbt.com/reference/warehouse-setups/vertica-setup - https://airflow.apache.org/docs/apache-airflow-providers-vertica/stable/connections/vertica.html + .. note:: + Use Airflow connection ``schema`` for vertica ``database`` to keep it consistent with other connection types and profiles. \ + The Vertica Airflow provider hook `assumes this `_. + This seems to be a common approach also for `Postgres `_, \ + Redshift and Exasol since there is no ``database`` field in Airflow connection and ``schema`` is not required for the database connection. + .. seealso:: + https://docs.getdbt.com/reference/warehouse-setups/vertica-setup + https://airflow.apache.org/docs/apache-airflow-providers-vertica/stable/connections/vertica.html """ airflow_connection_type: str = "vertica"