Pass X-Trino-Client-Info in trino hook

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,

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',