Fix bugs in URI constructor for MySQL connection (#24320)

* Fix bugs in URI constructor for MySQL connection

* Update unit tests

(cherry picked from commit ea54faf)

airflow/models/connection.py

@@ -231,10 +231,10 @@ def get_uri(self) -> str:
             host_block += quote(self.host, safe='')
 
         if self.port:
-            if host_block > '':
-                host_block += f':{self.port}'
-            else:
+            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='')}"
@@ -247,9 +247,9 @@ def get_uri(self) -> str:
             except TypeError:
                 query = None
             if query and self.extra_dejson == dict(parse_qsl(query, keep_blank_values=True)):
-                uri += '?' + query
+                uri += ('?' if self.schema else '/?') + query
             else:
-                uri += '?' + urlencode({self.EXTRA_KEY: self.extra})
+                uri += ('?' if self.schema else '/?') + urlencode({self.EXTRA_KEY: self.extra})
 
         return uri
 

tests/cli/commands/test_connection_command.py

@@ -242,14 +242,14 @@ def test_cli_connections_export_should_export_as_yaml(self, tmp_path):
                 'uri',
                 [
                     "airflow_db=mysql://root:plainpassword@mysql/airflow",
-                    "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql",
+                    "druid_broker_default=druid://druid-broker:8082/?endpoint=druid%2Fv2%2Fsql",
                 ],
             ),
             (
                 None,  # tests that default is URI
                 [
                     "airflow_db=mysql://root:plainpassword@mysql/airflow",
-                    "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql",
+                    "druid_broker_default=druid://druid-broker:8082/?endpoint=druid%2Fv2%2Fsql",
                 ],
             ),
             (
@@ -287,7 +287,7 @@ def test_cli_connections_export_should_export_as_env_for_uppercase_file_extensio
         connection_command.connections_export(args)
         expected_connections = [
             "airflow_db=mysql://root:plainpassword@mysql/airflow",
-            "druid_broker_default=druid://druid-broker:8082?endpoint=druid%2Fv2%2Fsql",
+            "druid_broker_default=druid://druid-broker:8082/?endpoint=druid%2Fv2%2Fsql",
         ]
 
         assert output_filepath.read_text().splitlines() == expected_connections

tests/hooks/test_dbapi.py

@@ -17,6 +17,7 @@
 # under the License.
 #
 
+import json
 import unittest
 from unittest import mock
 
@@ -235,6 +236,111 @@ def test_get_uri_authority_none(self):
         )
         assert "conn-type://host:1/schema" == self.db_hook.get_uri()
 
+    def test_get_uri_extra(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login='login',
+                password='password',
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login:password@host/?charset=utf-8"
+
+    def test_get_uri_extra_with_schema(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login='login',
+                password='password',
+                schema="schema",
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login:password@host/schema?charset=utf-8"
+
+    def test_get_uri_extra_with_port(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login='login',
+                password='password',
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login:password@host:3306/?charset=utf-8"
+
+    def test_get_uri_extra_with_port_and_empty_host(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                login='login',
+                password='password',
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login:password@:3306/?charset=utf-8"
+
+    def test_get_uri_extra_with_port_and_schema(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login='login',
+                password='password',
+                schema="schema",
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login:password@host:3306/schema?charset=utf-8"
+
+    def test_get_uri_without_password(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login='login',
+                password=None,
+                schema="schema",
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://login@host:3306/schema?charset=utf-8"
+
+    def test_get_uri_without_auth(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                host="host",
+                login=None,
+                password=None,
+                schema="schema",
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://host:3306/schema?charset=utf-8"
+
+    def test_get_uri_without_auth_and_empty_host(self):
+        self.db_hook.get_connection = mock.MagicMock(
+            return_value=Connection(
+                conn_type="conn-type",
+                login=None,
+                password=None,
+                schema="schema",
+                port=3306,
+                extra=json.dumps({'charset': 'utf-8'}),
+            )
+        )
+        assert self.db_hook.get_uri() == "conn-type://@:3306/schema?charset=utf-8"
+
     def test_run_log(self):
         statement = 'SQL'
         self.db_hook.run(statement)

tests/secrets/test_local_filesystem.py

@@ -157,8 +157,8 @@ def test_env_file_should_load_connection(self, file_content, expected_connection
     @parameterized.expand(
         (
             (
-                "CONN_ID=mysql://host_1?param1=val1&param2=val2",
-                {"CONN_ID": "mysql://host_1?param1=val1&param2=val2"},
+                "CONN_ID=mysql://host_1/?param1=val1&param2=val2",
+                {"CONN_ID": "mysql://host_1/?param1=val1&param2=val2"},
             ),
         )
     )