Add Support for client_tags in TrinoHook Connections

[jjavieralonso]

Description

This feature introduces the client_tags parameter to the TrinoHook connections in Apache Airflow. It allows users to specify a list of client tags when establishing a connection to Trino.

Use case/motivation

The addition of the client_tags parameter enables users to tag their Trino clients with specific identifiers, which can be useful for tracking and managing queries, especially in multi-tenant environments or for monitoring purposes. This feature enhances the flexibility and usability of Trino connections within Apache Airflow.

Related issues

Key Changes:
File: airflow/providers/trino/hooks/trino.py

Added client_tags parameter in the get_conn method.
Updated the connection initialization to include the client_tags parameter.
File: docs/apache-airflow-providers-trino/connections.rst

Updated documentation to include client_tags in the list of extra connection parameters.
File: tests/providers/trino/hooks/test_trino.py

Added test_get_conn_client_tags to verify the correct handling of client_tags.
Updated assert_connection_called_with to check for client_tags in connection parameters.

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[wei-juncheng]

Hi @potiuk ,
could you please assign this issue to me? I am also using TrinoHook and interested in this issue. I'm currently investigating it.

[potiuk]

Sure

[wei-juncheng]

Hi @jjavieralonso ,
I would like to confirm with you that currently, TrinoHook allows system administrators to input client_tags through the Extra parameter when creating a connection. Airflow will then pass the client tag into the Trino connection. In this issue, are you looking for more flexibility to input different client tags when the DAG is running (e.g., DAG information or function information)?

• Apache Trino Connection Document
• Trino's Official Document

[eladkal]

@wei-juncheng maybe we should customize the Trino UI connection form for easier setup?
https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html#custom-connection-fields

Similar to what we have for Snowflake:

airflow/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py

Lines 94 to 117 in 7e9b9ed

	@classmethod
	def get_connection_form_widgets(cls) -> dict[str, Any]:
	"""Return connection widgets to add to connection form."""
	from flask_appbuilder.fieldwidgets import (
	BS3PasswordFieldWidget,
	BS3TextFieldWidget,
	)
	from flask_babel import lazy_gettext
	from wtforms import BooleanField, PasswordField, StringField
	return {
	"account": StringField(lazy_gettext("Account"), widget=BS3TextFieldWidget()),
	"warehouse": StringField(lazy_gettext("Warehouse"), widget=BS3TextFieldWidget()),
	"database": StringField(lazy_gettext("Database"), widget=BS3TextFieldWidget()),
	"region": StringField(lazy_gettext("Region"), widget=BS3TextFieldWidget()),
	"role": StringField(lazy_gettext("Role"), widget=BS3TextFieldWidget()),
	"private_key_file": StringField(lazy_gettext("Private key (Path)"), widget=BS3TextFieldWidget()),
	"private_key_content": PasswordField(
	lazy_gettext("Private key (Text)"), widget=BS3PasswordFieldWidget()
	),
	"insecure_mode": BooleanField(
	label=lazy_gettext("Insecure mode"), description="Turns off OCSP certificate checks"
	),
	}

airflow/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py

Lines 119 to 150 in 7e9b9ed

	@classmethod
	def get_ui_field_behaviour(cls) -> dict[str, Any]:
	"""Return custom field behaviour."""
	import json
	return {
	"hidden_fields": ["port", "host"],
	"relabeling": {},
	"placeholders": {
	"extra": json.dumps(
	{
	"authenticator": "snowflake oauth",
	"private_key_file": "private key",
	"session_parameters": "session parameters",
	"client_request_mfa_token": "client request mfa token",
	"client_store_temporary_credential": "client store temporary credential (externalbrowser mode)",
	},
	indent=1,
	),
	"schema": "snowflake schema",
	"login": "snowflake username",
	"password": "snowflake password",
	"account": "snowflake account name",
	"warehouse": "snowflake warehouse name",
	"database": "snowflake db name",
	"region": "snowflake hosted region",
	"role": "snowflake role",
	"private_key_file": "Path of snowflake private key (PEM Format)",
	"private_key_content": "Content to snowflake private key (PEM format)",
	"insecure_mode": "insecure mode",
	},
	}