Build correct SQLAlchemy URI in hooks based on DbApiHook

[Taragolis]

Body

Original discussion: https://lists.apache.org/thread/8rhmz3qh30hvkondct4sfmgk4vd07mn5

tl;dr; DbApiHook assumes that Airflow Connection URI representation is a valid SA URI, but it is not true and never was, it might work in simple cases, however it will fail in some complex one.

airflow/airflow/providers/common/sql/hooks/sql.py

Lines 184 to 192 in 2ab6081

	def get_uri(self) -> str:
	"""
	Extract the URI from the connection.
	:return: the extracted uri.
	"""
	conn = self.get_connection(getattr(self, self.conn_name_attr))
	conn.schema = self.__schema or conn.schema
	return conn.get_uri()

This task intend to track fixes into the hooks which based on DbApiHook.

Hooks which don't overwrite invalid DbApiHook.get_uri implementation

• MySqlHook from mysql provider. This hook provide support for two different drivers.
  • mysqlclient: https://pypi.org/project/mysqlclient/
  • mysql-connector-python: https://pypi.org/project/mysql-connector-python/
• OracleHook from oracle provider
• JdbcHook from jdbc provider
• TrinoHook from trino provider
• HiveServer2Hook from apache.hive provider
• ImpalaHook from apache.impala provider
• DatabricksSqlHook from databricks provider
• ExasolHook from exasol provider
• PrestoHook from presto provider
• VerticaHook from vertica provider
• TeradataHook from teradata provider

Hooks which uses Connection.get_uri for construct SQLAlchemy URI

• PostgresHook from postgres provider. Connection to Postgres failed due to special characters in a password #38187, Connection id "extra" error #34456
• MsSqlHook from microsoft.mssql provider
• SqliteHook from sqlite provider

Committer

• I acknowledge that I am a maintainer/committer of the Apache Airflow project.

[Bowrna]

hello @Taragolis I was checking this issue and seeing if I could fix it.

airflow/airflow/models/connection.py

Lines 271 to 329 in 0e8f108

	def get_uri(self) -> str:
	"""Return connection in URI format."""
	if self.conn_type and "_" in self.conn_type:
	self.log.warning(
	"Connection schemes (type: %s) shall not contain '_' according to RFC3986.",
	self.conn_type,
	)
	if self.conn_type:
	uri = f"{self.conn_type.lower().replace('_', '-')}://"
	else:
	uri = "//"
	if self.host and "://" in self.host:
	protocol, host = self.host.split("://", 1)
	else:
	protocol, host = None, self.host
	if protocol:
	uri += f"{protocol}://"
	authority_block = ""
	if self.login is not None:
	authority_block += quote(self.login, safe="")
	if self.password is not None:
	authority_block += ":" + quote(self.password, safe="")
	if authority_block > "":
	authority_block += "@"
	uri += authority_block
	host_block = ""
	if host:
	host_block += quote(host, safe="")
	if self.port:
	if host_block == "" and authority_block == "":
	host_block += f"@:{self.port}"
	else:
	host_block += f":{self.port}"
	if self.schema:
	host_block += f"/{quote(self.schema, safe='')}"
	uri += host_block
	if self.extra:
	try:
	query: str | None = urlencode(self.extra_dejson)
	except TypeError:
	query = None
	if query and self.extra_dejson == dict(parse_qsl(query, keep_blank_values=True)):
	uri += ("?" if self.schema else "/?") + query
	else:
	uri += ("?" if self.schema else "/?") + urlencode({self.EXTRA_KEY: self.extra})
	return uri

Is fixing the get_uri part so that Airflow Connection URI representation becomes a valid SA URI the right way to fix this issue?

[Taragolis]

As far as i remember, we recently discuss this case in slack. I guess better introduce sa_uri property (or some similar naming) which returns sqlalchemy.engine.URL into the DbApiHook without any implementation, e.g. raise NotImplementedError and implements it per hook which have implementation for SQAlchemy.

get_uri - not clear which URI it is implements, DBApi? SQLAlchemy or even some internal (ElasticsearchSQLHook)

cc: @potiuk @dstandish

[rawwar]

@Taragolis , I am working on this PR - #38871 . This is a separate PR to just update the DbApiHook to include a method to return a SqlAlchemy URL that can be reused.

[dabla]

As far as i remember, we recently discuss this case in slack. I guess better introduce sa_uri property (or some similar naming) which returns sqlalchemy.engine.URL into the DbApiHook without any implementation, e.g. raise NotImplementedError and implements it per hook which have implementation for SQAlchemy.

get_uri - not clear which URI it is implements, DBApi? SQLAlchemy or even some internal (ElasticsearchSQLHook)

cc: @potiuk @dstandish

@Taragolis Why not implement another method named get_url and use that one to instantiate the SQLALchemy engine, as the engine can accept an SQLAlchemy URL object. Also the SQLAlchemy URL class has a factory method named create, which easily allows you to create the URL without the hastle of string manipulations, this is how I did it for my custom SQL operator (the one I mentioned on Slack):

    def get_url(self) -> URL:
        return URL.create(
            self.scheme,
            username=self.connection.login,
            password=self.connection.password,
            host=self.connection.host,
            port=self.connection.port,
        )

