-
Notifications
You must be signed in to change notification settings - Fork 170
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added new profile mapping configuration for Teradata (#1077)
Teradata has [Provider](https://airflow.apache.org/docs/apache-airflow-providers-teradata/stable/index.html) in airflow and [adapter](https://github.com/Teradata/dbt-teradata) in dbt. The cosmos library doesn't have profile configuration with mapping support. This PR address this issue. Closes: #1053
- Loading branch information
Showing
6 changed files
with
232 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
"""Maps Airflow Snowflake connections to dbt profiles if they use a user/password.""" | ||
|
||
from __future__ import annotations | ||
|
||
from typing import Any | ||
|
||
from ..base import BaseProfileMapping | ||
|
||
|
||
class TeradataUserPasswordProfileMapping(BaseProfileMapping): | ||
""" | ||
Maps Airflow Teradata connections using user + password authentication to dbt profiles. | ||
https://docs.getdbt.com/docs/core/connect-data-platform/teradata-setup | ||
https://airflow.apache.org/docs/apache-airflow-providers-teradata/stable/connections/teradata.html | ||
""" | ||
|
||
airflow_connection_type: str = "teradata" | ||
dbt_profile_type: str = "teradata" | ||
is_community = True | ||
|
||
required_fields = [ | ||
"host", | ||
"user", | ||
"password", | ||
] | ||
secret_fields = [ | ||
"password", | ||
] | ||
airflow_param_mapping = { | ||
"host": "host", | ||
"user": "login", | ||
"password": "password", | ||
"schema": "schema", | ||
"tmode": "extra.tmode", | ||
} | ||
|
||
@property | ||
def profile(self) -> dict[str, Any]: | ||
"""Gets profile. The password is stored in an environment variable.""" | ||
profile = { | ||
**self.mapped_params, | ||
**self.profile_args, | ||
# password should always get set as env var | ||
"password": self.get_env_var_format("password"), | ||
} | ||
# schema is not mandatory in teradata. In teradata user itself a database so if schema is not mentioned | ||
# in both airflow connection and profile_args then treating user as schema. | ||
if "schema" not in self.profile_args and self.mapped_params.get("schema") is None: | ||
profile["schema"] = profile["user"] | ||
|
||
return self.filter_null(profile) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,176 @@ | ||
"""Tests for the postgres profile.""" | ||
|
||
from unittest.mock import patch | ||
|
||
import pytest | ||
from airflow.models.connection import Connection | ||
|
||
from cosmos.profiles import get_automatic_profile_mapping | ||
from cosmos.profiles.teradata.user_pass import TeradataUserPasswordProfileMapping | ||
|
||
|
||
@pytest.fixture() | ||
def mock_teradata_conn(): # type: ignore | ||
""" | ||
Sets the connection as an environment variable. | ||
""" | ||
conn = Connection( | ||
conn_id="my_teradata_connection", | ||
conn_type="teradata", | ||
host="my_host", | ||
login="my_user", | ||
password="my_password", | ||
) | ||
|
||
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): | ||
yield conn | ||
|
||
|
||
@pytest.fixture() | ||
def mock_teradata_conn_custom_tmode(): # type: ignore | ||
""" | ||
Sets the connection as an environment variable. | ||
""" | ||
conn = Connection( | ||
conn_id="my_teradata_connection", | ||
conn_type="teradata", | ||
host="my_host", | ||
login="my_user", | ||
password="my_password", | ||
schema="my_database", | ||
extra='{"tmode": "TERA"}', | ||
) | ||
|
||
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): | ||
yield conn | ||
|
||
|
||
def test_connection_claiming() -> None: | ||
""" | ||
Tests that the teradata profile mapping claims the correct connection type. | ||
""" | ||
# should only claim when: | ||
# - conn_type == teradata | ||
# and the following exist: | ||
# - host | ||
# - user | ||
# - password | ||
potential_values: dict[str, str] = { | ||
"conn_type": "teradata", | ||
"host": "my_host", | ||
"login": "my_user", | ||
"password": "my_password", | ||
} | ||
|
||
# if we're missing any of the values, it shouldn't claim | ||
for key in potential_values: | ||
values = potential_values.copy() | ||
del values[key] | ||
conn = Connection(**values) # type: ignore | ||
|
||
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): | ||
profile_mapping = TeradataUserPasswordProfileMapping(conn) | ||
assert not profile_mapping.can_claim_connection() | ||
|
||
# Even there is no schema, making user as schema as user itself schema in teradata | ||
conn = Connection(**{k: v for k, v in potential_values.items() if k != "schema"}) | ||
with patch("airflow.hooks.base.BaseHook.get_connection", return_value=conn): | ||
profile_mapping = TeradataUserPasswordProfileMapping(conn, {"schema": None}) | ||
assert 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 = TeradataUserPasswordProfileMapping(conn, {"schema": "my_schema"}) | ||
assert profile_mapping.can_claim_connection() | ||
|
||
|
||
def test_profile_mapping_selected( | ||
mock_teradata_conn: Connection, | ||
) -> None: | ||
""" | ||
Tests that the correct profile mapping is selected. | ||
""" | ||
profile_mapping = get_automatic_profile_mapping( | ||
mock_teradata_conn.conn_id, | ||
) | ||
assert isinstance(profile_mapping, TeradataUserPasswordProfileMapping) | ||
|
||
|
||
def test_profile_mapping_keeps_port(mock_teradata_conn: Connection) -> None: | ||
# port is not handled in airflow connection so adding it as profile_args | ||
profile = TeradataUserPasswordProfileMapping(mock_teradata_conn.conn_id, profile_args={"port": 1025}) | ||
assert profile.profile["port"] == 1025 | ||
|
||
|
||
def test_profile_mapping_keeps_custom_tmode(mock_teradata_conn_custom_tmode: Connection) -> None: | ||
profile = TeradataUserPasswordProfileMapping(mock_teradata_conn_custom_tmode.conn_id) | ||
assert profile.profile["tmode"] == "TERA" | ||
|
||
|
||
def test_profile_args( | ||
mock_teradata_conn: Connection, | ||
) -> None: | ||
""" | ||
Tests that the profile values get set correctly. | ||
""" | ||
profile_mapping = get_automatic_profile_mapping( | ||
mock_teradata_conn.conn_id, | ||
profile_args={"schema": "my_database"}, | ||
) | ||
assert profile_mapping.profile_args == { | ||
"schema": "my_database", | ||
} | ||
|
||
assert profile_mapping.profile == { | ||
"type": mock_teradata_conn.conn_type, | ||
"host": mock_teradata_conn.host, | ||
"user": mock_teradata_conn.login, | ||
"password": "{{ env_var('COSMOS_CONN_TERADATA_PASSWORD') }}", | ||
"schema": "my_database", | ||
} | ||
|
||
|
||
def test_profile_args_overrides( | ||
mock_teradata_conn: Connection, | ||
) -> None: | ||
""" | ||
Tests that you can override the profile values. | ||
""" | ||
profile_mapping = get_automatic_profile_mapping( | ||
mock_teradata_conn.conn_id, | ||
profile_args={"schema": "my_schema"}, | ||
) | ||
assert profile_mapping.profile_args == { | ||
"schema": "my_schema", | ||
} | ||
|
||
assert profile_mapping.profile == { | ||
"type": mock_teradata_conn.conn_type, | ||
"host": mock_teradata_conn.host, | ||
"user": mock_teradata_conn.login, | ||
"password": "{{ env_var('COSMOS_CONN_TERADATA_PASSWORD') }}", | ||
"schema": "my_schema", | ||
} | ||
|
||
|
||
def test_profile_env_vars( | ||
mock_teradata_conn: Connection, | ||
) -> None: | ||
""" | ||
Tests that the environment variables get set correctly. | ||
""" | ||
profile_mapping = get_automatic_profile_mapping( | ||
mock_teradata_conn.conn_id, | ||
profile_args={"schema": "my_schema"}, | ||
) | ||
assert profile_mapping.env_vars == { | ||
"COSMOS_CONN_TERADATA_PASSWORD": mock_teradata_conn.password, | ||
} | ||
|
||
|
||
def test_mock_profile() -> None: | ||
""" | ||
Tests that the mock profile port value get set correctly. | ||
""" | ||
profile = TeradataUserPasswordProfileMapping("mock_conn_id") | ||
assert profile.mock_profile.get("host") == "mock_value" |