chore: encapsulate driver and engine validation

GoogleCloudPlatform · Jun 3, 2024 · 0dd112b · 0dd112b
1 parent f59810c
commit 0dd112b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py
@@ -33,7 +33,6 @@
 from google.cloud.sql.connector.enums import DriverMapping
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
 from google.cloud.sql.connector.exceptions import DnsNameResolutionError
-from google.cloud.sql.connector.exceptions import IncompatibleDriverError
 from google.cloud.sql.connector.instance import IPTypes
 from google.cloud.sql.connector.instance import RefreshAheadCache
 from google.cloud.sql.connector.instance import RefreshStrategy
@@ -335,13 +334,7 @@ async def connect_async(
         try:
             conn_info = await cache.connect_info()
             # validate driver matches intended database engine
-            mapping = DriverMapping[driver.upper()]
-            if not conn_info.database_version.startswith(mapping.value):
-                raise IncompatibleDriverError(
-                    f"Database driver '{driver}' is incompatible with database "
-                    f"version '{conn_info.database_version}'. Given driver can "
-                    f"only be used with Cloud SQL {mapping.value} databases."
-                )
+            DriverMapping.validate_engine(driver, conn_info.database_version)
             ip_address = conn_info.get_preferred_ip(ip_type)
             # resolve DNS name into IP address for PSC
             if ip_type.value == "PSC":

diff --git a/google/cloud/sql/connector/enums.py b/google/cloud/sql/connector/enums.py
@@ -14,6 +14,8 @@
 
 from enum import Enum
 
+from google.cloud.sql.connector.exceptions import IncompatibleDriverError
+
 
 class DriverMapping(Enum):
     """Maps a given database driver to it's corresponding database engine."""
@@ -22,3 +24,23 @@ class DriverMapping(Enum):
     PG8000 = "POSTGRES"
     PYMYSQL = "MYSQL"
     PYTDS = "SQLSERVER"
+
+    @staticmethod
+    def validate_engine(driver: str, engine_version: str) -> None:
+        """Validate that the given driver is compatible with the given engine.
+
+        Args:
+            driver (str): Database driver being used. (i.e. "pg8000")
+            engine_version (str): Database engine version. (i.e. "POSTGRES_16")
+
+        Raises:
+            IncompatibleDriverError: If the given driver is not compatible with
+                the given engine.
+        """
+        mapping = DriverMapping[driver.upper()]
+        if not engine_version.startswith(mapping.value):
+            raise IncompatibleDriverError(
+                f"Database driver '{driver}' is incompatible with database "
+                f"version '{engine_version}'. Given driver can "
+                f"only be used with Cloud SQL {mapping.value} databases."
+            )