feat: support basic user/password auth with MySQLEngine (#17)

Author: jackwotherspoon

integration.cloudbuild.yaml

@@ -28,6 +28,14 @@ steps:
       - 'DB_NAME=$_DB_NAME'
       - 'TABLE_NAME=test-$BUILD_ID'
       - 'REGION=$_REGION'
+    secretEnv: ['DB_USER', 'DB_PASSWORD']
+
+availableSecrets:
+  secretManager:
+  - versionName: projects/$PROJECT_ID/secrets/langchain-test-mysql-username/versions/1
+    env: 'DB_USER'
+  - versionName: projects/$PROJECT_ID/secrets/langchain-test-mysql-password/versions/1
+    env: 'DB_PASSWORD'
 
 substitutions:
   _INSTANCE_ID: test-instance

src/langchain_google_cloud_sql_mysql/mysql_engine.py

@@ -84,13 +84,17 @@ def from_instance(
         region: str,
         instance: str,
         database: str,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
     ) -> MySQLEngine:
         """Create an instance of MySQLEngine from Cloud SQL instance
         details.
 
         This method uses the Cloud SQL Python Connector to connect to Cloud SQL
         using automatic IAM database authentication with the Google ADC
-        credentials sourced from the environment.
+        credentials sourced from the environment by default. If user and
+        password arguments are given, basic database authentication will be
+        used for database login.
 
         More details can be found at https://github.com/GoogleCloudPlatform/cloud-sql-python-connector#credentials
 
@@ -101,42 +105,74 @@ def from_instance(
             instance (str): The name of the Cloud SQL instance.
             database (str): The name of the database to connect to on the
                 Cloud SQL instance.
+            user (str, optional): Database user to use for basic database
+                authentication and login. Defaults to None.
+            password (str, optional): Database password for 'user' to use for
+                basic database authentication and login. Defaults to None.
 
         Returns:
             (MySQLEngine): The engine configured to connect to a
                 Cloud SQL instance database.
         """
+        # error if only one of user or password is set, must be both or neither
+        if bool(user) ^ bool(password):
+            raise ValueError(
+                "Only one of 'user' or 'password' were specified. Either "
+                "both should be specified to use basic user/password "
+                "authentication or neither for IAM DB authentication."
+            )
         engine = cls._create_connector_engine(
             instance_connection_name=f"{project_id}:{region}:{instance}",
             database=database,
+            user=user,
+            password=password,
         )
         return cls(engine=engine)
 
     @classmethod
     def _create_connector_engine(
-        cls, instance_connection_name: str, database: str
+        cls,
+        instance_connection_name: str,
+        database: str,
+        user: Optional[str],
+        password: Optional[str],
     ) -> sqlalchemy.engine.Engine:
         """Create a SQLAlchemy engine using the Cloud SQL Python Connector.
 
         Defaults to use "pymysql" driver and to connect using automatic IAM
         database authentication with the IAM principal associated with the
-        environment's Google Application Default Credentials.
+        environment's Google Application Default Credentials. If user and
+        password arguments are given, basic database authentication will be
+        used for database login.
 
         Args:
             instance_connection_name (str): The instance connection
                 name of the Cloud SQL instance to establish a connection to.
                 (ex. "project-id:instance-region:instance-name")
             database (str): The name of the database to connect to on the
                 Cloud SQL instance.
+            user (str, optional): Database user to use for basic database
+                authentication and login. Defaults to None.
+            password (str, optional): Database password for 'user' to use for
+                basic database authentication and login. Defaults to None.
+
         Returns:
             (sqlalchemy.engine.Engine): Engine configured using the Cloud SQL
                 Python Connector.
         """
-        # get application default credentials
-        credentials, _ = google.auth.default(
-            scopes=["https://www.googleapis.com/auth/userinfo.email"]
-        )
-        iam_database_user = _get_iam_principal_email(credentials)
+        # if user and password are given, use basic auth
+        if user and password:
+            enable_iam_auth = False
+            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 = _get_iam_principal_email(credentials)
+            enable_iam_auth = True
+
         if cls._connector is None:
             cls._connector = Connector()
 
@@ -145,9 +181,10 @@ def getconn() -> pymysql.Connection:
             conn = cls._connector.connect(  # type: ignore
                 instance_connection_name,
                 "pymysql",
-                user=iam_database_user,
+                user=db_user,
+                password=password,
                 db=database,
-                enable_iam_auth=True,
+                enable_iam_auth=enable_iam_auth,
             )
             return conn
 

tests/integration/test_mysql_engine.py

@@ -0,0 +1,46 @@
+# Copyright 2024 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.
+import os
+
+import sqlalchemy
+
+from langchain_google_cloud_sql_mysql import MySQLEngine
+
+project_id = os.environ["PROJECT_ID"]
+region = os.environ["REGION"]
+instance_id = os.environ["INSTANCE_ID"]
+db_name = os.environ["DB_NAME"]
+db_user = os.environ["DB_USER"]
+db_password = os.environ["DB_PASSWORD"]
+
+
+def test_mysql_engine_with_basic_auth() -> None:
+    """Test MySQLEngine works with basic user/password auth."""
+    # override MySQLEngine._connector to allow a new Connector to be initiated
+    MySQLEngine._connector = None
+    engine = MySQLEngine.from_instance(
+        project_id=project_id,
+        region=region,
+        instance=instance_id,
+        database=db_name,
+        user=db_user,
+        password=db_password,
+    )
+    # test connection with query
+    with engine.connect() as conn:
+        res = conn.execute(sqlalchemy.text("SELECT 1")).fetchone()
+        conn.commit()
+        assert res[0] == 1  # type: ignore
+    # reset MySQLEngine._connector to allow a new Connector to be initiated
+    MySQLEngine._connector = None

tests/unit/test_engine.py

@@ -0,0 +1,49 @@
+# Copyright 2024 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.
+
+import pytest
+
+from langchain_google_cloud_sql_mysql import MySQLEngine
+
+
+def test_mysql_engine_with_invalid_arg_pattern() -> None:
+    """Test MySQLEngine errors when only one of user or password is given.
+
+    Both user and password must be specified (basic authentication)
+    or neither (IAM authentication).
+    """
+    expected_error_msg = "Only one of 'user' or 'password' were specified. Either both should be specified to use basic user/password authentication or neither for IAM DB authentication."
+    # test password not set
+    with pytest.raises(ValueError) as exc_info:
+        MySQLEngine.from_instance(
+            project_id="my-project",
+            region="my-region",
+            instance="my-instance",
+            database="my-db",
+            user="my-user",
+        )
+        # assert custom error is present
+        assert exc_info.value.args[0] == expected_error_msg
+
+    # test user not set
+    with pytest.raises(ValueError) as exc_info:
+        MySQLEngine.from_instance(
+            project_id="my-project",
+            region="my-region",
+            instance="my-instance",
+            database="my-db",
+            password="my-pass",
+        )
+        # assert custom error is present
+        assert exc_info.value.args[0] == expected_error_msg