Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add Support for client_tags in TrinoHook Connections #45841

Open
1 of 2 tasks
jjavieralonso opened this issue Jan 21, 2025 · 5 comments · May be fixed by #46996
Open
1 of 2 tasks

Add Support for client_tags in TrinoHook Connections #45841

jjavieralonso opened this issue Jan 21, 2025 · 5 comments · May be fixed by #46996
Assignees
@wei-juncheng
Labels
area:providers good first issue kind:feature Feature Requests provider:trino

Comments

@jjavieralonso
Copy link

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
@jjavieralonso jjavieralonso added kind:feature Feature Requests needs-triage label for new issues that we didn't triage yet labels Jan 21, 2025
@boring-cyborg boring-cyborg
Copy link

boring-cyborg bot commented Jan 21, 2025

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.

@potiuk potiuk added good first issue and removed needs-triage label for new issues that we didn't triage yet labels Jan 21, 2025
@wei-juncheng
Copy link

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
Copy link
Member

potiuk commented Jan 25, 2025

Sure

@wei-juncheng
Copy link

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)?

@eladkal
Copy link
Contributor

eladkal commented Feb 17, 2025

@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",
},
}

@wei-juncheng wei-juncheng linked a pull request Feb 23, 2025 that will close this issue
Open
2 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@wei-juncheng wei-juncheng

Labels
area:providers good first issue kind:feature Feature Requests provider:trino
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Add Trino Connection UI Fields Placeholders wei-juncheng/airflow
4 participants
@potiuk @wei-juncheng @eladkal @jjavieralonso