feat(cloudsql): add MySQL Auto IAM AuthN sample (GoogleCloudPlatform#8526)

jackwotherspoon · web-flow · commit bcb2b2748eda · 2022-12-07T15:37:42.000-05:00
diff --git a/cloud-sql/mysql/sqlalchemy/app.py b/cloud-sql/mysql/sqlalchemy/app.py
@@ -22,6 +22,7 @@
 import sqlalchemy
 
 from connect_connector import connect_with_connector
+from connect_connector_auto_iam_authn import connect_with_connector_auto_iam_authn
 from connect_tcp import connect_tcp_socket
 from connect_unix import connect_unix_socket
 
@@ -41,7 +42,9 @@ def init_connection_pool() -> sqlalchemy.engine.base.Engine:
 
     # use the connector when INSTANCE_CONNECTION_NAME (e.g. project:region:instance) is defined
     if os.environ.get("INSTANCE_CONNECTION_NAME"):
-        return connect_with_connector()
+        # Either a DB_USER or a DB_IAM_USER should be defined. If both are
+        # defined, DB_IAM_USER takes precedence.
+        return connect_with_connector_auto_iam_authn() if os.environ.get("DB_IAM_USER") else connect_with_connector()
 
     raise ValueError(
         "Missing database connection type. Please define one of INSTANCE_HOST, INSTANCE_UNIX_SOCKET, or INSTANCE_CONNECTION_NAME"
diff --git a/cloud-sql/mysql/sqlalchemy/connect_connector_auto_iam_authn.py b/cloud-sql/mysql/sqlalchemy/connect_connector_auto_iam_authn.py
@@ -0,0 +1,80 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START cloud_sql_mysql_sqlalchemy_auto_iam_authn]
+import os
+
+from google.cloud.sql.connector import Connector, IPTypes
+import pymysql
+
+import sqlalchemy
+
+
+# connect_with_connector_auto_iam_authn initializes a connection pool for
+# a Cloud SQL instance of MySQL using the Cloud SQL Python Connector
+# with Automatic IAM Database Authentication.
+def connect_with_connector_auto_iam_authn() -> sqlalchemy.engine.base.Engine:
+    # Note: Saving credentials in environment variables is convenient, but not
+    # secure - consider a more secure solution such as
+    # Cloud Secret Manager (https://cloud.google.com/secret-manager) to help
+    # keep secrets safe.
+    instance_connection_name = os.environ["INSTANCE_CONNECTION_NAME"]  # e.g. 'project:region:instance'
+    db_iam_user = os.environ["DB_IAM_USER"]  # e.g. 'sa-name@project-id.iam'
+    db_name = os.environ["DB_NAME"]  # e.g. 'my-database'
+
+    ip_type = IPTypes.PRIVATE if os.environ.get("PRIVATE_IP") else IPTypes.PUBLIC
+
+    # initialize Cloud SQL Python Connector object
+    connector = Connector()
+
+    def getconn() -> pymysql.connections.Connection:
+        conn: pymysql.connections.Connection = connector.connect(
+            instance_connection_name,
+            "pymysql",
+            user=db_iam_user,
+            db=db_name,
+            enable_iam_auth=True,
+            ip_type=ip_type,
+        )
+        return conn
+
+    # The Cloud SQL Python Connector can be used with SQLAlchemy
+    # using the 'creator' argument to 'create_engine'
+    pool = sqlalchemy.create_engine(
+        "mysql+pymysql://",
+        creator=getconn,
+        # [START_EXCLUDE]
+        # Pool size is the maximum number of permanent connections to keep.
+        pool_size=5,
+
+        # Temporarily exceeds the set pool_size if no connections are available.
+        max_overflow=2,
+
+        # The total number of concurrent connections for your application will be
+        # a total of pool_size and max_overflow.
+
+        # 'pool_timeout' is the maximum number of seconds to wait when retrieving a
+        # new connection from the pool. After the specified amount of time, an
+        # exception will be thrown.
+        pool_timeout=30,  # 30 seconds
+
+        # 'pool_recycle' is the maximum number of seconds a connection can persist.
+        # Connections that live longer than the specified amount of time will be
+        # re-established
+        pool_recycle=1800,  # 30 minutes
+        # [END_EXCLUDE]
+    )
+    return pool
+
+# [END cloud_sql_mysql_sqlalchemy_auto_iam_authn]
diff --git a/cloud-sql/mysql/sqlalchemy/requirements.txt b/cloud-sql/mysql/sqlalchemy/requirements.txt
@@ -2,4 +2,4 @@ Flask==2.1.0
 SQLAlchemy==1.4.38
 PyMySQL==1.0.2
 gunicorn==20.1.0
-cloud-sql-python-connector==0.6.0
+cloud-sql-python-connector==1.0.0
diff --git a/cloud-sql/postgres/sqlalchemy/requirements.txt b/cloud-sql/postgres/sqlalchemy/requirements.txt
@@ -1,5 +1,5 @@
 Flask==2.1.0
 pg8000==1.24.2
 SQLAlchemy==1.4.38
-cloud-sql-python-connector==0.6.1
+cloud-sql-python-connector==1.0.0
 gunicorn==20.1.0
diff --git a/cloud-sql/sql-server/sqlalchemy/requirements.txt b/cloud-sql/sql-server/sqlalchemy/requirements.txt
@@ -3,5 +3,5 @@ gunicorn==20.1.0
 python-tds==1.11.0
 pyopenssl==22.1.0
 SQLAlchemy==1.4.38
-cloud-sql-python-connector==0.6.1
+cloud-sql-python-connector==1.0.0
 sqlalchemy-pytds==0.3.4