Also see here the SQL documentation concerning creating URL's programmatically: https://docs.sqlalchemy.org/en/20/core/engines.html#creating-urls-programmatically

I see @rawwar was first mentioning this, I think this would be a good idea and would simplify code and be less bug prone.

[Taragolis]

There is new property which intends to returns SQL Alchemy URL, current behaviour of get_url is ambiguous and non-all hook return SA URI here, for some of the DB there is no even SA dialects exists, so better keep get_url as is and step-by-step migrate to appropriate ones. So if dialect exists and it installed with provider I do not have any objections to add it implementation as sqlalchemy_url

[dabla]

There is new property which intends to returns SQL Alchemy URL, current behaviour of get_url is ambiguous and non-all hook return SA URI here, for some of the DB there is no even SA dialects exists, so better keep get_url as is and step-by-step migrate to appropriate ones. So if dialect exists and it installed with provider I do not have any objections to add it implementation as sqlalchemy_url

Hey Andrey think there is some confusion here. Yes we should definitly leave the get_uri method as it is (not get_url as it doesn’t exist yet) and implement a new get_url method which would then be called by the existing get_sqlalchemy_engine.

[dabla]

This is how I done it in my custom SQLRowsInsertOperator:

    @cached_property
    def connection(self) -> Connection:
        return BaseHook.get_connection(self.conn_id)  # type: ignore

    @cached_property
    def scheme(self) -> str:
        return self.connection.extra_dejson["sqlalchemy_scheme"]

    @cached_property
    def table_name(self) -> str:
        return self.get_columns().table_name_with_schema()

    @cached_property
    def driver(self) -> str:
        return self.connection.extra_dejson["driver"]

    # TODO: Maybe this should be moved to DbAPiHook
    def get_url(self) -> URL:
        self.log.info("Connection schema: %s", self.schema)
        self.log.info("Connection scheme: %s", self.scheme)
        self.log.info("Connection driver: %s", self.driver)
        self.log.info("Connection type: %s", self.connection.conn_type)
        self.log.info("Connection login: %s", self.connection.login)
        self.log.info("Connection password: %s", self.connection.password)
        self.log.info("Connection host: %s", self.connection.host)
        self.log.info("Connection port: %s", self.connection.port)

        return URL.create(
            self.scheme,
            username=self.connection.login,
            password=self.connection.password,
            host=self.connection.host,
            port=self.connection.port,
        )

   # TODO: Maybe this method in DbAPiHook should be changed to this instead of calling get_uri() method
    def get_sqlalchemy_engine(self, engine_kwargs=None):
        if engine_kwargs is None:
            engine_kwargs = {}
        return create_engine(self.get_url(), **engine_kwargs)

    # TODO: Maybe this should be moved to DbAPiHook
    @cached_property
    def inspector(self):
        engine = self.get_sqlalchemy_engine()
        self.log.debug("Engine drivername: %s", engine.url.drivername)
        self.log.debug("Engine username: %s", engine.url.username)
        self.log.debug("Engine password: %s", engine.url.password)
        self.log.debug("Engine host: %s", engine.url.host)
        self.log.debug("Engine port: %s", engine.url.port)
        self.log.debug("Engine database: %s", engine.url.database)
        self.log.debug("Engine dialect: %s", engine.dialect.name)

        return Inspector.from_engine(engine)

[potiuk]

Hey Andrey think there is some confusion here. Yes we should definitly leave the get_uri method as it is (not get_url as it doesn’t exist yet) and implement a new get_url method which would then be called by the existing get_sqlalchemy_engine.

Yes, adding new method is the way to go. How to do it, should be proposed and discussed in the PR, we have 20+ sql providers, and I think the proposal should be complete and explain how to deal with cases that can/cannot be mapped. Also the method should explicily - IMHO - call our sqlalachemy, (getSqlAlchemyURI ?) and likely allow variations - from what I understand SQLAlchemy can have different URLs not only for different databases, but also for different drivers. Also some of our Connections allows to add different kinds of extras and parameters - not only hosts but other parameters that might be used and passed differently to different drivers.

If we decide to introduce a new "standard" method, in a product like Airflow and common API like DBApiHook, the proposal should provide answers to all the questions that users might ask in order to be mergable. It does not have to handle all of those but at least get enough error handlling and interface specification to explain what is the transition from the Connection form (in all those cases we have DBAPI implementation) -> SQLAlchemy URL. Which variations are handled, how it should work, how to map engine args etc.

I don't think it can be handled by a single implementation, it's more a scaffloiding + implementation for all the implementations and it should be designed nicely with al least a handful (if not all) implementation from the above list to be quite sure that we will be able to handle those cases.

[rawwar]

Hey Andrey think there is some confusion here. Yes we should definitly leave the get_uri method as it is (not get_url as it doesn’t exist yet) and implement a new get_url method which would then be called by the existing get_sqlalchemy_engine.

@dabla , I think, we are indeed adding get_url functionality using sqlalchemy_url property.