[AIRFLOW-5768] GCP cloud sql don't store ephemeral connection in db

airflow/gcp/hooks/cloud_sql.py

@@ -934,57 +934,16 @@ def _get_sqlproxy_instance_specification(self) -> str:
             instance_specification += "=tcp:" + str(self.sql_proxy_tcp_port)
         return instance_specification
 
-    @provide_session
-    def create_connection(self, session: Optional[Session] = None) -> None:
+    def create_connection(self) -> Connection:
         """
-        Create connection in the Connection table, according to whether it uses
-        proxy, TCP, UNIX sockets, SSL. Connection ID will be randomly generated.
-
-        :param session: Session of the SQL Alchemy ORM (automatically generated with
-                        decorator).
+        Create Connection object, according to whether it uses proxy, TCP, UNIX sockets, SSL.
+        Connection ID will be randomly generated.
         """
-        assert session is not None
         connection = Connection(conn_id=self.db_conn_id)
         uri = self._generate_connection_uri()
         self.log.info("Creating connection %s", self.db_conn_id)
         connection.parse_from_uri(uri)
-        session.add(connection)
-        session.commit()
-
-    @provide_session
-    def retrieve_connection(self, session: Optional[Session] = None) -> Optional[Connection]:
-        """
-        Retrieves the dynamically created connection from the Connection table.
-
-        :param session: Session of the SQL Alchemy ORM (automatically generated with
-                        decorator).
-        """
-        assert session is not None
-        self.log.info("Retrieving connection %s", self.db_conn_id)
-        connections = session.query(Connection).filter(
-            Connection.conn_id == self.db_conn_id)
-        if connections.count():
-            return connections[0]
-        return None
-
-    @provide_session
-    def delete_connection(self, session: Optional[Session] = None) -> None:
-        """
-        Delete the dynamically created connection from the Connection table.
-
-        :param session: Session of the SQL Alchemy ORM (automatically generated with
-                        decorator).
-        """
-        assert session is not None
-        self.log.info("Deleting connection %s", self.db_conn_id)
-        connections = session.query(Connection).filter(
-            Connection.conn_id == self.db_conn_id)
-        if connections.count():
-            connection = connections[0]
-            session.delete(connection)
-            session.commit()
-        else:
-            self.log.info("Connection was already deleted!")
+        return connection
 
     def get_sqlproxy_runner(self) -> CloudSqlProxyRunner:
         """
@@ -1006,17 +965,15 @@ def get_sqlproxy_runner(self) -> CloudSqlProxyRunner:
             gcp_conn_id=self.gcp_conn_id
         )
 
-    def get_database_hook(self) -> Union[PostgresHook, MySqlHook]:
+    def get_database_hook(self, connection: Connection) -> Union[PostgresHook, MySqlHook]:
         """
         Retrieve database hook. This is the actual Postgres or MySQL database hook
         that uses proxy or connects directly to the Google Cloud SQL database.
         """
         if self.database_type == 'postgres':
-            self.db_hook = PostgresHook(postgres_conn_id=self.db_conn_id,
-                                        schema=self.database)
+            self.db_hook = PostgresHook(connection=connection, schema=self.database)
         else:
-            self.db_hook = MySqlHook(mysql_conn_id=self.db_conn_id,
-                                     schema=self.database)
+            self.db_hook = MySqlHook(connection=connection, schema=self.database)
         return self.db_hook
 
     def cleanup_database_hook(self) -> None:

airflow/gcp/operators/cloud_sql.py

@@ -835,13 +835,10 @@ def execute(self, context):
                 'extra__google_cloud_platform__project')
         )
         hook.validate_ssl_certs()
-        hook.create_connection()
+        connection = hook.create_connection()
+        hook.validate_socket_path_length()
+        database_hook = hook.get_database_hook(connection=connection)
         try:
-            hook.validate_socket_path_length()
-            database_hook = hook.get_database_hook()
-            try:
-                self._execute_query(hook, database_hook)
-            finally:
-                hook.cleanup_database_hook()
+            self._execute_query(hook, database_hook)
         finally:
-            hook.delete_connection()
+            hook.cleanup_database_hook()

airflow/hooks/mysql_hook.py

@@ -47,6 +47,7 @@ class MySqlHook(DbApiHook):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.schema = kwargs.pop("schema", None)
+        self.connection = kwargs.pop("connection", None)
 
     def set_autocommit(self, conn, autocommit):
         """
@@ -69,7 +70,7 @@ def get_conn(self):
         """
         Returns a mysql connection object
         """
-        conn = self.get_connection(self.mysql_conn_id)
+        conn = self.connection or self.get_connection(self.mysql_conn_id)
 
         conn_config = {
             "user": conn.login,

airflow/hooks/postgres_hook.py

@@ -56,6 +56,7 @@ class PostgresHook(DbApiHook):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.schema = kwargs.pop("schema", None)
+        self.connection = kwargs.pop("connection", None)
 
     def _get_cursor(self, raw_cursor):
         _cursor = raw_cursor.lower()
@@ -68,8 +69,9 @@ def _get_cursor(self, raw_cursor):
         raise ValueError('Invalid cursor passed {}'.format(_cursor))
 
     def get_conn(self):
+
         conn_id = getattr(self, self.conn_name_attr)
-        conn = self.get_connection(conn_id)
+        conn = self.connection or self.get_connection(conn_id)
 
         # check for authentication via AWS IAM
         if conn.extra_dejson.get('iam', False):