feat: Support IAM account override (#160)

* feat: Support IAM account override

* rename to `iam_account_email`

* add integration test for iam override

* add self

* add sync fixture

* clean up

Author: duwenxin99

integration.cloudbuild.yaml

@@ -32,14 +32,16 @@ steps:
       - "INSTANCE_ID=$_INSTANCE_ID"
       - "DATABASE_ID=$_DATABASE_ID"
       - "REGION=$_REGION"
-    secretEnv: ["DB_USER", "DB_PASSWORD"]
+    secretEnv: ["DB_USER", "DB_PASSWORD", "IAM_ACCOUNT"]
 
 availableSecrets:
   secretManager:
     - versionName: projects/$PROJECT_ID/secrets/langchain-test-pg-username/versions/1
       env: "DB_USER"
     - versionName: projects/$PROJECT_ID/secrets/langchain-test-pg-password/versions/1
       env: "DB_PASSWORD"
+    - versionName: projects/$PROJECT_ID/secrets/service_account_email/versions/1
+      env: "IAM_ACCOUNT"
 
 substitutions:
   _DATABASE_ID: test-database

src/langchain_google_cloud_sql_pg/engine.py

@@ -136,6 +136,7 @@ def from_instance(
         password: Optional[str] = None,
         ip_type: Union[str, IPTypes] = IPTypes.PUBLIC,
         quota_project: Optional[str] = None,
+        iam_account_email: Optional[str] = None,
     ) -> PostgresEngine:
         """Create a PostgresEngine from a Postgres instance.
 
@@ -169,6 +170,7 @@ def from_instance(
             loop=loop,
             thread=thread,
             quota_project=quota_project,
+            iam_account_email=iam_account_email,
         )
         return asyncio.run_coroutine_threadsafe(coro, loop).result()
 
@@ -185,6 +187,7 @@ async def _create(
         loop: Optional[asyncio.AbstractEventLoop] = None,
         thread: Optional[Thread] = None,
         quota_project: Optional[str] = None,
+        iam_account_email: Optional[str] = None,
     ) -> PostgresEngine:
         """Create a PostgresEngine instance.
 
@@ -227,12 +230,15 @@ async def _create(
             db_user = user
         # otherwise use automatic IAM database authentication
         else:
-            # get application default credentials
-            credentials, _ = google.auth.default(
-                scopes=["https://www.googleapis.com/auth/userinfo.email"]
-            )
-            db_user = await _get_iam_principal_email(credentials)
             enable_iam_auth = True
+            if iam_account_email:
+                db_user = iam_account_email
+            else:
+                # get application default credentials
+                credentials, _ = google.auth.default(
+                    scopes=["https://www.googleapis.com/auth/userinfo.email"]
+                )
+                db_user = await _get_iam_principal_email(credentials)
 
         # anonymous function to be used for SQLAlchemy 'creator' argument
         async def getconn() -> asyncpg.Connection:
@@ -264,6 +270,7 @@ async def afrom_instance(
         password: Optional[str] = None,
         ip_type: Union[str, IPTypes] = IPTypes.PUBLIC,
         quota_project: Optional[str] = None,
+        iam_account_email: Optional[str] = None,
     ) -> PostgresEngine:
         """Create a PostgresEngine from a Postgres instance.
 
@@ -290,6 +297,7 @@ async def afrom_instance(
             user,
             password,
             quota_project=quota_project,
+            iam_account_email=iam_account_email,
         )
 
     @classmethod

tests/test_postgresql_engine.py

@@ -65,6 +65,10 @@ def user(self) -> str:
     def password(self) -> str:
         return get_env_var("DB_PASSWORD", "database password for cloud sql")
 
+    @pytest.fixture(scope="module")
+    def iam_account(self) -> str:
+        return get_env_var("IAM_ACCOUNT", "Cloud SQL IAM account email")
+
     @pytest_asyncio.fixture
     async def engine(self, db_project, db_region, db_instance, db_name):
         engine = await PostgresEngine.afrom_instance(
@@ -176,6 +180,26 @@ async def test_column(self, engine):
         with pytest.raises(ValueError):
             Column(1, "INTEGER")
 
+    async def test_iam_account_override(
+        self,
+        db_project,
+        db_instance,
+        db_region,
+        db_name,
+        iam_account,
+    ):
+        engine = await PostgresEngine.afrom_instance(
+            project_id=db_project,
+            instance=db_instance,
+            region=db_region,
+            database=db_name,
+            iam_account_email=iam_account,
+        )
+        assert engine
+        await engine._aexecute("SELECT 1")
+        await engine._connector.close_async()
+        await engine._engine.dispose()
+
 
 @pytest.mark.asyncio
 class TestEngineSync:
@@ -203,6 +227,10 @@ def user(self) -> str:
     def password(self) -> str:
         return get_env_var("DB_PASSWORD", "database password for cloud sql")
 
+    @pytest.fixture(scope="module")
+    def iam_account(self) -> str:
+        return get_env_var("IAM_ACCOUNT", "Cloud SQL IAM account email")
+
     @pytest_asyncio.fixture
     def engine(self, db_project, db_region, db_instance, db_name):
         engine = PostgresEngine.from_instance(
@@ -285,3 +313,23 @@ async def test_engine_constructor_key(
         key = object()
         with pytest.raises(Exception):
             PostgresEngine(key, engine)
+
+    def test_iam_account_override(
+        self,
+        db_project,
+        db_instance,
+        db_region,
+        db_name,
+        iam_account,
+    ):
+        engine = PostgresEngine.from_instance(
+            project_id=db_project,
+            instance=db_instance,
+            region=db_region,
+            database=db_name,
+            iam_account_email=iam_account,
+        )
+        assert engine
+        engine._execute("SELECT 1")
+        engine._connector.close()
+        engine._engine.dispose()