diff --git a/airflow_dbt_python/hooks/target.py b/airflow_dbt_python/hooks/target.py index ac74f55..0a8d581 100644 --- a/airflow_dbt_python/hooks/target.py +++ b/airflow_dbt_python/hooks/target.py @@ -492,3 +492,38 @@ class DbtSparkHook(DbtConnectionHook): ), ] conn_extra_params = [] + + +class DbtTrinoHook(DbtConnectionHook): + """A hook to interact with dbt using a Trino connection.""" + + conn_type = "trino" + hook_name = "dbt Trino Hook" + airflow_conn_types = (conn_type,) + + conn_params = [ + "host", + DbtConnectionParam("schema", default="public"), + DbtConnectionParam("login", "user"), + "password", + DbtConnectionParam("port", default=443), + ] + + conn_extra_params = [ + DbtConnectionParam("type", default="trino"), + DbtConnectionParam("method", default="none"), + DbtConnectionParam("http_scheme", default="https"), + DbtConnectionParam("verify", default=True), + DbtConnectionParam("database", "catalog", "hive"), + "catalog", + "host", + "port", + DbtConnectionParam("user", "user"), + "password", + "schema", + "session_properties", + "http_headers", + "role", + "client_tags", + "source", + ] diff --git a/tests/hooks/test_trino.py b/tests/hooks/test_trino.py new file mode 100644 index 0000000..2bdb3d0 --- /dev/null +++ b/tests/hooks/test_trino.py @@ -0,0 +1,65 @@ +"""Tests for the DbtTrinoHook.""" + +from airflow.models.connection import Connection + +from airflow_dbt_python.hooks.target import DbtConnectionHook, DbtTrinoHook + + +def test_trino_registration(): + """Ensure that a Trino connection is mapped to a valid dbt profile dict.""" + conn = Connection( + conn_id="my_trino", + conn_type="trino", + extra=""" + { + "type": "trino", + "method": "ldap", + "user": "user", + "password": "pass", + "host": "trino.example", + "port": 443, + "http_scheme": "https", + "database": "example", + "schema": "example", + "verify": false + } + """, + ) + + trino_hook = DbtTrinoHook(conn=conn) + details = trino_hook.get_dbt_details_from_connection(conn) + + assert details["type"] == "trino" + assert details["method"] == "ldap" + assert details["host"] == "trino.example" + assert details["port"] == 443 + assert details["user"] == "user" + assert details["password"] == "pass" + assert details["schema"] == "example" + assert details["catalog"] == "example" + assert details["verify"] is False + + +def test_trino_catalog_overrides_database(): + """Ensure that explicit catalog overrides database → catalog fallback.""" + conn = Connection( + conn_id="my_trino_catalog", + conn_type="trino", + extra=""" + { + "type": "trino", + "user": "u", + "password": "p", + "host": "h", + "port": 8443, + "database": "fallback_db", + "catalog": "explicit_catalog", + "schema": "s" + } + """, + ) + + trino_hook = DbtTrinoHook(conn=conn) + details = trino_hook.get_dbt_details_from_connection(conn) + + assert details["catalog"] == "explicit_catalog"