diff --git a/airflow/providers/trino/hooks/trino.py b/airflow/providers/trino/hooks/trino.py index 2401c327d79b2..37bc575c6850f 100644 --- a/airflow/providers/trino/hooks/trino.py +++ b/airflow/providers/trino/hooks/trino.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json import os import warnings from typing import Any, Callable, Iterable, Optional, overload @@ -27,6 +28,34 @@ from airflow.configuration import conf from airflow.hooks.dbapi import DbApiHook from airflow.models import Connection +from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING + +try: + from airflow.utils.operator_helpers import DEFAULT_FORMAT_PREFIX +except ImportError: + # This is from airflow.utils.operator_helpers, + # For the sake of provider backward compatibility, this is hardcoded if import fails + # https://github.com/apache/airflow/pull/22416#issuecomment-1075531290 + DEFAULT_FORMAT_PREFIX = 'airflow.ctx.' + + +def generate_trino_client_info() -> str: + """Return json string with dag_id, task_id, execution_date and try_number""" + context_var = { + format_map['default'].replace(DEFAULT_FORMAT_PREFIX, ''): os.environ.get( + format_map['env_var_format'], '' + ) + for format_map in AIRFLOW_VAR_NAME_FORMAT_MAPPING.values() + } + task_info = { + 'dag_id': context_var['dag_id'], + 'task_id': context_var['task_id'], + 'execution_date': context_var['execution_date'], + 'try_number': context_var['try_number'], + 'dag_run_id': context_var['dag_run_id'], + 'dag_owner': context_var['dag_owner'], + } + return json.dumps(task_info, sort_keys=True) class TrinoException(Exception): @@ -82,12 +111,15 @@ def get_conn(self) -> Connection: delegate=_boolify(extra.get('kerberos__delegate', False)), ca_bundle=extra.get('kerberos__ca_bundle'), ) + + http_headers = {"X-Trino-Client-Info": generate_trino_client_info()} trino_conn = trino.dbapi.connect( host=db.host, port=db.port, user=db.login, source=extra.get('source', 'airflow'), http_scheme=extra.get('protocol', 'http'), + http_headers=http_headers, catalog=extra.get('catalog', 'hive'), schema=db.schema, auth=auth, diff --git a/tests/providers/trino/hooks/test_trino.py b/tests/providers/trino/hooks/test_trino.py index 09d5aa592536d..883794bf0f21a 100644 --- a/tests/providers/trino/hooks/test_trino.py +++ b/tests/providers/trino/hooks/test_trino.py @@ -46,6 +46,41 @@ def test_get_conn_basic_auth(self, mock_get_connection, mock_connect, mock_basic self.assert_connection_called_with(mock_connect, auth=mock_basic_auth) mock_basic_auth.assert_called_once_with('login', 'password') + @patch('airflow.providers.trino.hooks.trino.generate_trino_client_info') + @patch(BASIC_AUTHENTICATION) + @patch(TRINO_DBAPI_CONNECT) + @patch(HOOK_GET_CONNECTION) + def test_http_headers( + self, + mock_get_connection, + mock_connect, + mock_basic_auth, + mocked_generate_airflow_trino_client_info_header, + ): + mock_get_connection.return_value = Connection( + login='login', password='password', host='host', schema='hive' + ) + client = json.dumps( + { + "dag_id": "dag-id", + "execution_date": "2022-01-01T00:00:00", + "task_id": "task-id", + "try_number": "1", + "dag_run_id": "dag-run-id", + "dag_owner": "dag-owner", + }, + sort_keys=True, + ) + http_headers = {'X-Trino-Client-Info': client} + + mocked_generate_airflow_trino_client_info_header.return_value = http_headers['X-Trino-Client-Info'] + + conn = TrinoHook().get_conn() + self.assert_connection_called_with(mock_connect, auth=mock_basic_auth, http_headers=http_headers) + + mock_basic_auth.assert_called_once_with('login', 'password') + assert mock_connect.return_value == conn + @patch(HOOK_GET_CONNECTION) def test_get_conn_invalid_auth(self, mock_get_connection): extras = {'auth': 'kerberos'} @@ -108,12 +143,13 @@ def set_get_connection_return_value(mock_get_connection, extra=None, password=No mock_get_connection.return_value = mocked_connection @staticmethod - def assert_connection_called_with(mock_connect, auth=None, verify=True): + def assert_connection_called_with(mock_connect, http_headers=mock.ANY, auth=None, verify=True): mock_connect.assert_called_once_with( catalog='hive', host='host', port=None, http_scheme='http', + http_headers=http_headers, schema='hive', source='airflow', user='login',