chore: update tests and README to use lambda over getconn func

Author: rhatgadkar-goog

README.md

@@ -134,9 +134,10 @@ import sqlalchemy
 # initialize Connector object
 connector = Connector()
 
-# function to return the database connection
-def getconn():
-    conn = connector.connect(
+# initialize SQLAlchemy connection pool with Connector
+pool = sqlalchemy.create_engine(
+    "postgresql+pg8000://",
+    creator=lambda: connector.connect(
         "projects/<YOUR_PROJECT>/locations/<YOUR_REGION>/clusters/<YOUR_CLUSTER>/instances/<YOUR_INSTANCE>",
         "pg8000",
         user="my-user",
@@ -145,13 +146,7 @@ def getconn():
         # NOTE: this assumes private IP by default.
         # Add the following keyword arg to use public IP:
         # ip_type="PUBLIC"
-    )
-    return conn
-
-# create connection pool
-pool = sqlalchemy.create_engine(
-    "postgresql+pg8000://",
-    creator=getconn,
+    ),
 )
 ```
 
@@ -196,30 +191,19 @@ Connector as a context manager:
 from google.cloud.alloydb.connector import Connector
 import sqlalchemy
 
-# helper function to return SQLAlchemy connection pool
-def init_connection_pool(connector: Connector) -> sqlalchemy.engine.Engine:
-    # function used to generate database connection
-    def getconn():
-        conn = connector.connect(
+# initialize Connector as context manager
+with Connector() as connector:
+    # initialize SQLAlchemy connection pool with Connector
+    pool = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=lambda: connector.connect(
             "projects/<YOUR_PROJECT>/locations/<YOUR_REGION>/clusters/<YOUR_CLUSTER>/instances/<YOUR_INSTANCE>",
             "pg8000",
             user="my-user",
             password="my-password",
             db="my-db-name"
-        )
-        return conn
-
-    # create connection pool
-    pool = sqlalchemy.create_engine(
-        "postgresql+pg8000://",
-        creator=getconn,
+        ),
     )
-    return pool
-
-# initialize Connector as context manager
-with Connector() as connector:
-    # initialize connection pool
-    pool = init_connection_pool(connector)
     # insert statement
     insert_stmt = sqlalchemy.text(
         "INSERT INTO my_table (id, title) VALUES (:id, :title)",
@@ -258,30 +242,26 @@ import asyncpg
 from google.cloud.alloydb.connector import AsyncConnector
 
 async def main():
-    # initialize Connector object for connections to AlloyDB
+    # initialize AsyncConnector object for connections to AlloyDB
     connector = AsyncConnector()
 
-    # creation function to generate asyncpg connections as the 'connect' arg
-    async def getconn(instance_connection_name, **kwargs) -> asyncpg.Connection:
-        return await connector.connect(
+    # initialize asyncpg connection pool with AsyncConnector
+    pool = await asyncpg.create_pool(
+        "projects/<YOUR_PROJECT>/locations/<YOUR_REGION>/clusters/<YOUR_CLUSTER>/instances/<YOUR_INSTANCE>",
+        connect=lambda: connector.connect(
             instance_connection_name,
             "asyncpg",
             user="my-user",
             password="my-password",
             db="my-db",
-        )
-
-    # initialize connection pool
-    pool = await asyncpg.create_pool(
-        "projects/<YOUR_PROJECT>/locations/<YOUR_REGION>/clusters/<YOUR_CLUSTER>/instances/<YOUR_INSTANCE>",
-        connect=getconn,
+        ),
     )
 
     # acquire connection and query AlloyDB database
     async with pool.acquire() as conn:
         res = await conn.fetch("SELECT NOW()")
 
-    # close Connector
+    # close AsyncConnector
     await connector.close()
 ```
 

tests/system/test_asyncpg_connection.py

@@ -66,20 +66,16 @@ async def create_sqlalchemy_engine(
     """
     connector = AsyncConnector(refresh_strategy=refresh_strategy)
 
-    async def getconn() -> asyncpg.Connection:
-        conn: asyncpg.Connection = await connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.ext.asyncio.create_async_engine(
+        "postgresql+asyncpg://",
+        async_creator=lambda: connector.connect(
             inst_uri,
             "asyncpg",
             user=user,
             password=password,
             db=db,
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.ext.asyncio.create_async_engine(
-        "postgresql+asyncpg://",
-        async_creator=getconn,
+        ),
         execution_options={"isolation_level": "AUTOCOMMIT"},
     )
     return engine, connector
@@ -129,20 +125,15 @@ async def create_asyncpg_pool(
     """
     connector = AsyncConnector(refresh_strategy=refresh_strategy)
 
-    async def getconn(
-        instance_connection_name: str, **kwargs: Any
-    ) -> asyncpg.Connection:
-        conn: asyncpg.Connection = await connector.connect(
+    # create native asyncpg pool (requires asyncpg version >=0.30.0)
+    pool = await asyncpg.create_pool(instance_connection_name, connect=lambda: connector.connect(
             instance_connection_name,
             "asyncpg",
             user=user,
             password=password,
             db=db,
         )
-        return conn
-
-    # create native asyncpg pool (requires asyncpg version >=0.30.0)
-    pool = await asyncpg.create_pool(instance_connection_name, connect=getconn)
+    )
     return pool, connector
 
 

tests/system/test_asyncpg_iam_authn.py

@@ -16,7 +16,6 @@
 import os
 
 # [START alloydb_sqlalchemy_connect_async_connector_iam_authn]
-import asyncpg
 import sqlalchemy
 import sqlalchemy.ext.asyncio
 
@@ -61,20 +60,16 @@ async def create_sqlalchemy_engine(
     """
     connector = AsyncConnector(refresh_strategy=refresh_strategy)
 
-    async def getconn() -> asyncpg.Connection:
-        conn: asyncpg.Connection = await connector.connect(
+    # create async SQLAlchemy connection pool
+    engine = sqlalchemy.ext.asyncio.create_async_engine(
+        "postgresql+asyncpg://",
+        async_creator=lambda: connector.connect(
             inst_uri,
             "asyncpg",
             user=user,
             db=db,
             enable_iam_auth=True,
-        )
-        return conn
-
-    # create async SQLAlchemy connection pool
-    engine = sqlalchemy.ext.asyncio.create_async_engine(
-        "postgresql+asyncpg://",
-        async_creator=getconn,
+        ),
         execution_options={"isolation_level": "AUTOCOMMIT"},
     )
     return engine, connector

tests/system/test_asyncpg_psc.py

@@ -14,7 +14,6 @@
 
 import os
 
-import asyncpg
 import pytest
 import sqlalchemy
 import sqlalchemy.ext.asyncio
@@ -60,21 +59,17 @@ async def create_sqlalchemy_engine(
     """
     connector = AsyncConnector()
 
-    async def getconn() -> asyncpg.Connection:
-        conn: asyncpg.Connection = await connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.ext.asyncio.create_async_engine(
+        "postgresql+asyncpg://",
+        async_creator=lambda: connector.connect(
             inst_uri,
             "asyncpg",
             user=user,
             password=password,
             db=db,
             ip_type="PSC",
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.ext.asyncio.create_async_engine(
-        "postgresql+asyncpg://",
-        async_creator=getconn,
+        ),
         execution_options={"isolation_level": "AUTOCOMMIT"},
     )
     return engine, connector

tests/system/test_asyncpg_public_ip.py

@@ -15,7 +15,6 @@
 import os
 
 # [START alloydb_sqlalchemy_connect_async_connector_public_ip]
-import asyncpg
 import pytest
 import sqlalchemy
 import sqlalchemy.ext.asyncio
@@ -61,21 +60,17 @@ async def create_sqlalchemy_engine(
     """
     connector = AsyncConnector()
 
-    async def getconn() -> asyncpg.Connection:
-        conn: asyncpg.Connection = await connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.ext.asyncio.create_async_engine(
+        "postgresql+asyncpg://",
+        async_creator=lambda: connector.connect(
             inst_uri,
             "asyncpg",
             user=user,
             password=password,
             db=db,
             ip_type="PUBLIC",
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.ext.asyncio.create_async_engine(
-        "postgresql+asyncpg://",
-        async_creator=getconn,
+        ),
         execution_options={"isolation_level": "AUTOCOMMIT"},
     )
     return engine, connector

tests/system/test_pg8000_connection.py

@@ -16,7 +16,6 @@
 import os
 
 # [START alloydb_sqlalchemy_connect_connector]
-import pg8000
 import sqlalchemy
 
 from google.cloud.alloydb.connector import Connector
@@ -66,20 +65,16 @@ def create_sqlalchemy_engine(
     """
     connector = Connector(refresh_strategy=refresh_strategy)
 
-    def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=lambda: connector.connect(
             inst_uri,
             "pg8000",
             user=user,
             password=password,
             db=db,
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.create_engine(
-        "postgresql+pg8000://",
-        creator=getconn,
+        ),
     )
     engine.dialect.description_encoding = None
     return engine, connector

tests/system/test_pg8000_iam_authn.py

@@ -16,7 +16,6 @@
 import os
 
 # [START alloydb_sqlalchemy_connect_connector_iam_authn]
-import pg8000
 import sqlalchemy
 
 from google.cloud.alloydb.connector import Connector
@@ -60,20 +59,16 @@ def create_sqlalchemy_engine(
     """
     connector = Connector(refresh_strategy=refresh_strategy)
 
-    def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=lambda: connector.connect(
             inst_uri,
             "pg8000",
             user=user,
             db=db,
             enable_iam_auth=True,
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.create_engine(
-        "postgresql+pg8000://",
-        creator=getconn,
+        ),
     )
     return engine, connector
 

tests/system/test_pg8000_psc.py

@@ -15,7 +15,6 @@
 from datetime import datetime
 import os
 
-import pg8000
 import sqlalchemy
 
 from google.cloud.alloydb.connector import Connector
@@ -60,21 +59,17 @@ def create_sqlalchemy_engine(
     """
     connector = Connector()
 
-    def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=lambda: connector.connect(
             inst_uri,
             "pg8000",
             user=user,
             password=password,
             db=db,
             ip_type="PSC",
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.create_engine(
-        "postgresql+pg8000://",
-        creator=getconn,
+        ),
     )
     return engine, connector
 

tests/system/test_pg8000_public_ip.py

@@ -16,7 +16,6 @@
 import os
 
 # [START alloydb_sqlalchemy_connect_connector_public_ip]
-import pg8000
 import sqlalchemy
 
 from google.cloud.alloydb.connector import Connector
@@ -61,21 +60,17 @@ def create_sqlalchemy_engine(
     """
     connector = Connector()
 
-    def getconn() -> pg8000.dbapi.Connection:
-        conn: pg8000.dbapi.Connection = connector.connect(
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=lambda: connector.connect(
             inst_uri,
             "pg8000",
             user=user,
             password=password,
             db=db,
             ip_type="PUBLIC",
-        )
-        return conn
-
-    # create SQLAlchemy connection pool
-    engine = sqlalchemy.create_engine(
-        "postgresql+pg8000://",
-        creator=getconn,
+        ),
     )
     return engine, connector