Airflow3. URI field does not exist for Connection imported from task-sdk.

[MaksYermak]

Apache Airflow version

3.0.0

If "Other Airflow 2 version" selected, which one?

No response

What happened?

I was trying to replace the import for Connection from from airflow.models import Connection to from airflow.sdk import Connection for making cloud_sql Hook compatible with Airflow 3 and Airflow 2. And I have gotten the error that the uri field does not exist for Connection imported from task-sdk.

It means for me that in the current time we can't replace the code for cloud_sql hook here: https://github.com/apache/airflow/blob/main/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py#L1122. This parameter is an important parameter for establishing a connection for CloudSQL operators. And when in the future we remove the ability for importing Connection from the models these operators will be broke.

What you think should happen instead?

In my opinion Connection from task-sdk and models should have the same fields and interface.

How to reproduce

For cloud_sql hook replace the import from from airflow.models import Connection to from airflow.sdk import Connection. After that run test_cloud_sql unit test.

Operating System

Ubuntu

Versions of Apache Airflow Providers

No response

Deployment

Other

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

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

[tirkarthi]

cc: @amoghrajesh @kaxil

[kaxil]

Given that we are doing all this (linked below) to generate URI, we should just simply construct a Connection object in the google provider

airflow/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py

Lines 1048 to 1103 in 75140f6

	def _generate_connection_uri(self) -> str:
	if self.use_proxy:
	if self.sql_proxy_use_tcp:
	if not self.sql_proxy_tcp_port:
	self.reserve_free_tcp_port()
	if not self.sql_proxy_unique_path:
	self.sql_proxy_unique_path = self._generate_unique_path()
	if not self.database_type:
	raise ValueError("The database_type should be set")
	database_uris = CONNECTION_URIS[self.database_type]
	ssl_spec = None
	socket_path = None
	if self.use_proxy:
	proxy_uris = database_uris["proxy"]
	if self.sql_proxy_use_tcp:
	format_string = proxy_uris["tcp"]
	else:
	format_string = proxy_uris["socket"]
	socket_path = f"{self.sql_proxy_unique_path}/{self._get_instance_socket_name()}"
	else:
	public_uris = database_uris["public"]
	if self.use_ssl:
	format_string = public_uris["ssl"]
	ssl_spec = {"cert": self.sslcert, "key": self.sslkey, "ca": self.sslrootcert}
	else:
	format_string = public_uris["non-ssl"]
	if not self.user:
	raise AirflowException("The login parameter needs to be set in connection")
	if not self.public_ip:
	raise AirflowException("The location parameter needs to be set in connection")
	if not self.password:
	raise AirflowException("The password parameter needs to be set in connection")
	if not self.database:
	raise AirflowException("The database parameter needs to be set in connection")
	connection_uri = format_string.format(
	user=quote_plus(self.user) if self.user else "",
	password=quote_plus(self.password) if self.password else "",
	database=quote_plus(self.database) if self.database else "",
	public_ip=self.public_ip,
	public_port=self.public_port,
	proxy_port=self.sql_proxy_tcp_port,
	socket_path=self._quote(socket_path),
	ssl_spec=self._quote(json.dumps(ssl_spec)) if ssl_spec else "",
	client_cert_file=self._quote(self.sslcert) if self.sslcert else "",
	client_key_file=self._quote(self.sslkey) if self.sslcert else "",
	server_ca_file=self._quote(self.sslrootcert if self.sslcert else ""),
	)
	self.log.info(
	"DB connection URI %s",
	connection_uri.replace(
	quote_plus(self.password) if self.password else "PASSWORD", "XXXXXXXXXXXX"
	),
	)
	return connection_uri

[MaksYermak]

Given that we are doing all this (linked below) to generate URI, we should just simply construct a Connection object in the google provider

Sorry I don't understand what you mean just simply construct a Connection object, could you please clarify?

In cloud_sql the code firstly generates a URI, using a function which you shared, and, then, creates a Connection object with this generated URI. I do not understand how we can create a Connection object, using task-sdk, with this unique URI when the code does not have the URI field.

[kaxil]

The Connection object in Airflow 2.x could use either uri or other args to construct the object. You can just do the following as an example:

diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
index 8ffdd7c1a4..b5accdcb0f 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -1118,8 +1118,45 @@ class CloudSQLDatabaseHook(BaseHook):
         Connection ID will be randomly generated according to whether it uses
         proxy, TCP, UNIX sockets, SSL.
         """
-        uri = self._generate_connection_uri()
-        connection = Connection(conn_id=self.db_conn_id, uri=uri)
+        if not self.user:
+            raise AirflowException("The login parameter needs to be set in connection")
+        if not self.public_ip:
+            raise AirflowException("The location parameter needs to be set in connection")
+        if not self.password:
+            raise AirflowException("The password parameter needs to be set in connection")
+        if not self.database:
+            raise AirflowException("The database parameter needs to be set in connection")
+
+        if self.use_proxy:
+            if self.sql_proxy_use_tcp:
+                host = "127.0.0.1"
+                port = self.sql_proxy_tcp_port
+            else:
+                host = self.sql_proxy_unique_path
+                port = None
+        else:
+            host = self.public_ip
+            port = self.public_port
+
+        extra = {}
+        if self.use_ssl:
+            extra.update({
+                "sslmode": "verify-ca",
+                "sslcert": self.sslcert,
+                "sslkey": self.sslkey,
+                "sslrootcert": self.sslrootcert
+            })
+
+        connection = Connection(
+            conn_id=self.db_conn_id,
+            conn_type=self.database_type,
+            host=host,
+            port=port,
+            login=self.user,
+            password=self.password,
+            schema=self.database,
+            extra=json.dumps(extra) if extra else None
+        )
         self.log.info("Creating connection %s", self.db_conn_id)
         return connection

[amoghrajesh]

@MaksYermak i think we can also try and leverage the get_uri method we expose?
https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/connection.py#L61

[amoghrajesh]

@MaksYermak did you get a chance to get around this one?

[MaksYermak]

@amoghrajesh

@MaksYermak did you get a chance to get around this one?

I didn't have time to check @kaxil's solution. I think this solution may suit us, but I am not sure that it will cover all functionality and I want to check it.

@MaksYermak i think we can also try and leverage the get_uri method we expose? https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/connection.py#L61

This solution does not suit us, because we need to create a Connection object. In the current implementation we create a Connection object using the uri parameter. In the solution which your proposed get_uri method exists only for an already created Connection object.

[amoghrajesh]

Cool. Let us know how it goes

[amoghrajesh]

@MaksYermak this has been fixed in one of the patch releases (3.0.4) likely and it works as expected on main too:

conn = Connection(
    conn_id='test_conn',
    conn_type='postgres',
    host='localhost',
    port=5432,
    login='user',
    password='pass',
    extra='{"ssl": true}'
)
conn.get_uri()
Out[4]: 'postgres://user:pass@localhost:5432/?__extra__=%7B%22ssl%22%3A+true%7D'

Closing this issue in that favour. Request you to consume the latest patch release for Airflow 3 (3.0.5/3.0.6) and it should work

[MaksYermak]

@amoghrajesh this issue wasn't fixed by adding the get_uri() method to the Connection object. This issue is about different interfaces for Connection classes from models and task-sdk. Previously for Connection from models we can create an object using only the uri parameter. For Connection from task-sdk we do not have this option anymore now we need to specify each parameter itself.

In the google provider for cloud_sql we generate a different URIs depending on different types of DB and then create a Connection using this generated URI. The absence of this parameter in the interface for Connection from task-sdk blocks us for smooth migration for cloud_sql hook.

If we are not planning to add this uri parameter for creating a Connection then for cloud_sql hook we need to do a big refactoring by using this new approach for creating a Connection. As I see this potential refactoring may cause some problems.

@kaxil @potiuk do we have some dates or Airflow versions when we are planning to completely remove the Connection from models?

[potiuk]

In the google provider for cloud_sql we generate a different URIs depending on different types of DB and then create a Connection using this generated URI. The absence of this parameter in the interface for Connection from task-sdk blocks us for smooth migration for cloud_sql hook.

Actually I implemented original version of it :). Many years ago. And as I see it now, this is not really a big issue to refactor the method. Currently the method _generate_connection_uri does this:

        connection_uri = format_string.format(
            user=quote_plus(self.user) if self.user else "",
            password=quote_plus(self.password) if self.password else "",
            database=quote_plus(self.database) if self.database else "",
            public_ip=self.public_ip,
            public_port=self.public_port,
            proxy_port=self.sql_proxy_tcp_port,
            socket_path=self._quote(socket_path),
            ssl_spec=self._quote(json.dumps(ssl_spec)) if ssl_spec else "",
            client_cert_file=self._quote(self.sslcert) if self.sslcert else "",
            client_key_file=self._quote(self.sslkey) if self.sslcert else "",
            server_ca_file=self._quote(self.sslrootcert if self.sslcert else ""),
        )
        self.log.info(
            "DB connection URI %s",
            connection_uri.replace(
                quote_plus(self.password) if self.password else "PASSWORD", "XXXXXXXXXXXX"
            ),
        )
        return connection_uri

and it's only used in one place (and one test)

        uri = self._generate_connection_uri()
        connection = Connection(conn_id=self.db_conn_id, uri=uri)

But going through URI is not necessary - neither in Airflow 2 nor 3, you could build dict with kwargs and initialize Connection with it. It does not seem to outreagous to change the methot to:

_generate_connection() -> task.sdk.Connection | model.Connection

And do

• if airflow 2 -> create connection via uri
• if aorflow 3 -> create connection via default parameters used

Or even

• build kwarg dict for Connection
• if airlfow 2 -> return models.Connection(**kwargs)
• if airflow 3 -> return task.sdk.Connection(**kwargs)

Then we could drop the Airflow 2 path when we drop Airflow 2 support.