chore: split out Python Connector IAM authn sample and region tags (GoogleCloudPlatform#8164)

jackwotherspoon · web-flow · commit 3558b7caac0b · 2022-07-11T17:05:10.000-04:00
diff --git a/cloud-sql/postgres/sqlalchemy/app.py b/cloud-sql/postgres/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/postgres/sqlalchemy/connect_connector.py b/cloud-sql/postgres/sqlalchemy/connect_connector.py
@@ -29,18 +29,15 @@ def connect_with_connector() -> sqlalchemy.engine.base.Engine:
     # Cloud Secret Manager (https://cloud.google.com/secret-manager) to help
     # keep secrets safe.
 
-    # Either a DB_USER or a DB_IAM_USER should be defined. If both are
-    # defined, DB_IAM_USER takes precedence.
     instance_connection_name = os.environ["INSTANCE_CONNECTION_NAME"]  # e.g. 'project:region:instance'
-    db_user = os.environ.get("DB_USER", "")  # e.g. 'my-db-user'
-    db_iam_user = os.environ.get("DB_IAM_USER", "")  # e.g. 'sa-name@project-id.iam'
+    db_user = os.environ["DB_USER"]  # e.g. 'my-db-user'
     db_pass = os.environ["DB_PASS"]  # e.g. 'my-db-password'
     db_name = os.environ["DB_NAME"]  # e.g. 'my-database'
 
-    enable_iam_auth, db_user = (True, db_iam_user) if db_iam_user else (False, db_user)
     ip_type = IPTypes.PRIVATE if os.environ.get("PRIVATE_IP") else IPTypes.PUBLIC
 
-    connector = Connector(ip_type, enable_iam_auth)
+    # initialize Cloud SQL Python Connector object
+    connector = Connector()
 
     def getconn() -> pg8000.dbapi.Connection:
         conn: pg8000.dbapi.Connection = connector.connect(
@@ -49,9 +46,12 @@ def getconn() -> pg8000.dbapi.Connection:
             user=db_user,
             password=db_pass,
             db=db_name,
+            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(
         "postgresql+pg8000://",
         creator=getconn,
diff --git a/cloud-sql/postgres/sqlalchemy/connect_connector_auto_iam_authn.py b/cloud-sql/postgres/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_postgres_sqlalchemy_auto_iam_authn]
+import os
+
+from google.cloud.sql.connector import Connector, IPTypes
+import pg8000
+
+import sqlalchemy
+
+
+# connect_with_connector_auto_iam_authn initializes a connection pool for
+# a Cloud SQL instance of Postgres 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() -> pg8000.dbapi.Connection:
+        conn: pg8000.dbapi.Connection = connector.connect(
+            instance_connection_name,
+            "pg8000",
+            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(
+        "postgresql+pg8000://",
+        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_postgres_sqlalchemy_auto_iam_authn]
