Skip to content

Fix overwrite get_uri for Trino - #48917

Merged
Lee-W merged 1 commit into
apache:mainfrom
guan404ming:add-get-uri-for-trino
Apr 16, 2025
Merged

Lee-W merged 1 commit into
apache:mainfrom
guan404ming:add-get-uri-for-trino

Conversation

@guan404ming

@guan404ming guan404ming commented Apr 8, 2025 •

Copy link
Copy Markdown
Member

Related Issue

Towards #38195

Why

DbApiHook's get_uri() incorrectly assumes Airflow Connection URIs are valid SQLAlchemy URIs, which fails in complex cases, especially with special characters. For Trino, the connection URI requires proper URL encoding and specific handling of catalog/schema information.

This PR implements a robust get_uri() method in TrinoHook that correctly formats Trino connection URIs with proper URL encoding of special characters.

How

  • Implemented get_uri() in TrinoHook with proper URL
  • Handled schema information from both schema field and extra field
  • Added unit tests for various connection scenarios

reference: https://github.com/trinodb/trino-python-client?tab=readme-ov-file

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

@guan404ming guan404ming changed the title fix: overwrite get-uri for Trino fix: overwrite get_uri for Trino Apr 8, 2025

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR!
I left a question regarding whether we really need to override get_uri for TrinoHook, as it seems there’s no special case for it.

Comment thread providers/trino/src/airflow/providers/trino/hooks/trino.py

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we could consider implementing a more generic get_uri in DbApiHook.

Yes, indeed. We can create another PR after this one to introduce some private methods in DbApiHook.
( it's fine to have duplicate code in the current one, not to make the scope too large )

For example, comparing the previous MySQL implementation and the current Trino one, we have duplicated logic in:

  • uri_prefix
  • auth_part
  • host_part
  • schema_part
  • URI concatenation
  • extra connection parameters

The DbApiHook.get_uri could leverage these private methods to build the final URI.

We could define private methods for each of these duplicated parts, and allow providers to override only what’s necessary for special cases.

For instance, Trino would only need to override _get_schema_part, MySQL only need to override _get_uri_prefix part.

Comment thread providers/trino/src/airflow/providers/trino/hooks/trino.py
@guan404ming

guan404ming commented Apr 10, 2025 •

Copy link
Copy Markdown
Member Author

Thanks for looking into this issue and suggestion! It does make sense, to implement common private functions do help reduce code duplicate issue. I would try to open another PR to implement some private func in DbApiHook after this PR!

@guan404ming
guan404ming requested a review from jason810496 April 11, 2025 16:04
@guan404ming
guan404ming force-pushed the add-get-uri-for-trino branch from 0bc24c2 to bcecb31 Compare April 12, 2025 02:01

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After clearing the context, the PR looks good to me, thanks!

@guan404ming

Copy link
Copy Markdown
Member Author

Thanks!

@guan404ming
guan404ming force-pushed the add-get-uri-for-trino branch from bcecb31 to 125df75 Compare April 14, 2025 18:31
@jason810496
jason810496 requested a review from Lee-W April 15, 2025 08:44
@guan404ming
guan404ming force-pushed the add-get-uri-for-trino branch from 125df75 to 74bcc84 Compare April 15, 2025 19:19
@amoghrajesh

Copy link
Copy Markdown
Contributor

CC @eladkal @jason810496 do you think this one's good to land?

@jason810496

Copy link
Copy Markdown
Member

CC @eladkal @jason810496 do you think this one's good to land?

Sure, LGTM to include in current provider release.

@Lee-W
Lee-W merged commit 39a373c into apache:main Apr 16, 2025
@guan404ming guan404ming changed the title fix: overwrite get_uri for Trino Fix overwrite get_uri for Trino Apr 30, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants