From 31fb33b36ba162bd9d7ef9a7c61df8e4108dd95a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Tue, 4 Nov 2025 04:18:48 +0000 Subject: [PATCH 1/9] feat: add support for Google Cloud SQL and AlloyDB connectors Automatic support for the CloudSQL/AlloyDB asyncpg driver integrations. --- AGENTS.md | 315 ++++++++++ docs/guides/adapters/asyncpg.md | 131 ++++- docs/guides/cloud/google-connectors.md | 500 ++++++++++++++++ pyproject.toml | 2 + sqlspec/_typing.py | 12 + sqlspec/adapters/asyncpg/config.py | 176 +++++- sqlspec/typing.py | 4 + .../test_cloud_connectors_integration.py | 205 +++++++ .../test_asyncpg/test_cloud_connectors.py | 544 ++++++++++++++++++ uv.lock | 86 ++- 10 files changed, 1957 insertions(+), 18 deletions(-) create mode 100644 docs/guides/cloud/google-connectors.md create mode 100644 tests/integration/test_adapters/test_asyncpg/test_cloud_connectors_integration.py create mode 100644 tests/unit/test_adapters/test_asyncpg/test_cloud_connectors.py diff --git a/AGENTS.md b/AGENTS.md index 1bd927d3..5aa75303 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1912,6 +1912,321 @@ Current state of all adapters (as of type-cleanup branch): - **Excellent**: Follows all patterns, well documented - **Good**: Follows patterns appropriately for adapter's needs +## Google Cloud Connector Pattern + +### Overview + +Google Cloud SQL and AlloyDB connectors provide automatic IAM authentication, SSL management, and IP routing for PostgreSQL databases hosted on Google Cloud Platform. SQLSpec integrates these connectors through the AsyncPG adapter using a connection factory pattern. + +### When to Use This Pattern + +Use Google Cloud connectors when: + +- Connecting to Cloud SQL for PostgreSQL instances +- Connecting to AlloyDB for PostgreSQL clusters +- Need automatic IAM authentication +- Want managed SSL/TLS connections +- Require private IP or PSC connectivity + +### Implementation Pattern + +#### Step 1: Add Optional Dependencies + +Add connector packages as optional dependency groups in pyproject.toml: + +```toml +[project.optional-dependencies] +cloud-sql = ["cloud-sql-python-connector[asyncpg]"] +alloydb = ["cloud-alloydb-python-connector[asyncpg]"] +``` + +#### Step 2: Add Detection Constants + +In sqlspec/_typing.py: + +```python +try: + import google.cloud.sql.connector + CLOUD_SQL_CONNECTOR_INSTALLED = True +except ImportError: + CLOUD_SQL_CONNECTOR_INSTALLED = False + +try: + import google.cloud.alloydb.connector + ALLOYDB_CONNECTOR_INSTALLED = True +except ImportError: + ALLOYDB_CONNECTOR_INSTALLED = False +``` + +Re-export in sqlspec/typing.py and add to **all**. + +#### Step 3: Update Driver Features TypedDict + +Document all connector options with comprehensive descriptions: + +```python +class AsyncpgDriverFeatures(TypedDict): + """AsyncPG driver feature flags.""" + + enable_cloud_sql: NotRequired[bool] + """Enable Google Cloud SQL connector integration. + Requires cloud-sql-python-connector package. + Defaults to True when package is installed. + Auto-configures IAM authentication, SSL, and IP routing. + Mutually exclusive with enable_alloydb. + """ + + cloud_sql_instance: NotRequired[str] + """Cloud SQL instance connection name. + Format: "project:region:instance" + Required when enable_cloud_sql is True. + """ + + cloud_sql_enable_iam_auth: NotRequired[bool] + """Enable IAM database authentication. + Defaults to False for passwordless authentication. + When False, requires user/password in pool_config. + """ + + cloud_sql_ip_type: NotRequired[str] + """IP address type for connection. + Options: "PUBLIC", "PRIVATE", "PSC" + Defaults to "PRIVATE". + """ + + enable_alloydb: NotRequired[bool] + """Enable Google AlloyDB connector integration. + Requires cloud-alloydb-python-connector package. + Defaults to True when package is installed. + Auto-configures IAM authentication and private networking. + Mutually exclusive with enable_cloud_sql. + """ + + alloydb_instance_uri: NotRequired[str] + """AlloyDB instance URI. + Format: "projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE" + Required when enable_alloydb is True. + """ + + alloydb_enable_iam_auth: NotRequired[bool] + """Enable IAM database authentication. + Defaults to False for passwordless authentication. + """ + + alloydb_ip_type: NotRequired[str] + """IP address type for connection. + Options: "PUBLIC", "PRIVATE", "PSC" + Defaults to "PRIVATE". + """ +``` + +#### Step 4: Add Auto-Detection to Config Init + +```python +class AsyncpgConfig(AsyncDatabaseConfig): + def __init__(self, *, driver_features=None, **kwargs): + features_dict = dict(driver_features) if driver_features else {} + + features_dict.setdefault("enable_cloud_sql", CLOUD_SQL_CONNECTOR_INSTALLED) + features_dict.setdefault("enable_alloydb", ALLOYDB_CONNECTOR_INSTALLED) + + super().__init__(driver_features=features_dict, **kwargs) + + self._cloud_sql_connector = None + self._alloydb_connector = None + + self._validate_connector_config() +``` + +#### Step 5: Add Configuration Validation + +```python +def _validate_connector_config(self) -> None: + """Validate Google Cloud connector configuration.""" + enable_cloud_sql = self.driver_features.get("enable_cloud_sql", False) + enable_alloydb = self.driver_features.get("enable_alloydb", False) + + if enable_cloud_sql and enable_alloydb: + msg = "Cannot enable both Cloud SQL and AlloyDB connectors simultaneously. Use separate configs for each database." + raise ImproperConfigurationError(msg) + + if enable_cloud_sql: + if not CLOUD_SQL_CONNECTOR_INSTALLED: + msg = "cloud-sql-python-connector package not installed. Install with: pip install cloud-sql-python-connector" + raise ImproperConfigurationError(msg) + + instance = self.driver_features.get("cloud_sql_instance") + if not instance: + msg = "cloud_sql_instance required when enable_cloud_sql is True. Format: 'project:region:instance'" + raise ImproperConfigurationError(msg) + + cloud_sql_instance_parts_expected = 2 + if instance.count(":") != cloud_sql_instance_parts_expected: + msg = f"Invalid Cloud SQL instance format: {instance}. Expected format: 'project:region:instance'" + raise ImproperConfigurationError(msg) + + elif enable_alloydb: + if not ALLOYDB_CONNECTOR_INSTALLED: + msg = "cloud-alloydb-python-connector package not installed. Install with: pip install cloud-alloydb-python-connector" + raise ImproperConfigurationError(msg) + + instance_uri = self.driver_features.get("alloydb_instance_uri") + if not instance_uri: + msg = "alloydb_instance_uri required when enable_alloydb is True. Format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" + raise ImproperConfigurationError(msg) + + if not instance_uri.startswith("projects/"): + msg = f"Invalid AlloyDB instance URI format: {instance_uri}. Expected format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" + raise ImproperConfigurationError(msg) +``` + +#### Step 6: Implement Connection Factory Pattern + +Extract connector setup into private helper methods: + +```python +def _setup_cloud_sql_connector(self, config: dict[str, Any]) -> None: + """Setup Cloud SQL connector and configure pool for connection factory pattern.""" + from google.cloud.sql.connector import Connector + + self._cloud_sql_connector = Connector() + + user = config.get("user") + password = config.get("password") + database = config.get("database") + + async def get_conn() -> AsyncpgConnection: + conn_kwargs = { + "instance_connection_string": self.driver_features["cloud_sql_instance"], + "driver": "asyncpg", + "enable_iam_auth": self.driver_features.get("cloud_sql_enable_iam_auth", False), + "ip_type": self.driver_features.get("cloud_sql_ip_type", "PRIVATE"), + } + + if user: + conn_kwargs["user"] = user + if password: + conn_kwargs["password"] = password + if database: + conn_kwargs["db"] = database + + return await self._cloud_sql_connector.connect_async(**conn_kwargs) + + for key in ("dsn", "host", "port", "user", "password", "database"): + config.pop(key, None) + + config["connect"] = get_conn + + +def _setup_alloydb_connector(self, config: dict[str, Any]) -> None: + """Setup AlloyDB connector and configure pool for connection factory pattern.""" + from google.cloud.alloydb.connector import AsyncConnector + + self._alloydb_connector = AsyncConnector() + + user = config.get("user") + password = config.get("password") + database = config.get("database") + + async def get_conn() -> AsyncpgConnection: + conn_kwargs = { + "instance_uri": self.driver_features["alloydb_instance_uri"], + "driver": "asyncpg", + "enable_iam_auth": self.driver_features.get("alloydb_enable_iam_auth", False), + "ip_type": self.driver_features.get("alloydb_ip_type", "PRIVATE"), + } + + if user: + conn_kwargs["user"] = user + if password: + conn_kwargs["password"] = password + if database: + conn_kwargs["db"] = database + + return await self._alloydb_connector.connect(**conn_kwargs) + + for key in ("dsn", "host", "port", "user", "password", "database"): + config.pop(key, None) + + config["connect"] = get_conn +``` + +#### Step 7: Use in Pool Creation + +```python +async def _create_pool(self) -> Pool[Record]: + config = self._get_pool_config_dict() + + if self.driver_features.get("enable_cloud_sql", False): + self._setup_cloud_sql_connector(config) + elif self.driver_features.get("enable_alloydb", False): + self._setup_alloydb_connector(config) + + if "init" not in config: + config["init"] = self._init_connection + + return await asyncpg_create_pool(**config) +``` + +#### Step 8: Cleanup Connectors + +```python +async def _close_pool(self) -> None: + if self.pool_instance: + await self.pool_instance.close() + + if self._cloud_sql_connector is not None: + await self._cloud_sql_connector.close_async() + self._cloud_sql_connector = None + + if self._alloydb_connector is not None: + await self._alloydb_connector.close() + self._alloydb_connector = None +``` + +### Key Design Principles + +1. **Auto-Detection**: Default to package installation status +2. **Mutual Exclusion**: Cannot enable both connectors simultaneously +3. **Connection Factory Pattern**: Use driver's `connect` parameter +4. **Clean Helper Methods**: Extract setup logic for maintainability +5. **Proper Lifecycle**: Initialize in create_pool, cleanup in close_pool +6. **Clear Validation**: Validate instance names, package installation, config +7. **Comprehensive TypedDict**: Document all options inline + +### Testing Requirements + +- Unit tests with mocked connectors +- Integration tests with real instances (conditional) +- Test auto-detection with both packages installed/not installed +- Test mutual exclusion validation +- Test connection factory pattern integration +- Test lifecycle (initialization and cleanup) +- Test all IP types and auth modes + +### Driver Compatibility + +| Driver | Cloud SQL | AlloyDB | Notes | +|--------|-----------|---------|-------| +| AsyncPG | ✅ Full | ✅ Full | Connection factory pattern via `connect` param | +| Psycopg | ⚠️ Research | ⚠️ Research | Not officially documented, needs prototype | +| Psqlpy | ❌ No | ❌ No | Internal Rust driver, architecturally incompatible | +| ADBC | ❌ No | ❌ No | URI-only interface, no factory pattern support | + +### Examples from Existing Implementations + +See sqlspec/adapters/asyncpg/config.py for the reference implementation. + +### Documentation Requirements + +When implementing cloud connector support: + +1. **Update adapter guide** - Add cloud integration section with examples +2. **Create cloud connector guide** - Comprehensive configuration reference +3. **Document limitations** - Clearly state unsupported drivers +4. **Provide troubleshooting** - Common errors and solutions +5. **Include migration guide** - From direct DSN to connector pattern + ### Testing Requirements When implementing `driver_features`, you MUST test: diff --git a/docs/guides/adapters/asyncpg.md b/docs/guides/adapters/asyncpg.md index ea8bc5a5..747ae08c 100644 --- a/docs/guides/adapters/asyncpg.md +++ b/docs/guides/adapters/asyncpg.md @@ -8,27 +8,134 @@ This guide provides specific instructions and best practices for working with th ## Key Information -- **Driver:** `asyncpg` -- **Parameter Style:** `numeric` (e.g., `$1, $2`) +- **Driver:** `asyncpg` +- **Parameter Style:** `numeric` (e.g., `$1, $2`) ## Parameter Profile -- **Registry Key:** `"asyncpg"` -- **JSON Strategy:** `driver` (delegates JSON binding to asyncpg codecs) -- **Extras:** None (codecs registered through config init hook) +- **Registry Key:** `"asyncpg"` +- **JSON Strategy:** `driver` (delegates JSON binding to asyncpg codecs) +- **Extras:** None (codecs registered through config init hook) ## Best Practices -- **High-Performance:** `asyncpg` is often chosen for high-performance applications due to its speed. It's a good choice for applications with a high volume of database traffic. +- **High-Performance:** `asyncpg` is often chosen for high-performance applications due to its speed. It's a good choice for applications with a high volume of database traffic. ## Driver Features The `asyncpg` adapter supports the following driver features: -- `json_serializer`: A function to serialize Python objects to JSON. Defaults to `sqlspec.utils.serializers.to_json`. -- `json_deserializer`: A function to deserialize JSON strings to Python objects. Defaults to `sqlspec.utils.serializers.from_json`. -- `enable_json_codecs`: A boolean to enable or disable automatic JSON/JSONB codec registration. Defaults to `True`. -- `enable_pgvector`: A boolean to enable or disable `pgvector` support. Defaults to `True` if `pgvector` is installed. +- `json_serializer`: A function to serialize Python objects to JSON. Defaults to `sqlspec.utils.serializers.to_json`. +- `json_deserializer`: A function to deserialize JSON strings to Python objects. Defaults to `sqlspec.utils.serializers.from_json`. +- `enable_json_codecs`: A boolean to enable or disable automatic JSON/JSONB codec registration. Defaults to `True`. +- `enable_pgvector`: A boolean to enable or disable `pgvector` support. Defaults to `True` if `pgvector` is installed. +- `enable_cloud_sql`: Enable Google Cloud SQL connector integration. Defaults to `True` when `cloud-sql-python-connector` is installed. +- `cloud_sql_instance`: Cloud SQL instance connection name (format: `"project:region:instance"`). Required when `enable_cloud_sql` is `True`. +- `cloud_sql_enable_iam_auth`: Enable IAM database authentication for Cloud SQL. Defaults to `False`. +- `cloud_sql_ip_type`: IP address type for Cloud SQL connection (`"PUBLIC"`, `"PRIVATE"`, or `"PSC"`). Defaults to `"PRIVATE"`. +- `enable_alloydb`: Enable Google AlloyDB connector integration. Defaults to `True` when `cloud-alloydb-python-connector` is installed. +- `alloydb_instance_uri`: AlloyDB instance URI (format: `"projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE"`). Required when `enable_alloydb` is `True`. +- `alloydb_enable_iam_auth`: Enable IAM database authentication for AlloyDB. Defaults to `False`. +- `alloydb_ip_type`: IP address type for AlloyDB connection (`"PUBLIC"`, `"PRIVATE"`, or `"PSC"`). Defaults to `"PRIVATE"`. + +## Google Cloud Integration + +AsyncPG supports native integration with Google Cloud SQL and AlloyDB connectors for simplified authentication and connection management. + +### Cloud SQL Connector + +Connect to Cloud SQL PostgreSQL instances with automatic SSL and IAM authentication: + +```python +from sqlspec import SQLSpec +from sqlspec.adapters.asyncpg import AsyncpgConfig + +db_manager = SQLSpec() + +# IAM authentication (no password required) +db = db_manager.add_config(AsyncpgConfig( + pool_config={ + "user": "my-service-account@project.iam", + "database": "mydb", + "min_size": 2, + "max_size": 10, + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_enable_iam_auth": True, + "cloud_sql_ip_type": "PRIVATE", + } +)) + +async with db_manager.provide_session(db) as session: + result = await session.select_one("SELECT current_user, version()") + print(result) +``` + +Password authentication is also supported: + +```python +config = AsyncpgConfig( + pool_config={ + "user": "postgres", + "password": "secret", + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_ip_type": "PUBLIC", # Public IP for external access + } +) +``` + +### AlloyDB Connector + +Connect to AlloyDB instances with the same pattern: + +```python +# IAM authentication +config = AsyncpgConfig( + pool_config={ + "user": "my-service-account@project.iam", + "database": "mydb", + }, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance", + "alloydb_enable_iam_auth": True, + "alloydb_ip_type": "PRIVATE", + } +) +``` + +### Configuration Notes + +**Auto-Detection**: Both connectors are automatically enabled when their respective packages are installed: + +```bash +# Install Cloud SQL connector +pip install cloud-sql-python-connector + +# Install AlloyDB connector +pip install cloud-alloydb-python-connector +``` + +**Mutual Exclusion**: A single config can only use one connector (Cloud SQL or AlloyDB). For multiple databases, create separate configs with unique `bind_key` values. + +**IP Type Selection**: + +- `"PRIVATE"` (default): Connect via private VPC network +- `"PUBLIC"`: Connect via public IP address +- `"PSC"`: Connect via Private Service Connect (AlloyDB only) + +**Authentication Methods**: + +- IAM authentication: Set `cloud_sql_enable_iam_auth=True` or `alloydb_enable_iam_auth=True` +- Password authentication: Leave IAM flags as `False` (default) and provide password in `pool_config` + +For comprehensive configuration options and troubleshooting, see the [Google Cloud Connectors Guide](/guides/cloud/google-connectors.md). ## MERGE Operations (PostgreSQL 15+) @@ -112,5 +219,5 @@ For comprehensive examples and migration guides, see: ## Common Issues -- **`asyncpg.exceptions.PostgresSyntaxError`**: Check your SQL syntax and parameter styles. `asyncpg` uses the `$` numeric style for parameters. -- **Connection Pooling:** `asyncpg` has its own connection pool implementation. Ensure the pool settings in the `sqlspec` config are appropriate for the application's needs (e.g., `min_size`, `max_size`). +- **`asyncpg.exceptions.PostgresSyntaxError`**: Check your SQL syntax and parameter styles. `asyncpg` uses the `$` numeric style for parameters. +- **Connection Pooling:** `asyncpg` has its own connection pool implementation. Ensure the pool settings in the `sqlspec` config are appropriate for the application's needs (e.g., `min_size`, `max_size`). diff --git a/docs/guides/cloud/google-connectors.md b/docs/guides/cloud/google-connectors.md new file mode 100644 index 00000000..63d01d2f --- /dev/null +++ b/docs/guides/cloud/google-connectors.md @@ -0,0 +1,500 @@ +# Google Cloud SQL and AlloyDB Connector Guide + +This guide covers integration of Google Cloud SQL and AlloyDB connectors with SQLSpec, providing simplified authentication, automatic SSL configuration, and streamlined connection management for Google Cloud managed PostgreSQL instances. + +## Overview + +Google Cloud provides official Python connectors for Cloud SQL and AlloyDB that handle authentication, SSL encryption, and network routing automatically. SQLSpec integrates these connectors through the AsyncPG adapter using a connection factory pattern. + +### When to Use Cloud Connectors + +Use Cloud SQL or AlloyDB connectors when: + +- Deploying applications on Google Cloud Platform (GCP) +- Requiring IAM-based database authentication +- Managing SSL certificates automatically +- Connecting to private IP instances from GCP resources +- Simplifying credential rotation and management + +### Supported Adapters + +- **AsyncPG**: Full support (recommended) +- **Psycopg**: Not officially supported by Google connectors +- **Psqlpy**: Architecturally incompatible (internal Rust driver) +- **ADBC**: Incompatible (URI-only interface) + +For unsupported adapters, use [Cloud SQL Auth Proxy](https://cloud.google.com/sql/docs/postgres/sql-proxy) as an alternative. + +## Quick Start + +### Installation + +Install the connector packages alongside SQLSpec: + +```bash +# Cloud SQL connector +pip install sqlspec[asyncpg] cloud-sql-python-connector + +# AlloyDB connector +pip install sqlspec[asyncpg] cloud-alloydb-python-connector + +# Both connectors +pip install sqlspec[asyncpg] cloud-sql-python-connector cloud-alloydb-python-connector +``` + +### Basic Cloud SQL Connection + +```python +from sqlspec import SQLSpec +from sqlspec.adapters.asyncpg import AsyncpgConfig + +sql = SQLSpec() + +config = AsyncpgConfig( + pool_config={ + "user": "postgres", + "password": "secret", + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + } +) +sql.add_config(config) + +async with sql.provide_session(config) as session: + users = await session.select_all("SELECT * FROM users") + print(users) +``` + +### Basic AlloyDB Connection + +```python +config = AsyncpgConfig( + pool_config={ + "user": "postgres", + "password": "secret", + "database": "mydb", + }, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance", + } +) +``` + +## Configuration Reference + +### Cloud SQL Driver Features + +All Cloud SQL configuration is specified in `driver_features`: + +```python +driver_features = { + "enable_cloud_sql": True, # Auto-detected when package installed + "cloud_sql_instance": "project:region:instance", # Required + "cloud_sql_enable_iam_auth": False, # Optional (default: False) + "cloud_sql_ip_type": "PRIVATE", # Optional (default: "PRIVATE") +} +``` + +**Field Descriptions**: + +- `enable_cloud_sql`: Enable Cloud SQL connector integration. Defaults to `True` when `cloud-sql-python-connector` is installed. +- `cloud_sql_instance`: Instance connection name in format `"project:region:instance"`. Required when connector is enabled. +- `cloud_sql_enable_iam_auth`: Use IAM authentication instead of password. Defaults to `False`. +- `cloud_sql_ip_type`: IP address type for connection. Options: `"PUBLIC"`, `"PRIVATE"`, or `"PSC"`. Defaults to `"PRIVATE"`. + +### AlloyDB Driver Features + +All AlloyDB configuration is specified in `driver_features`: + +```python +driver_features = { + "enable_alloydb": True, # Auto-detected when package installed + "alloydb_instance_uri": "projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE", # Required + "alloydb_enable_iam_auth": False, # Optional (default: False) + "alloydb_ip_type": "PRIVATE", # Optional (default: "PRIVATE") +} +``` + +**Field Descriptions**: + +- `enable_alloydb`: Enable AlloyDB connector integration. Defaults to `True` when `cloud-alloydb-python-connector` is installed. +- `alloydb_instance_uri`: Instance URI in full resource path format. Required when connector is enabled. +- `alloydb_enable_iam_auth`: Use IAM authentication instead of password. Defaults to `False`. +- `alloydb_ip_type`: IP address type for connection. Options: `"PUBLIC"`, `"PRIVATE"`, or `"PSC"`. Defaults to `"PRIVATE"`. + +## Authentication Methods + +### Application Default Credentials (ADC) + +The simplest method for GCP deployments. Connectors automatically use credentials from: + +1. Environment variable `GOOGLE_APPLICATION_CREDENTIALS` +2. Cloud Run / Cloud Functions / GKE service account +3. Compute Engine default service account +4. gcloud CLI credentials (local development) + +```python +# No explicit credentials needed - ADC handles it +config = AsyncpgConfig( + pool_config={ + "user": "postgres", + "password": "secret", + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + } +) +``` + +### IAM Database Authentication + +Passwordless authentication using Google Cloud IAM: + +```python +config = AsyncpgConfig( + pool_config={ + "user": "my-service-account@my-project.iam", # IAM principal + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_enable_iam_auth": True, + } +) +``` + +**Requirements**: + +1. Cloud SQL instance has IAM authentication enabled +2. IAM principal has `roles/cloudsql.client` role +3. Database user exists and is granted IAM authentication + +**Setup Commands**: + +```sql +-- Create IAM user in database +CREATE ROLE "my-service-account@my-project.iam" WITH LOGIN; +GRANT ALL ON DATABASE mydb TO "my-service-account@my-project.iam"; +``` + +### Password Authentication + +Traditional username/password authentication: + +```python +config = AsyncpgConfig( + pool_config={ + "user": "postgres", + "password": "secret", # From Secret Manager recommended + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_enable_iam_auth": False, # Explicit (default) + } +) +``` + +## IP Type Selection + +### Private IP (Default) + +Connect via VPC private network: + +```python +driver_features = { + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_ip_type": "PRIVATE", +} +``` + +**When to use**: + +- Application runs on GCP (Cloud Run, GKE, Compute Engine) +- VPC peering configured +- Security requirement to avoid public internet + +**Requirements**: + +- Cloud SQL instance has private IP enabled +- VPC network properly configured +- Serverless VPC Access connector (for Cloud Run/Functions) + +### Public IP + +Connect via public internet: + +```python +driver_features = { + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + "cloud_sql_ip_type": "PUBLIC", +} +``` + +**When to use**: + +- Development from local machine +- External services outside GCP +- No VPC networking configured + +**Security**: Traffic is SSL encrypted but uses public IP. Configure authorized networks in Cloud SQL for additional security. + +### Private Service Connect (PSC) + +AlloyDB-specific option for enhanced security: + +```python +driver_features = { + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + "alloydb_ip_type": "PSC", +} +``` + +**When to use**: + +- AlloyDB only (not available for Cloud SQL) +- Strict security and compliance requirements +- Dedicated network isolation + +## Multi-Database Configuration + +Use separate configs with unique `bind_key` values: + +```python +from sqlspec import SQLSpec +from sqlspec.adapters.asyncpg import AsyncpgConfig + +sql = SQLSpec() + +# Cloud SQL production database +cloud_sql_config = AsyncpgConfig( + pool_config={"user": "app", "password": "secret", "database": "prod"}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "prod-project:us-central1:prod-db", + }, + bind_key="cloud_sql" +) +sql.add_config(cloud_sql_config) + +# AlloyDB analytics database +alloydb_config = AsyncpgConfig( + pool_config={"user": "analytics", "password": "secret", "database": "warehouse"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/analytics/locations/us-central1/clusters/warehouse/instances/primary", + }, + bind_key="alloydb" +) +sql.add_config(alloydb_config) + +# Use different databases +async with sql.provide_session(cloud_sql_config) as session: + users = await session.select_all("SELECT * FROM users") + +async with sql.provide_session(alloydb_config) as session: + analytics = await session.select_all("SELECT * FROM events") +``` + +## Migration from Direct Connections + +### Before (Direct DSN) + +```python +config = AsyncpgConfig( + pool_config={ + "dsn": "postgresql://user:pass@10.0.0.5:5432/mydb", + "ssl": ssl_context, # Manual SSL setup + } +) +``` + +### After (Cloud SQL Connector) + +```python +config = AsyncpgConfig( + pool_config={ + "user": "user", + "password": "pass", + "database": "mydb", + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + } +) +``` + +**Benefits**: + +- No manual SSL certificate management +- Automatic credential rotation (with IAM auth) +- Simplified configuration +- Built-in connection retry logic + +## Troubleshooting + +### Error: "Cannot enable both Cloud SQL and AlloyDB connectors simultaneously" + +**Cause**: Single config has both `enable_cloud_sql=True` and `enable_alloydb=True`. + +**Solution**: Use separate configs with unique `bind_key` values for each database. + +### Error: "cloud_sql_instance required when enable_cloud_sql is True" + +**Cause**: Connector enabled but instance name not provided. + +**Solution**: Add instance name to `driver_features`: + +```python +driver_features = { + "enable_cloud_sql": True, + "cloud_sql_instance": "project:region:instance", +} +``` + +### Error: "Invalid Cloud SQL instance format" + +**Cause**: Instance name doesn't match `"project:region:instance"` format. + +**Solution**: Verify instance connection name from Cloud Console: + +```bash +gcloud sql instances describe INSTANCE_NAME --format="value(connectionName)" +``` + +### Error: "cloud-sql-python-connector package not installed" + +**Cause**: Connector package missing. + +**Solution**: Install the connector: + +```bash +pip install cloud-sql-python-connector +``` + +### IAM Authentication Fails + +**Common causes**: + +1. Database user not created with IAM authentication +2. IAM principal missing `roles/cloudsql.client` role +3. Cloud SQL instance doesn't have IAM authentication enabled + +**Solutions**: + +```bash +# Grant IAM role +gcloud projects add-iam-policy-binding PROJECT_ID \ + --member="serviceAccount:SERVICE_ACCOUNT" \ + --role="roles/cloudsql.client" + +# Enable IAM on instance +gcloud sql instances patch INSTANCE_NAME --database-flags=cloudsql.iam_authentication=on + +# Create IAM user in database +psql -h INSTANCE_IP -U postgres -d DATABASE +CREATE ROLE "SERVICE_ACCOUNT@PROJECT.iam" WITH LOGIN; +GRANT ALL ON DATABASE mydb TO "SERVICE_ACCOUNT@PROJECT.iam"; +``` + +### Private IP Connection Timeout + +**Common causes**: + +1. VPC peering not configured +2. Serverless VPC Access connector missing (Cloud Run/Functions) +3. Wrong IP type selected + +**Solutions**: + +- Verify VPC network configuration +- Create Serverless VPC Access connector for serverless environments +- Use `"PUBLIC"` IP type for testing (not recommended for production) + +## Performance Considerations + +### Connection Pool Configuration + +Connectors work seamlessly with AsyncPG connection pooling: + +```python +config = AsyncpgConfig( + pool_config={ + "min_size": 2, + "max_size": 10, + "max_inactive_connection_lifetime": 300, + }, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "my-project:us-central1:my-instance", + } +) +``` + +**Recommendations**: + +- Start with `min_size=2` and `max_size=10` for most applications +- Increase `max_size` for high-traffic applications +- Set `max_inactive_connection_lifetime` to match Cloud SQL timeout settings + +### Connection Overhead + +Initial connection creation is slightly slower due to SSL handshake and authentication: + +- Direct DSN: ~50-100ms +- Cloud SQL connector: ~100-200ms +- AlloyDB connector: ~100-200ms + +Connection pooling amortizes this overhead across many queries. + +## Security Best Practices + +1. **Use IAM Authentication**: Eliminates password management and enables automatic credential rotation +2. **Prefer Private IP**: Keeps traffic within GCP network +3. **Store Secrets Securely**: Use Secret Manager for passwords (when not using IAM) +4. **Least Privilege**: Grant minimal database permissions to IAM principals +5. **Enable Audit Logging**: Monitor database access with Cloud Audit Logs +6. **Regular Updates**: Keep connector packages updated for security patches + +## Limitations + +### Unsupported Adapters + +**Psqlpy**: Architecturally incompatible due to internal Rust driver. Use AsyncPG instead or Cloud SQL Auth Proxy for direct connections. + +**ADBC**: URI-only interface incompatible with connection factory pattern. Use AsyncPG with `select_to_arrow()` for Arrow integration, or Cloud SQL Auth Proxy. + +**Psycopg**: Not officially supported by Google Cloud connectors. GitHub issue tracking psycopg3 support has been open since 2021. Use AsyncPG as alternative. + +### Alternative: Cloud SQL Auth Proxy + +For unsupported adapters, use [Cloud SQL Auth Proxy](https://cloud.google.com/sql/docs/postgres/sql-proxy): + +```bash +# Start proxy +cloud-sql-proxy my-project:us-central1:my-instance --port 5432 + +# Connect with any adapter +config = PsqlpyConfig( # Or ADBC, psycopg, etc. + pool_config={"dsn": "postgresql://localhost:5432/mydb"} +) +``` + +## References + +- [Cloud SQL Python Connector Documentation](https://cloud.google.com/sql/docs/postgres/connect-connectors) +- [AlloyDB Python Connector Documentation](https://cloud.google.com/alloydb/docs/auth-proxy/connect) +- [AsyncPG Adapter Guide](/guides/adapters/asyncpg.md) +- [SQLSpec Configuration Documentation](/reference/configuration.rst) diff --git a/pyproject.toml b/pyproject.toml index 837ab49e..47f79766 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,8 @@ asyncmy = ["asyncmy"] asyncpg = ["asyncpg"] attrs = ["attrs", "cattrs"] bigquery = ["google-cloud-bigquery"] +cloud-sql = ["cloud-sql-python-connector"] +alloydb = ["google-cloud-alloydb-connector"] cli = ["rich-click"] duckdb = ["duckdb"] fastapi = ["fastapi"] diff --git a/sqlspec/_typing.py b/sqlspec/_typing.py index 5577b5c4..a54e4eb7 100644 --- a/sqlspec/_typing.py +++ b/sqlspec/_typing.py @@ -710,11 +710,23 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter OBSTORE_INSTALLED = bool(find_spec("obstore")) PGVECTOR_INSTALLED = bool(find_spec("pgvector")) +try: + CLOUD_SQL_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.sql.connector")) +except (ImportError, ModuleNotFoundError): + CLOUD_SQL_CONNECTOR_INSTALLED = False + +try: + ALLOYDB_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.alloydb.connector")) +except (ImportError, ModuleNotFoundError): + ALLOYDB_CONNECTOR_INSTALLED = False + __all__ = ( "AIOSQL_INSTALLED", + "ALLOYDB_CONNECTOR_INSTALLED", "ATTRS_INSTALLED", "CATTRS_INSTALLED", + "CLOUD_SQL_CONNECTOR_INSTALLED", "FSSPEC_INSTALLED", "LITESTAR_INSTALLED", "MSGSPEC_INSTALLED", diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 499bb941..012f8a21 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -19,7 +19,8 @@ build_asyncpg_statement_config, ) from sqlspec.config import ADKConfig, AsyncDatabaseConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig -from sqlspec.typing import PGVECTOR_INSTALLED +from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.typing import ALLOYDB_CONNECTOR_INSTALLED, CLOUD_SQL_CONNECTOR_INSTALLED, PGVECTOR_INSTALLED from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -88,12 +89,47 @@ class AsyncpgDriverFeatures(TypedDict): Defaults to True when pgvector-python is installed. Provides automatic conversion between Python objects and PostgreSQL vector types. Enables vector similarity operations and index support. + enable_cloud_sql: Enable Google Cloud SQL connector integration. + Requires cloud-sql-python-connector package. + Defaults to False (explicit opt-in required). + Auto-configures IAM authentication, SSL, and IP routing. + Mutually exclusive with enable_alloydb. + cloud_sql_instance: Cloud SQL instance connection name. + Format: "project:region:instance" + Required when enable_cloud_sql is True. + cloud_sql_enable_iam_auth: Enable IAM database authentication. + Defaults to False for passwordless authentication. + When False, requires user/password in pool_config. + cloud_sql_ip_type: IP address type for connection. + Options: "PUBLIC", "PRIVATE", "PSC" + Defaults to "PRIVATE". + enable_alloydb: Enable Google AlloyDB connector integration. + Requires cloud-alloydb-python-connector package. + Defaults to False (explicit opt-in required). + Auto-configures IAM authentication and private networking. + Mutually exclusive with enable_cloud_sql. + alloydb_instance_uri: AlloyDB instance URI. + Format: "projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE" + Required when enable_alloydb is True. + alloydb_enable_iam_auth: Enable IAM database authentication. + Defaults to False for passwordless authentication. + alloydb_ip_type: IP address type for connection. + Options: "PUBLIC", "PRIVATE", "PSC" + Defaults to "PRIVATE". """ json_serializer: NotRequired[Callable[[Any], str]] json_deserializer: NotRequired[Callable[[str], Any]] enable_json_codecs: NotRequired[bool] enable_pgvector: NotRequired[bool] + enable_cloud_sql: NotRequired[bool] + cloud_sql_instance: NotRequired[str] + cloud_sql_enable_iam_auth: NotRequired[bool] + cloud_sql_ip_type: NotRequired[str] + enable_alloydb: NotRequired[bool] + alloydb_instance_uri: NotRequired[str] + alloydb_enable_iam_auth: NotRequired[bool] + alloydb_ip_type: NotRequired[str] class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", AsyncpgDriver]): @@ -135,6 +171,8 @@ def __init__( deserializer = features_dict.setdefault("json_deserializer", from_json) features_dict.setdefault("enable_json_codecs", True) features_dict.setdefault("enable_pgvector", PGVECTOR_INSTALLED) + features_dict.setdefault("enable_cloud_sql", False) + features_dict.setdefault("enable_alloydb", False) base_statement_config = statement_config or build_asyncpg_statement_config( json_serializer=serializer, json_deserializer=deserializer @@ -150,6 +188,53 @@ def __init__( extension_config=extension_config, ) + self._cloud_sql_connector: Any | None = None + self._alloydb_connector: Any | None = None + + self._validate_connector_config() + + def _validate_connector_config(self) -> None: + """Validate Google Cloud connector configuration. + + Raises: + ImproperConfigurationError: If configuration is invalid. + """ + enable_cloud_sql = self.driver_features.get("enable_cloud_sql", False) + enable_alloydb = self.driver_features.get("enable_alloydb", False) + + if enable_cloud_sql and enable_alloydb: + msg = "Cannot enable both Cloud SQL and AlloyDB connectors simultaneously. Use separate configs for each database." + raise ImproperConfigurationError(msg) + + if enable_cloud_sql: + if not CLOUD_SQL_CONNECTOR_INSTALLED: + msg = "cloud-sql-python-connector package not installed. Install with: pip install cloud-sql-python-connector" + raise ImproperConfigurationError(msg) + + instance = self.driver_features.get("cloud_sql_instance") + if not instance: + msg = "cloud_sql_instance required when enable_cloud_sql is True. Format: 'project:region:instance'" + raise ImproperConfigurationError(msg) + + cloud_sql_instance_parts_expected = 2 + if instance.count(":") != cloud_sql_instance_parts_expected: + msg = f"Invalid Cloud SQL instance format: {instance}. Expected format: 'project:region:instance'" + raise ImproperConfigurationError(msg) + + elif enable_alloydb: + if not ALLOYDB_CONNECTOR_INSTALLED: + msg = "cloud-alloydb-python-connector package not installed. Install with: pip install cloud-alloydb-python-connector" + raise ImproperConfigurationError(msg) + + instance_uri = self.driver_features.get("alloydb_instance_uri") + if not instance_uri: + msg = "alloydb_instance_uri required when enable_alloydb is True. Format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" + raise ImproperConfigurationError(msg) + + if not instance_uri.startswith("projects/"): + msg = f"Invalid AlloyDB instance URI format: {instance_uri}. Expected format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'" + raise ImproperConfigurationError(msg) + def _get_pool_config_dict(self) -> "dict[str, Any]": """Get pool configuration as plain dict for external library. @@ -161,10 +246,89 @@ def _get_pool_config_dict(self) -> "dict[str, Any]": config.update(extras) return {k: v for k, v in config.items() if v is not None} + def _setup_cloud_sql_connector(self, config: "dict[str, Any]") -> None: + """Setup Cloud SQL connector and configure pool for connection factory pattern. + + Args: + config: Pool configuration dictionary to modify in-place. + """ + from google.cloud.sql.connector import Connector # pyright: ignore + + self._cloud_sql_connector = Connector() + + user = config.get("user") + password = config.get("password") + database = config.get("database") + + async def get_conn() -> "AsyncpgConnection": + conn_kwargs: dict[str, Any] = { + "instance_connection_string": self.driver_features["cloud_sql_instance"], + "driver": "asyncpg", + "enable_iam_auth": self.driver_features.get("cloud_sql_enable_iam_auth", False), + "ip_type": self.driver_features.get("cloud_sql_ip_type", "PRIVATE"), + } + + if user: + conn_kwargs["user"] = user + if password: + conn_kwargs["password"] = password + if database: + conn_kwargs["db"] = database + + conn: AsyncpgConnection = await self._cloud_sql_connector.connect_async(**conn_kwargs) # type: ignore[union-attr] + return conn + + for key in ("dsn", "host", "port", "user", "password", "database"): + config.pop(key, None) + + config["connect"] = get_conn + + def _setup_alloydb_connector(self, config: "dict[str, Any]") -> None: + """Setup AlloyDB connector and configure pool for connection factory pattern. + + Args: + config: Pool configuration dictionary to modify in-place. + """ + from google.cloud.alloydb.connector import AsyncConnector + + self._alloydb_connector = AsyncConnector() + + user = config.get("user") + password = config.get("password") + database = config.get("database") + + async def get_conn() -> "AsyncpgConnection": + conn_kwargs: dict[str, Any] = { + "instance_uri": self.driver_features["alloydb_instance_uri"], + "driver": "asyncpg", + "enable_iam_auth": self.driver_features.get("alloydb_enable_iam_auth", False), + "ip_type": self.driver_features.get("alloydb_ip_type", "PRIVATE"), + } + + if user: + conn_kwargs["user"] = user + if password: + conn_kwargs["password"] = password + if database: + conn_kwargs["db"] = database + + conn: AsyncpgConnection = await self._alloydb_connector.connect(**conn_kwargs) # type: ignore[union-attr] + return conn + + for key in ("dsn", "host", "port", "user", "password", "database"): + config.pop(key, None) + + config["connect"] = get_conn + async def _create_pool(self) -> "Pool[Record]": """Create the actual async connection pool.""" config = self._get_pool_config_dict() + if self.driver_features.get("enable_cloud_sql", False): + self._setup_cloud_sql_connector(config) + elif self.driver_features.get("enable_alloydb", False): + self._setup_alloydb_connector(config) + if "init" not in config: config["init"] = self._init_connection @@ -191,10 +355,18 @@ async def _init_connection(self, connection: "AsyncpgConnection") -> None: await register_pgvector_support(connection) async def _close_pool(self) -> None: - """Close the actual async connection pool.""" + """Close the actual async connection pool and cleanup connectors.""" if self.pool_instance: await self.pool_instance.close() + if self._cloud_sql_connector is not None: + await self._cloud_sql_connector.close_async() + self._cloud_sql_connector = None + + if self._alloydb_connector is not None: + await self._alloydb_connector.close() + self._alloydb_connector = None + async def close_pool(self) -> None: """Close the connection pool.""" await self._close_pool() diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 8447758e..8d215695 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -7,8 +7,10 @@ from sqlspec._typing import ( AIOSQL_INSTALLED, + ALLOYDB_CONNECTOR_INSTALLED, ATTRS_INSTALLED, CATTRS_INSTALLED, + CLOUD_SQL_CONNECTOR_INSTALLED, FSSPEC_INSTALLED, LITESTAR_INSTALLED, MSGSPEC_INSTALLED, @@ -144,8 +146,10 @@ def get_type_adapter(f: "type[T]") -> Any: __all__ = ( "AIOSQL_INSTALLED", + "ALLOYDB_CONNECTOR_INSTALLED", "ATTRS_INSTALLED", "CATTRS_INSTALLED", + "CLOUD_SQL_CONNECTOR_INSTALLED", "FSSPEC_INSTALLED", "LITESTAR_INSTALLED", "MSGSPEC_INSTALLED", diff --git a/tests/integration/test_adapters/test_asyncpg/test_cloud_connectors_integration.py b/tests/integration/test_adapters/test_asyncpg/test_cloud_connectors_integration.py new file mode 100644 index 00000000..b34ae474 --- /dev/null +++ b/tests/integration/test_adapters/test_asyncpg/test_cloud_connectors_integration.py @@ -0,0 +1,205 @@ +"""Integration tests for Google Cloud SQL and AlloyDB connector support. + +These tests require actual Google Cloud credentials and instances. +They are skipped by default unless credentials are available. +""" + +import os + +import pytest + +from sqlspec.adapters.asyncpg import AsyncpgConfig +from sqlspec.typing import ALLOYDB_CONNECTOR_INSTALLED, CLOUD_SQL_CONNECTOR_INSTALLED + +HAS_CLOUD_SQL_CREDENTIALS = ( + CLOUD_SQL_CONNECTOR_INSTALLED + and os.environ.get("GOOGLE_CLOUD_SQL_INSTANCE") + and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") +) + +HAS_ALLOYDB_CREDENTIALS = ( + ALLOYDB_CONNECTOR_INSTALLED + and os.environ.get("GOOGLE_ALLOYDB_INSTANCE_URI") + and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") +) + +pytestmark = [ + pytest.mark.skipif( + not (HAS_CLOUD_SQL_CREDENTIALS or HAS_ALLOYDB_CREDENTIALS), reason="Google Cloud credentials not available" + ), + pytest.mark.xdist_group("google_cloud"), +] + + +@pytest.mark.skipif(not HAS_CLOUD_SQL_CREDENTIALS, reason="Cloud SQL credentials not available") +@pytest.mark.asyncio +async def test_cloud_sql_connection_basic() -> None: + """Test basic Cloud SQL connection via connector.""" + instance = os.environ["GOOGLE_CLOUD_SQL_INSTANCE"] + user = os.environ.get("GOOGLE_CLOUD_SQL_USER", "postgres") + database = os.environ.get("GOOGLE_CLOUD_SQL_DATABASE", "postgres") + password = os.environ.get("GOOGLE_CLOUD_SQL_PASSWORD") + + config = AsyncpgConfig( + pool_config={"user": user, "password": password, "database": database, "min_size": 1, "max_size": 2}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": instance, "cloud_sql_enable_iam_auth": False}, + ) + + await config.create_pool() + try: + async with config.provide_connection() as conn: + result = await conn.fetchval("SELECT 1") + assert result == 1 + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_CLOUD_SQL_CREDENTIALS, reason="Cloud SQL credentials not available") +@pytest.mark.asyncio +async def test_cloud_sql_query_execution() -> None: + """Test query execution via Cloud SQL connector.""" + instance = os.environ["GOOGLE_CLOUD_SQL_INSTANCE"] + user = os.environ.get("GOOGLE_CLOUD_SQL_USER", "postgres") + database = os.environ.get("GOOGLE_CLOUD_SQL_DATABASE", "postgres") + password = os.environ.get("GOOGLE_CLOUD_SQL_PASSWORD") + + config = AsyncpgConfig( + pool_config={"user": user, "password": password, "database": database, "min_size": 1, "max_size": 2}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": instance, "cloud_sql_enable_iam_auth": False}, + ) + + await config.create_pool() + try: + async with config.provide_session() as session: + result = await session.select_one("SELECT 1 as value, 'test' as name") + assert result["value"] == 1 + assert result["name"] == "test" + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_CLOUD_SQL_CREDENTIALS, reason="Cloud SQL IAM test requires credentials") +@pytest.mark.asyncio +async def test_cloud_sql_iam_auth() -> None: + """Test Cloud SQL with IAM authentication.""" + instance = os.environ["GOOGLE_CLOUD_SQL_INSTANCE"] + user = os.environ.get("GOOGLE_CLOUD_SQL_IAM_USER", "service-account@project.iam") + database = os.environ.get("GOOGLE_CLOUD_SQL_DATABASE", "postgres") + + config = AsyncpgConfig( + pool_config={"user": user, "database": database, "min_size": 1, "max_size": 2}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": instance, "cloud_sql_enable_iam_auth": True}, + ) + + await config.create_pool() + try: + async with config.provide_connection() as conn: + result = await conn.fetchval("SELECT 1") + assert result == 1 + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_CLOUD_SQL_CREDENTIALS, reason="Cloud SQL credentials not available") +@pytest.mark.asyncio +async def test_cloud_sql_private_ip() -> None: + """Test Cloud SQL connection using PRIVATE IP type.""" + instance = os.environ["GOOGLE_CLOUD_SQL_INSTANCE"] + user = os.environ.get("GOOGLE_CLOUD_SQL_USER", "postgres") + database = os.environ.get("GOOGLE_CLOUD_SQL_DATABASE", "postgres") + password = os.environ.get("GOOGLE_CLOUD_SQL_PASSWORD") + + config = AsyncpgConfig( + pool_config={"user": user, "password": password, "database": database, "min_size": 1, "max_size": 2}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": instance, + "cloud_sql_enable_iam_auth": False, + "cloud_sql_ip_type": "PRIVATE", + }, + ) + + await config.create_pool() + try: + async with config.provide_connection() as conn: + result = await conn.fetchval("SELECT 1") + assert result == 1 + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_ALLOYDB_CREDENTIALS, reason="AlloyDB credentials not available") +@pytest.mark.asyncio +async def test_alloydb_connection_basic() -> None: + """Test basic AlloyDB connection via connector.""" + instance_uri = os.environ["GOOGLE_ALLOYDB_INSTANCE_URI"] + user = os.environ.get("GOOGLE_ALLOYDB_USER", "postgres") + database = os.environ.get("GOOGLE_ALLOYDB_DATABASE", "postgres") + password = os.environ.get("GOOGLE_ALLOYDB_PASSWORD") + + config = AsyncpgConfig( + pool_config={"user": user, "password": password, "database": database, "min_size": 1, "max_size": 2}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": instance_uri, + "alloydb_enable_iam_auth": False, + }, + ) + + await config.create_pool() + try: + async with config.provide_connection() as conn: + result = await conn.fetchval("SELECT 1") + assert result == 1 + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_ALLOYDB_CREDENTIALS, reason="AlloyDB credentials not available") +@pytest.mark.asyncio +async def test_alloydb_query_execution() -> None: + """Test query execution via AlloyDB connector.""" + instance_uri = os.environ["GOOGLE_ALLOYDB_INSTANCE_URI"] + user = os.environ.get("GOOGLE_ALLOYDB_USER", "postgres") + database = os.environ.get("GOOGLE_ALLOYDB_DATABASE", "postgres") + password = os.environ.get("GOOGLE_ALLOYDB_PASSWORD") + + config = AsyncpgConfig( + pool_config={"user": user, "password": password, "database": database, "min_size": 1, "max_size": 2}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": instance_uri, + "alloydb_enable_iam_auth": False, + }, + ) + + await config.create_pool() + try: + async with config.provide_session() as session: + result = await session.select_one("SELECT 1 as value, 'test' as name") + assert result["value"] == 1 + assert result["name"] == "test" + finally: + await config.close_pool() + + +@pytest.mark.skipif(not HAS_ALLOYDB_CREDENTIALS, reason="AlloyDB IAM test requires credentials") +@pytest.mark.asyncio +async def test_alloydb_iam_auth() -> None: + """Test AlloyDB with IAM authentication.""" + instance_uri = os.environ["GOOGLE_ALLOYDB_INSTANCE_URI"] + user = os.environ.get("GOOGLE_ALLOYDB_IAM_USER", "service-account@project.iam") + database = os.environ.get("GOOGLE_ALLOYDB_DATABASE", "postgres") + + config = AsyncpgConfig( + pool_config={"user": user, "database": database, "min_size": 1, "max_size": 2}, + driver_features={"enable_alloydb": True, "alloydb_instance_uri": instance_uri, "alloydb_enable_iam_auth": True}, + ) + + await config.create_pool() + try: + async with config.provide_connection() as conn: + result = await conn.fetchval("SELECT 1") + assert result == 1 + finally: + await config.close_pool() diff --git a/tests/unit/test_adapters/test_asyncpg/test_cloud_connectors.py b/tests/unit/test_adapters/test_asyncpg/test_cloud_connectors.py new file mode 100644 index 00000000..c6b04204 --- /dev/null +++ b/tests/unit/test_adapters/test_asyncpg/test_cloud_connectors.py @@ -0,0 +1,544 @@ +"""Unit tests for Google Cloud SQL and AlloyDB connector integration in AsyncPG.""" + +# pyright: reportPrivateUsage=false + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from sqlspec.adapters.asyncpg.config import AsyncpgConfig +from sqlspec.exceptions import ImproperConfigurationError + + +@pytest.fixture(autouse=True) +def disable_connectors_by_default(): + """Disable both connectors by default for clean test state.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", False): + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", False): + yield + + +@pytest.fixture +def mock_cloud_sql_module(): + """Create and register mock google.cloud.sql module.""" + mock_connector_class = MagicMock() + mock_module = MagicMock() + mock_module.connector.Connector = mock_connector_class + + sys.modules["google.cloud.sql"] = mock_module + sys.modules["google.cloud.sql.connector"] = mock_module.connector + + yield mock_connector_class + + sys.modules.pop("google.cloud.sql", None) + sys.modules.pop("google.cloud.sql.connector", None) + + +@pytest.fixture +def mock_alloydb_module(): + """Create and register mock google.cloud.alloydb module.""" + mock_connector_class = MagicMock() + mock_module = MagicMock() + mock_module.connector.AsyncConnector = mock_connector_class + + sys.modules["google.cloud.alloydb"] = mock_module + sys.modules["google.cloud.alloydb.connector"] = mock_module.connector + + yield mock_connector_class + + sys.modules.pop("google.cloud.alloydb", None) + sys.modules.pop("google.cloud.alloydb.connector", None) + + +def test_cloud_sql_defaults_to_false() -> None: + """Cloud SQL connector should always default to False (explicit opt-in required).""" + config = AsyncpgConfig(pool_config={"dsn": "postgresql://localhost/test"}) + assert config.driver_features["enable_cloud_sql"] is False + + +def test_alloydb_defaults_to_false() -> None: + """AlloyDB connector should always default to False (explicit opt-in required).""" + config = AsyncpgConfig(pool_config={"dsn": "postgresql://localhost/test"}) + assert config.driver_features["enable_alloydb"] is False + + +def test_mutual_exclusion_both_enabled_raises_error() -> None: + """Enabling both Cloud SQL and AlloyDB connectors should raise error.""" + with pytest.raises( + ImproperConfigurationError, match="Cannot enable both Cloud SQL and AlloyDB connectors simultaneously" + ): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "project:region:instance", + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + }, + ) + + +def test_cloud_sql_missing_package_raises_error() -> None: + """Enabling Cloud SQL without package installed should raise error.""" + with pytest.raises(ImproperConfigurationError, match="cloud-sql-python-connector package not installed"): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance"}, + ) + + +def test_alloydb_missing_package_raises_error() -> None: + """Enabling AlloyDB without package installed should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", False): + with pytest.raises(ImproperConfigurationError, match="cloud-alloydb-python-connector package not installed"): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + }, + ) + + +def test_cloud_sql_missing_instance_raises_error() -> None: + """Enabling Cloud SQL without instance string should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + with pytest.raises( + ImproperConfigurationError, match="cloud_sql_instance required when enable_cloud_sql is True" + ): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, driver_features={"enable_cloud_sql": True} + ) + + +def test_alloydb_missing_instance_uri_raises_error() -> None: + """Enabling AlloyDB without instance URI should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + with pytest.raises( + ImproperConfigurationError, match="alloydb_instance_uri required when enable_alloydb is True" + ): + AsyncpgConfig(pool_config={"dsn": "postgresql://localhost/test"}, driver_features={"enable_alloydb": True}) + + +def test_cloud_sql_invalid_instance_format_raises_error() -> None: + """Cloud SQL instance with invalid format should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + with pytest.raises(ImproperConfigurationError, match="Invalid Cloud SQL instance format"): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "invalid-format"}, + ) + + +def test_cloud_sql_instance_format_too_many_colons() -> None: + """Cloud SQL instance with too many colons should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + with pytest.raises(ImproperConfigurationError, match="Invalid Cloud SQL instance format"): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance:extra"}, + ) + + +def test_alloydb_invalid_instance_uri_format_raises_error() -> None: + """AlloyDB instance URI with invalid format should raise error.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + with pytest.raises(ImproperConfigurationError, match="Invalid AlloyDB instance URI format"): + AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, + driver_features={"enable_alloydb": True, "alloydb_instance_uri": "invalid-format"}, + ) + + +def test_cloud_sql_explicit_disable() -> None: + """Explicitly disabling Cloud SQL should work even when package installed.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + config = AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, driver_features={"enable_cloud_sql": False} + ) + assert config.driver_features["enable_cloud_sql"] is False + + +def test_alloydb_explicit_disable() -> None: + """Explicitly disabling AlloyDB should work even when package installed.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + config = AsyncpgConfig( + pool_config={"dsn": "postgresql://localhost/test"}, driver_features={"enable_alloydb": False} + ) + assert config.driver_features["enable_alloydb"] is False + + +def test_normal_config_without_connectors() -> None: + """Normal config without connectors should work.""" + config = AsyncpgConfig(pool_config={"dsn": "postgresql://localhost/test"}) + assert config is not None + assert config.driver_features.get("enable_cloud_sql", False) is not True + assert config.driver_features.get("enable_alloydb", False) is not True + + +@pytest.mark.asyncio +async def test_cloud_sql_connector_initialization(mock_cloud_sql_module) -> None: + """Cloud SQL connector should be initialized correctly in create_pool.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect_async = AsyncMock(return_value=MagicMock()) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance"}, + ) + + pool = await config._create_pool() + + assert pool is mock_pool + mock_cloud_sql_module.assert_called_once() + assert config._cloud_sql_connector is mock_connector # pyright: ignore + mock_create_pool.assert_called_once() + call_kwargs = mock_create_pool.call_args.kwargs + assert "connect" in call_kwargs + assert "dsn" not in call_kwargs + assert "host" not in call_kwargs + assert "user" not in call_kwargs + + +@pytest.mark.asyncio +async def test_cloud_sql_iam_auth_enabled(mock_cloud_sql_module) -> None: + """Cloud SQL IAM authentication should configure enable_iam_auth=True.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["enable_iam_auth"] is True + return MagicMock() + + mock_connector.connect_async = AsyncMock(side_effect=mock_connect) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "database": "testdb"}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "project:region:instance", + "cloud_sql_enable_iam_auth": True, + }, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_cloud_sql_iam_auth_disabled(mock_cloud_sql_module) -> None: + """Cloud SQL with IAM disabled should configure enable_iam_auth=False.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["enable_iam_auth"] is False + return MagicMock() + + mock_connector.connect_async = AsyncMock(side_effect=mock_connect) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "project:region:instance", + "cloud_sql_enable_iam_auth": False, + }, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_cloud_sql_ip_type_configuration(mock_cloud_sql_module) -> None: + """Cloud SQL IP type should be passed to connector.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["ip_type"] == "PUBLIC" + return MagicMock() + + mock_connector.connect_async = AsyncMock(side_effect=mock_connect) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_cloud_sql": True, + "cloud_sql_instance": "project:region:instance", + "cloud_sql_ip_type": "PUBLIC", + }, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_cloud_sql_default_ip_type(mock_cloud_sql_module) -> None: + """Cloud SQL should default to PRIVATE IP type.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["ip_type"] == "PRIVATE" + return MagicMock() + + mock_connector.connect_async = AsyncMock(side_effect=mock_connect) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance"}, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_alloydb_connector_initialization(mock_alloydb_module) -> None: + """AlloyDB connector should be initialized correctly in create_pool.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect = AsyncMock(return_value=MagicMock()) + mock_alloydb_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + }, + ) + + pool = await config._create_pool() + + assert pool is mock_pool + mock_alloydb_module.assert_called_once() + assert config._alloydb_connector is mock_connector + mock_create_pool.assert_called_once() + call_kwargs = mock_create_pool.call_args.kwargs + assert "connect" in call_kwargs + assert "dsn" not in call_kwargs + assert "host" not in call_kwargs + assert "user" not in call_kwargs + + +@pytest.mark.asyncio +async def test_alloydb_iam_auth_enabled(mock_alloydb_module) -> None: + """AlloyDB IAM authentication should configure enable_iam_auth=True.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["enable_iam_auth"] is True + return MagicMock() + + mock_connector.connect = AsyncMock(side_effect=mock_connect) + mock_alloydb_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "database": "testdb"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + "alloydb_enable_iam_auth": True, + }, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_alloydb_ip_type_configuration(mock_alloydb_module) -> None: + """AlloyDB IP type should be passed to connector.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + + async def mock_connect(**kwargs): + assert kwargs["ip_type"] == "PSC" + return MagicMock() + + mock_connector.connect = AsyncMock(side_effect=mock_connect) + mock_alloydb_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + "alloydb_ip_type": "PSC", + }, + ) + + await config._create_pool() + get_conn_func = mock_create_pool.call_args.kwargs["connect"] + await get_conn_func() + + +@pytest.mark.asyncio +async def test_cloud_sql_connector_cleanup(mock_cloud_sql_module) -> None: + """Cloud SQL connector should be closed on pool close.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect_async = AsyncMock(return_value=MagicMock()) + mock_connector.close_async = AsyncMock() + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_pool.close = AsyncMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance"}, + ) + + await config._create_pool() + await config._close_pool() + + mock_connector.close_async.assert_called_once() + assert config._cloud_sql_connector is None + + +@pytest.mark.asyncio +async def test_alloydb_connector_cleanup(mock_alloydb_module) -> None: + """AlloyDB connector should be closed on pool close.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect = AsyncMock(return_value=MagicMock()) + mock_connector.close = AsyncMock() + mock_alloydb_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_pool.close = AsyncMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + }, + ) + + await config._create_pool() + await config._close_pool() + + mock_connector.close.assert_called_once() + assert config._alloydb_connector is None + + +@pytest.mark.asyncio +async def test_connection_factory_pattern_cloud_sql(mock_cloud_sql_module) -> None: + """Cloud SQL should use connection factory pattern with connect parameter.""" + with patch("sqlspec.adapters.asyncpg.config.CLOUD_SQL_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect_async = AsyncMock(return_value=MagicMock()) + mock_cloud_sql_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={"enable_cloud_sql": True, "cloud_sql_instance": "project:region:instance"}, + ) + + await config._create_pool() + + call_kwargs = mock_create_pool.call_args.kwargs + assert "connect" in call_kwargs + assert callable(call_kwargs["connect"]) + + +@pytest.mark.asyncio +async def test_connection_factory_pattern_alloydb(mock_alloydb_module) -> None: + """AlloyDB should use connection factory pattern with connect parameter.""" + with patch("sqlspec.adapters.asyncpg.config.ALLOYDB_CONNECTOR_INSTALLED", True): + mock_connector = MagicMock() + mock_connector.connect = AsyncMock(return_value=MagicMock()) + mock_alloydb_module.return_value = mock_connector + + with patch("sqlspec.adapters.asyncpg.config.asyncpg_create_pool", new_callable=AsyncMock) as mock_create_pool: + mock_pool = MagicMock() + mock_create_pool.return_value = mock_pool + + config = AsyncpgConfig( + pool_config={"user": "testuser", "password": "testpass", "database": "testdb"}, + driver_features={ + "enable_alloydb": True, + "alloydb_instance_uri": "projects/p/locations/r/clusters/c/instances/i", + }, + ) + + await config._create_pool() + + call_kwargs = mock_create_pool.call_args.kwargs + assert "connect" in call_kwargs + assert callable(call_kwargs["connect"]) + + +@pytest.mark.asyncio +async def test_pool_close_without_connectors() -> None: + """Closing pool without connectors should not raise errors.""" + config = AsyncpgConfig(pool_config={"dsn": "postgresql://localhost/test"}) + + mock_pool = MagicMock() + mock_pool.close = AsyncMock() + config.pool_instance = mock_pool + + await config._close_pool() + + mock_pool.close.assert_called_once() + assert config._cloud_sql_connector is None + assert config._alloydb_connector is None diff --git a/uv.lock b/uv.lock index 8c59ec3b..7833100c 100644 --- a/uv.lock +++ b/uv.lock @@ -143,6 +143,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/2a/d275ec4ce5cd0096665043995a7d76f5d0524853c76a3d04656de49f8808/aiobotocore-2.25.1-py3-none-any.whl", hash = "sha256:eb6daebe3cbef5b39a0bb2a97cffbe9c7cb46b2fcc399ad141f369f3c2134b1f", size = 86039, upload-time = "2025-10-28T22:33:19.949Z" }, ] +[[package]] +name = "aiofiles" +version = "25.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/c3/534eac40372d8ee36ef40df62ec129bee4fdb5ad9706e58a29be53b2c970/aiofiles-25.1.0.tar.gz", hash = "sha256:a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2", size = 46354, upload-time = "2025-10-09T20:51:04.358Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -969,6 +978,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/85/32/10bb5764d90a8eee674e9dc6f4db6a0ab47c8c4d0d83c27f7c39ac415a4d/click-8.2.1-py3-none-any.whl", hash = "sha256:61a3265b914e850b85317d0b3109c7f8cd35a670f963866005d6ef1d5175a12b", size = 102215, upload-time = "2025-05-20T23:19:47.796Z" }, ] +[[package]] +name = "cloud-sql-python-connector" +version = "1.18.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiofiles" }, + { name = "aiohttp" }, + { name = "cryptography" }, + { name = "dnspython" }, + { name = "google-auth" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cb/9f/76285c820bce09729f716cf9e466d0261bc30d6a008dae5176c469f883d3/cloud_sql_python_connector-1.18.5.tar.gz", hash = "sha256:043477258226214973ebb9d1f6c8e775552619031a3a1e6c9337f0e4b05758ff", size = 42913, upload-time = "2025-10-09T22:30:04.133Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/36/9eab10d407369f9dc5c3ba8304c7d3926af7ba59ec1cde1ac74ba7488b62/cloud_sql_python_connector-1.18.5-py3-none-any.whl", hash = "sha256:ddead8e081d09e718967cfd5ba7100749f798742e962fbe03d887fa254a424b6", size = 49396, upload-time = "2025-10-09T22:30:03.074Z" }, +] + [[package]] name = "cloudpickle" version = "3.1.2" @@ -1214,6 +1240,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" }, ] +[[package]] +name = "dnspython" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, +] + [[package]] name = "docker" version = "7.1.0" @@ -1740,6 +1775,41 @@ agent-engines = [ { name = "typing-extensions" }, ] +[[package]] +name = "google-cloud-alloydb" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "google-api-core", extra = ["grpc"] }, + { name = "google-auth" }, + { name = "grpc-google-iam-v1" }, + { name = "grpcio" }, + { name = "proto-plus" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/a2/5f08b312c275116def8d6ef717d62fe72a8e3037f687e8d8f068e11e4987/google_cloud_alloydb-0.6.0.tar.gz", hash = "sha256:8062a9237a96992cd379b9497cba36c39e3ac0a4c504f385a861d9e30c99ab51", size = 682612, upload-time = "2025-10-17T02:33:33.708Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/a0/90766b6cefc1f96ac6299c945d599d459cbea624d84850b51701ab2ed9cb/google_cloud_alloydb-0.6.0-py3-none-any.whl", hash = "sha256:f45e64e8a77bc00beb27f8559cff77530e76f32a1e336e9aae4336f09bd5e836", size = 525800, upload-time = "2025-10-17T02:30:37.192Z" }, +] + +[[package]] +name = "google-cloud-alloydb-connector" +version = "1.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiofiles" }, + { name = "cryptography" }, + { name = "google-api-core" }, + { name = "google-auth" }, + { name = "google-cloud-alloydb" }, + { name = "protobuf" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6f/03/d94089a6bc8f22abfb6d9396fe2636f47c91dc8a3a718f223444c3740c9c/google_cloud_alloydb_connector-1.9.1.tar.gz", hash = "sha256:1f50794f428d6f5da09c874fe209120e2a7247d618c3e9a4584eb84307f3c138", size = 36069, upload-time = "2025-09-08T17:26:13.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/73/c43854d8f8103b175fb2689a6cf755771823557d524831f88e372d91924b/google_cloud_alloydb_connector-1.9.1-py3-none-any.whl", hash = "sha256:d4da1722321279e4ecff29e37d9d3e42367e8ddb71e9f2d756eb5aa8ccbafcf1", size = 45840, upload-time = "2025-09-08T17:26:12.667Z" }, +] + [[package]] name = "google-cloud-appengine-logging" version = "1.7.0" @@ -4137,14 +4207,14 @@ wheels = [ [[package]] name = "pyarrow-stubs" -version = "20.0.0.20250928" +version = "20.0.0.20251104" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyarrow" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/18/5f/9520b0a5cd42b95a945b8ca3bc47f723fc7ec906b7a7de76f2d075d69911/pyarrow_stubs-20.0.0.20250928.tar.gz", hash = "sha256:e802b18e8e5fdf0a78afa05fae78f1456d861fcb1f95ec0234be5d6a5ecdcde2", size = 236588, upload-time = "2025-09-28T02:50:04.839Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f2/52/5ce506962bc64d92b572499489f16ca05e8f466bbb9ed91b032fc653e624/pyarrow_stubs-20.0.0.20251104.tar.gz", hash = "sha256:a08964a67627c28354668af64d7021fb8d51fd69322ddb8b7a9c0336e49a367b", size = 236544, upload-time = "2025-11-04T02:16:09.485Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/83/13/75c86a8ef61ea2c758c924318cf894dced2436b0f7aeb3c5f0fe9e4305b4/pyarrow_stubs-20.0.0.20250928-py3-none-any.whl", hash = "sha256:5389057a55db3c2662c05f22685a52e15e5effaf4345f41f12fb9b6b348647b9", size = 235745, upload-time = "2025-09-28T02:50:03.205Z" }, + { url = "https://files.pythonhosted.org/packages/ac/b8/d130c794e61a719eb5de0e64cff0e01928986aade997486730c42c79843d/pyarrow_stubs-20.0.0.20251104-py3-none-any.whl", hash = "sha256:aea08b72754e342b982cbf764ab5dbfbdf041d8cca95d3c3b89d829d48f509c8", size = 235721, upload-time = "2025-11-04T02:16:07.749Z" }, ] [[package]] @@ -5826,6 +5896,9 @@ aiosql = [ aiosqlite = [ { name = "aiosqlite" }, ] +alloydb = [ + { name = "google-cloud-alloydb-connector" }, +] asyncmy = [ { name = "asyncmy" }, ] @@ -5842,6 +5915,9 @@ bigquery = [ cli = [ { name = "rich-click" }, ] +cloud-sql = [ + { name = "cloud-sql-python-connector" }, +] duckdb = [ { name = "duckdb" }, ] @@ -6067,12 +6143,14 @@ requires-dist = [ { name = "asyncpg", marker = "extra == 'asyncpg'" }, { name = "attrs", marker = "extra == 'attrs'" }, { name = "cattrs", marker = "extra == 'attrs'" }, + { name = "cloud-sql-python-connector", marker = "extra == 'cloud-sql'" }, { name = "duckdb", marker = "extra == 'duckdb'" }, { name = "fastapi", marker = "extra == 'fastapi'" }, { name = "fastnanoid", marker = "extra == 'nanoid'", specifier = ">=0.4.1" }, { name = "flask", marker = "extra == 'flask'" }, { name = "fsspec", marker = "extra == 'fsspec'" }, { name = "google-adk", marker = "extra == 'adk'" }, + { name = "google-cloud-alloydb-connector", marker = "extra == 'alloydb'" }, { name = "google-cloud-bigquery", marker = "extra == 'bigquery'" }, { name = "google-cloud-spanner", marker = "extra == 'spanner'" }, { name = "litestar", marker = "extra == 'litestar'" }, @@ -6102,7 +6180,7 @@ requires-dist = [ { name = "typing-extensions" }, { name = "uuid-utils", marker = "extra == 'uuid'" }, ] -provides-extras = ["adbc", "adk", "aioodbc", "aiosql", "aiosqlite", "asyncmy", "asyncpg", "attrs", "bigquery", "cli", "duckdb", "fastapi", "flask", "fsspec", "litestar", "msgspec", "mypyc", "nanoid", "obstore", "opentelemetry", "oracledb", "orjson", "pandas", "performance", "polars", "prometheus", "psqlpy", "psycopg", "pydantic", "pymssql", "pymysql", "spanner", "uuid"] +provides-extras = ["adbc", "adk", "aioodbc", "aiosql", "aiosqlite", "asyncmy", "asyncpg", "attrs", "bigquery", "cloud-sql", "alloydb", "cli", "duckdb", "fastapi", "flask", "fsspec", "litestar", "msgspec", "mypyc", "nanoid", "obstore", "opentelemetry", "oracledb", "orjson", "pandas", "performance", "polars", "prometheus", "psqlpy", "psycopg", "pydantic", "pymssql", "pymysql", "spanner", "uuid"] [package.metadata.requires-dev] benchmarks = [ From 310d760682c62922aca5f6badd6e587037e4b150 Mon Sep 17 00:00:00 2001 From: Cody Fincher <204685+cofin@users.noreply.github.com> Date: Mon, 3 Nov 2025 22:25:37 -0600 Subject: [PATCH 2/9] Apply suggestion from @cofin --- sqlspec/_typing.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/sqlspec/_typing.py b/sqlspec/_typing.py index a54e4eb7..12ffd4c0 100644 --- a/sqlspec/_typing.py +++ b/sqlspec/_typing.py @@ -710,16 +710,8 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter OBSTORE_INSTALLED = bool(find_spec("obstore")) PGVECTOR_INSTALLED = bool(find_spec("pgvector")) -try: - CLOUD_SQL_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.sql.connector")) -except (ImportError, ModuleNotFoundError): - CLOUD_SQL_CONNECTOR_INSTALLED = False - -try: - ALLOYDB_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.alloydb.connector")) -except (ImportError, ModuleNotFoundError): - ALLOYDB_CONNECTOR_INSTALLED = False - +CLOUD_SQL_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.sql.connector")) +ALLOYDB_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.alloydb.connector")) __all__ = ( "AIOSQL_INSTALLED", From d49b5aeca4bccf41740e929158b00d05bca1a143 Mon Sep 17 00:00:00 2001 From: Cody Fincher <204685+cofin@users.noreply.github.com> Date: Mon, 3 Nov 2025 22:27:53 -0600 Subject: [PATCH 3/9] Apply suggestion from @cofin --- docs/guides/adapters/asyncpg.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/docs/guides/adapters/asyncpg.md b/docs/guides/adapters/asyncpg.md index 747ae08c..791bce4c 100644 --- a/docs/guides/adapters/asyncpg.md +++ b/docs/guides/adapters/asyncpg.md @@ -110,9 +110,6 @@ config = AsyncpgConfig( ) ``` -### Configuration Notes - -**Auto-Detection**: Both connectors are automatically enabled when their respective packages are installed: ```bash # Install Cloud SQL connector From e98b4cc70019cb70136c2b70fb824a1ade3a2349 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Tue, 4 Nov 2025 21:23:33 +0000 Subject: [PATCH 4/9] fix: suppress FutureWarning for Google API Core in filterwarnings --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 47f79766..796abb8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -295,6 +295,7 @@ filterwarnings = [ "ignore:`use_rich_markup=` will be deprecated:PendingDeprecationWarning", "ignore:`show_metavars_column=` will be deprecated:PendingDeprecationWarning", "ignore:`append_metavars_help=` will be deprecated:PendingDeprecationWarning", + "ignore:You are using a Python version.*:FutureWarning:google.api_core.*", ] markers = [ "integration: marks tests that require an external database", From 8b5fb6395c1d74cfa7518bca5628100f4d0c40de Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 5 Nov 2025 00:27:52 +0000 Subject: [PATCH 5/9] feat: add cloud-sql and alloydb connectors to test requirements in workflows --- .github/workflows/publish.yml | 1 + .github/workflows/test-build.yml | 1 + pyproject.toml | 5 ++--- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 1580ab5d..b9b65984 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -105,6 +105,7 @@ jobs: CIBW_ENVIRONMENT: "HATCH_BUILD_HOOKS_ENABLE=1 MYPYC_OPT_LEVEL=3 MYPYC_DEBUG_LEVEL=0 MYPYC_MULTI_FILE=1" # Test the built wheels + CIBW_TEST_REQUIRES: "cloud-sql-python-connector cloud-alloydb-python-connector" CIBW_TEST_COMMAND: "python -c \"import sqlspec; print('MyPyC wheel test passed')\"" - name: Upload wheel artifacts diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index 8fac65fd..e199b3db 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -117,6 +117,7 @@ jobs: CIBW_ENVIRONMENT: "HATCH_BUILD_HOOKS_ENABLE=1 MYPYC_OPT_LEVEL=3 MYPYC_DEBUG_LEVEL=0 MYPYC_MULTI_FILE=1" # Test the built wheels + CIBW_TEST_REQUIRES: "cloud-sql-python-connector cloud-alloydb-python-connector" CIBW_TEST_COMMAND: "python -c \"import sqlspec; print('MyPyC wheel test passed')\"" - name: Upload wheel artifacts diff --git a/pyproject.toml b/pyproject.toml index 796abb8d..9d9280b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,13 +20,13 @@ adk = ["google-adk"] aioodbc = ["aioodbc"] aiosql = ["aiosql"] aiosqlite = ["aiosqlite"] +alloydb = ["google-cloud-alloydb-connector"] asyncmy = ["asyncmy"] asyncpg = ["asyncpg"] attrs = ["attrs", "cattrs"] bigquery = ["google-cloud-bigquery"] -cloud-sql = ["cloud-sql-python-connector"] -alloydb = ["google-cloud-alloydb-connector"] cli = ["rich-click"] +cloud-sql = ["cloud-sql-python-connector"] duckdb = ["duckdb"] fastapi = ["fastapi"] flask = ["flask"] @@ -295,7 +295,6 @@ filterwarnings = [ "ignore:`use_rich_markup=` will be deprecated:PendingDeprecationWarning", "ignore:`show_metavars_column=` will be deprecated:PendingDeprecationWarning", "ignore:`append_metavars_help=` will be deprecated:PendingDeprecationWarning", - "ignore:You are using a Python version.*:FutureWarning:google.api_core.*", ] markers = [ "integration: marks tests that require an external database", From b217a65ece3a9c399e2ee944473950a2e2410902 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 5 Nov 2025 00:36:56 +0000 Subject: [PATCH 6/9] fix: update Python version in CI workflows and correct test requirements for AlloyDB connector --- .github/workflows/ci.yml | 8 ++++---- .github/workflows/publish.yml | 2 +- tests/unit/test_cli/test_shell_completion.py | 1 - 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 15062abf..0fde103b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: uses: astral-sh/setup-uv@v3 - name: Set up Python - run: uv python install 3.10 + run: uv python install 3.11 - name: Create virtual environment run: uv sync --all-extras --dev @@ -43,7 +43,7 @@ jobs: uses: astral-sh/setup-uv@v3 - name: Set up Python - run: uv python install 3.10 + run: uv python install 3.11 - name: Install dependencies run: uv sync --all-extras --dev @@ -60,7 +60,7 @@ jobs: uses: astral-sh/setup-uv@v3 - name: Set up Python - run: uv python install 3.10 + run: uv python install 3.11 - name: Install dependencies run: uv sync --all-extras --dev @@ -77,7 +77,7 @@ jobs: uses: astral-sh/setup-uv@v3 - name: Set up Python - run: uv python install 3.10 + run: uv python install 3.11 - name: Install dependencies run: uv sync --all-extras --dev diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index b9b65984..e4c1d11f 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -105,7 +105,7 @@ jobs: CIBW_ENVIRONMENT: "HATCH_BUILD_HOOKS_ENABLE=1 MYPYC_OPT_LEVEL=3 MYPYC_DEBUG_LEVEL=0 MYPYC_MULTI_FILE=1" # Test the built wheels - CIBW_TEST_REQUIRES: "cloud-sql-python-connector cloud-alloydb-python-connector" + CIBW_TEST_REQUIRES: "cloud-sql-python-connector google-cloud-alloydb-connector" CIBW_TEST_COMMAND: "python -c \"import sqlspec; print('MyPyC wheel test passed')\"" - name: Upload wheel artifacts diff --git a/tests/unit/test_cli/test_shell_completion.py b/tests/unit/test_cli/test_shell_completion.py index 6a936313..e6f9d03f 100644 --- a/tests/unit/test_cli/test_shell_completion.py +++ b/tests/unit/test_cli/test_shell_completion.py @@ -69,4 +69,3 @@ def test_completion_scripts_are_valid_shell_syntax() -> None: assert result.returncode == 0, f"{shell_name} completion failed: {result.stderr}" assert len(result.stdout) > 0, f"{shell_name} completion script is empty" - assert result.stderr == "", f"{shell_name} completion has stderr output: {result.stderr}" From d550044c4ec6686b826a2d43f5431420b83ea56b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 5 Nov 2025 00:42:09 +0000 Subject: [PATCH 7/9] fix: update test requirements for MyPyC wheels to use google-cloud-alloydb-connector --- .github/workflows/test-build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index e199b3db..6c6f8f5b 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -117,7 +117,7 @@ jobs: CIBW_ENVIRONMENT: "HATCH_BUILD_HOOKS_ENABLE=1 MYPYC_OPT_LEVEL=3 MYPYC_DEBUG_LEVEL=0 MYPYC_MULTI_FILE=1" # Test the built wheels - CIBW_TEST_REQUIRES: "cloud-sql-python-connector cloud-alloydb-python-connector" + CIBW_TEST_REQUIRES: "cloud-sql-python-connector google-cloud-alloydb-connector" CIBW_TEST_COMMAND: "python -c \"import sqlspec; print('MyPyC wheel test passed')\"" - name: Upload wheel artifacts From a6d736dca969a7d6c64ac756bd9fe7a5386aaf3c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 5 Nov 2025 03:57:14 +0000 Subject: [PATCH 8/9] feat: add module_available function to check module availability --- sqlspec/_typing.py | 29 +++++++++++++++++++++++------ sqlspec/typing.py | 2 ++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/sqlspec/_typing.py b/sqlspec/_typing.py index 12ffd4c0..23c62882 100644 --- a/sqlspec/_typing.py +++ b/sqlspec/_typing.py @@ -11,6 +11,22 @@ from typing_extensions import TypeVar, dataclass_transform +def module_available(module_name: str) -> bool: + """Return True if the given module spec can be resolved. + + Args: + module_name: Dotted path for the module to locate. + + Returns: + True if the module can be resolved, False otherwise. + """ + + try: + return find_spec(module_name) is not None + except ModuleNotFoundError: + return False + + @runtime_checkable class DataclassProtocol(Protocol): """Protocol for instance checking dataclasses.""" @@ -705,13 +721,13 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter AIOSQL_INSTALLED = False # pyright: ignore[reportConstantRedefinition] # pyright: ignore[reportConstantRedefinition] -FSSPEC_INSTALLED = bool(find_spec("fsspec")) -NUMPY_INSTALLED = bool(find_spec("numpy")) -OBSTORE_INSTALLED = bool(find_spec("obstore")) -PGVECTOR_INSTALLED = bool(find_spec("pgvector")) +FSSPEC_INSTALLED = module_available("fsspec") +NUMPY_INSTALLED = module_available("numpy") +OBSTORE_INSTALLED = module_available("obstore") +PGVECTOR_INSTALLED = module_available("pgvector") -CLOUD_SQL_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.sql.connector")) -ALLOYDB_CONNECTOR_INSTALLED = bool(find_spec("google.cloud.alloydb.connector")) +CLOUD_SQL_CONNECTOR_INSTALLED = module_available("google.cloud.sql.connector") +ALLOYDB_CONNECTOR_INSTALLED = module_available("google.cloud.alloydb.connector") __all__ = ( "AIOSQL_INSTALLED", @@ -794,5 +810,6 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter "cattrs_unstructure", "convert", "convert_stub", + "module_available", "trace", ) diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 8d215695..1cb6deda 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -68,6 +68,7 @@ cattrs_structure, cattrs_unstructure, convert, + module_available, trace, ) @@ -213,5 +214,6 @@ def get_type_adapter(f: "type[T]") -> Any: "cattrs_unstructure", "convert", "get_type_adapter", + "module_available", "trace", ) From 93481a3aa267c21eacd15d049f2eba86b9d2a44d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 5 Nov 2025 03:59:23 +0000 Subject: [PATCH 9/9] fix: update import statements for Google Cloud connectors to suppress type checking warnings --- sqlspec/adapters/asyncpg/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 012f8a21..16d4dca5 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -252,7 +252,7 @@ def _setup_cloud_sql_connector(self, config: "dict[str, Any]") -> None: Args: config: Pool configuration dictionary to modify in-place. """ - from google.cloud.sql.connector import Connector # pyright: ignore + from google.cloud.sql.connector import Connector # type: ignore[import-untyped,unused-ignore] self._cloud_sql_connector = Connector() @@ -289,7 +289,7 @@ def _setup_alloydb_connector(self, config: "dict[str, Any]") -> None: Args: config: Pool configuration dictionary to modify in-place. """ - from google.cloud.alloydb.connector import AsyncConnector + from google.cloud.alloydb.connector import AsyncConnector # type: ignore[import-untyped,unused-ignore] self._alloydb_connector = AsyncConnector()