datacommons-db 1.1.3__py3-none-any.whl

datacommons_db/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.metadata
+
+__version__ = importlib.metadata.version("datacommons-db")

datacommons_db/clients/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datacommons_db.clients.spanner_client import (
+    DdlResult,
+    DmlResult,
+    ExecutionStatus,
+    QueryResult,
+    SpannerClient,
+)
+
+__all__ = [
+    "DdlResult",
+    "DmlResult",
+    "ExecutionStatus",
+    "QueryResult",
+    "SpannerClient",
+]

datacommons_db/clients/spanner_client.py

@@ -0,0 +1,251 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any
+
+from google.auth.credentials import Credentials
+from google.cloud import spanner
+from google.cloud.spanner_v1.transaction import Transaction
+
+from datacommons_db.utils.validators import (
+    validate_resource_id,
+    validate_table_name,
+)
+
+
+class ExecutionStatus(StrEnum):
+    """Status of a Spanner database operation."""
+
+    SUCCESS = "SUCCESS"
+    ERROR = "ERROR"
+
+
+@dataclass(frozen=True)
+class DdlResult:
+    """Result of a DDL statement execution.
+
+    Attributes:
+        status: Execution status enum (SUCCESS or ERROR).
+        error_message: Error message string if execution failed, None otherwise.
+    """
+
+    status: ExecutionStatus
+    error_message: str | None = None
+
+
+@dataclass(frozen=True)
+class DmlResult:
+    """Result of a DML statement execution inside a read-write transaction.
+
+    Attributes:
+        status: Execution status enum (SUCCESS or ERROR).
+        rows_affected: Number of rows modified by the DML statement (0 on failure).
+        error_message: Error message string if execution failed, None otherwise.
+    """
+
+    status: ExecutionStatus
+    rows_affected: int = 0
+    error_message: str | None = None
+
+
+@dataclass(frozen=True)
+class QueryResult:
+    """Result of a snapshot read query.
+
+    Attributes:
+        status: Execution status enum (SUCCESS or ERROR).
+        rows: List of rows where each row is a list of column values ([] on failure).
+        error_message: Error message string if execution failed, None otherwise.
+    """
+
+    status: ExecutionStatus
+    rows: list[list[Any]] = field(default_factory=list)
+    error_message: str | None = None
+
+
+class SpannerClient:
+    """Client for Cloud Spanner operations, DDL execution, and query execution."""
+
+    def __init__(
+        self,
+        project_id: str,
+        instance_id: str,
+        database_id: str,
+        credentials: Credentials | None = None,
+        *,
+        disable_builtin_metrics: bool = True,
+    ) -> None:
+        """Initialize the SpannerClient.
+
+        Args:
+            project_id: GCP project ID.
+            instance_id: Cloud Spanner instance ID.
+            database_id: Cloud Spanner database ID.
+            credentials: Optional Google Cloud credentials object.
+            disable_builtin_metrics: Whether to disable built-in Cloud Monitoring metrics export.
+        """
+        validate_resource_id("project_id", project_id)
+        validate_resource_id("instance_id", instance_id)
+        validate_resource_id("database_id", database_id)
+
+        self.project_id = project_id
+        self.instance_id = instance_id
+        self.database_id = database_id
+
+        self.client = spanner.Client(
+            project=project_id,
+            credentials=credentials,
+            disable_builtin_metrics=disable_builtin_metrics,
+        )
+        self.instance = self.client.instance(self.instance_id)
+        self.database = self.instance.database(self.database_id)
+
+    def table_exists(self, table_name: str) -> bool:
+        """Check if a table exists in the Cloud Spanner database.
+
+        Args:
+            table_name: The table name to check.
+
+        Returns:
+            True if the table exists, False otherwise.
+        """
+        validate_table_name(table_name)
+
+        query = (
+            "SELECT 1 FROM information_schema.tables "
+            "WHERE table_schema IN ('', 'public') AND table_name = @table_name LIMIT 1"
+        )
+
+        params = {"table_name": table_name}
+        param_types = {"table_name": spanner.param_types.STRING}
+
+        result = self.execute_query(query, params=params, param_types=param_types)
+        return bool(result.status == ExecutionStatus.SUCCESS and result.rows)
+
+    def execute_ddl(self, ddl_statements: list[str]) -> DdlResult:
+        """Execute DDL statements and wait for completion.
+
+        Handles operations like CREATE TABLE, ALTER TABLE, etc.
+
+        Args:
+            ddl_statements: A non-empty list of DDL statement strings.
+
+        Returns:
+            DdlResult with execution status and optional error message.
+        """
+        if not ddl_statements or not isinstance(ddl_statements, list):
+            return DdlResult(
+                status=ExecutionStatus.ERROR,
+                error_message="ddl_statements must be a non-empty list of str.",
+            )
+
+        try:
+            operation = self.database.update_ddl(ddl_statements)
+            operation.result()
+            return DdlResult(status=ExecutionStatus.SUCCESS)
+        except Exception as e:  # noqa: BLE001 - must catch all exceptions to ensure a DdlResult is always returned
+            return DdlResult(status=ExecutionStatus.ERROR, error_message=str(e))
+
+    def execute_dml(
+        self,
+        query: str,
+        params: dict[str, Any] | None = None,
+        param_types: dict[str, Any] | None = None,
+    ) -> DmlResult:
+        """Execute a DML statement inside a Spanner read-write transaction.
+
+        Handles operations like INSERT, UPDATE, DELETE
+
+        Args:
+            query: The parameterized DML statement.
+            params: Dictionary of parameters.
+            param_types: Dictionary of parameter types.
+
+        Returns:
+            DmlResult with execution status, rows affected, and optional error message.
+        """
+
+        def _unit_of_work(transaction: Transaction) -> int:
+            return transaction.execute_update(
+                query, params=params, param_types=param_types
+            )
+
+        try:
+            rows_affected = self.database.run_in_transaction(_unit_of_work)
+            return DmlResult(
+                status=ExecutionStatus.SUCCESS, rows_affected=rows_affected
+            )
+        except Exception as e:  # noqa: BLE001 - must catch all exceptions to ensure a DmlResult is always returned
+            return DmlResult(
+                status=ExecutionStatus.ERROR,
+                rows_affected=0,
+                error_message=str(e),
+            )
+
+    def _execute_query_stream(
+        self,
+        query: str,
+        params: dict[str, Any] | None = None,
+        param_types: dict[str, Any] | None = None,
+    ) -> Iterator[list[Any]]:
+        """Stream query results lazily row-by-row from a point-in-time snapshot.
+        This prevents OOM errors when querying large tables.
+
+        Args:
+            query: The parameterized SQL query string.
+            params: Dictionary of query parameters.
+            param_types: Dictionary of parameter types.
+
+        Yields:
+            Rows as lists of column values.
+        """
+        with self.database.snapshot() as snapshot:
+            results = snapshot.execute_sql(
+                query, params=params, param_types=param_types
+            )
+            for row in results:
+                yield list(row)
+
+    def execute_query(
+        self,
+        query: str,
+        params: dict[str, Any] | None = None,
+        param_types: dict[str, Any] | None = None,
+    ) -> QueryResult:
+        """Execute a read query within a snapshot transaction.
+
+        Handles operations like SELECT.
+
+        Args:
+            query: The parameterized SQL query.
+            params: Dictionary of parameters.
+            param_types: Dictionary of parameter types.
+
+        Returns:
+            QueryResult with execution status, rows, and optional error message.
+        """
+        try:
+            rows = list(
+                self._execute_query_stream(
+                    query, params=params, param_types=param_types
+                )
+            )
+            return QueryResult(status=ExecutionStatus.SUCCESS, rows=rows)
+        except Exception as e:  # noqa: BLE001 - must catch all exceptions to ensure a QueryResult is always returned
+            return QueryResult(
+                status=ExecutionStatus.ERROR, rows=[], error_message=str(e)
+            )

datacommons_db/migrations/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datacommons_db.migrations.base import SchemaMigration
+from datacommons_db.migrations.migration_runner import (
+    MigrationResult,
+    MigrationRunner,
+)
+
+__all__ = [
+    "MigrationResult",
+    "MigrationRunner",
+    "SchemaMigration",
+]

datacommons_db/migrations/base.py

@@ -0,0 +1,42 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+
+from datacommons_db.clients.spanner_client import SpannerClient
+
+
+class SchemaMigration(ABC):
+    """Base class for defining Spanner schema migration scripts."""
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Human-readable description of this schema migration."""
+
+    @property
+    @abstractmethod
+    def creation_timestamp(self) -> str:
+        """UTC ISO-8601 creation timestamp (e.g. '2026-08-17T15:30:00Z')."""
+
+    @abstractmethod
+    def upgrade(self, spanner_client: SpannerClient) -> None:
+        """Executes forward schema changes to upgrade the database.
+
+        Args:
+            spanner_client: SpannerClient instance to execute DDL / DML.
+
+        Raises:
+            RuntimeError: If any DDL or DML operation fails.
+        """

datacommons_db/migrations/migration_runner.py

@@ -0,0 +1,297 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import inspect
+import logging
+import pkgutil
+from dataclasses import dataclass
+
+from google.cloud import spanner
+
+import datacommons_db.migrations.migration_scripts
+from datacommons_db.clients.spanner_client import ExecutionStatus, SpannerClient
+from datacommons_db.migrations.base import SchemaMigration
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MigrationResult:
+    """Result of applying a schema migration.
+
+    Attributes:
+        status: ExecutionStatus enum (SUCCESS or ERROR).
+        creation_timestamp: UTC ISO-8601 creation timestamp of the migration.
+        description: Description of the schema migration.
+        error_message: Error message string if migration failed, None otherwise.
+    """
+
+    status: ExecutionStatus
+    creation_timestamp: str
+    description: str
+    error_message: str | None = None
+
+
+class MigrationRunner:
+    """Discovers, validates, plans, and executes Spanner database schema migrations."""
+
+    def __init__(
+        self,
+        spanner_client: SpannerClient,
+        migrations: list[SchemaMigration] | None = None,
+    ) -> None:
+        """Initialize the MigrationRunner.
+
+        Args:
+            spanner_client: SpannerClient instance.
+            migrations: Optional explicit list of SchemaMigration instances.
+                If omitted, migrations will be automatically discovered.
+        """
+        self.spanner_client = spanner_client
+        if migrations is not None:
+            self.migrations = self.validate_migrations(migrations)
+        else:
+            self.migrations = self.discover_migrations()
+
+    @staticmethod
+    def validate_migrations(
+        migrations: list[SchemaMigration],
+    ) -> list[SchemaMigration]:
+        """Validate and sort migrations chronologically by creation_timestamp.
+
+        Args:
+            migrations: List of SchemaMigration instances.
+
+        Returns:
+            Sorted list of SchemaMigration instances.
+
+        Raises:
+            ValueError: If migrations have duplicate creation_timestamps.
+        """
+        if not migrations:
+            return []
+
+        sorted_migrations = sorted(migrations, key=lambda m: m.creation_timestamp)
+
+        # Check that there are no duplicates
+        seen_timestamps: set[str] = set()
+        for migration in sorted_migrations:
+            ts = migration.creation_timestamp
+            if ts in seen_timestamps:
+                raise ValueError(
+                    f"Duplicate migration creation_timestamp detected: {ts}"
+                )
+            seen_timestamps.add(ts)
+
+        return sorted_migrations
+
+    @classmethod
+    def discover_migrations(cls) -> list[SchemaMigration]:
+        """Automatically discover and validate all SchemaMigration subclasses in migration_scripts.
+
+        Returns:
+            Sorted, validated list of SchemaMigration instances.
+        """
+        discovered: list[SchemaMigration] = []
+
+        # Only iterate through files from __path__.
+        # This prevents external path injection or local file execution when running a packaged release.
+        for module_info in pkgutil.iter_modules(
+            datacommons_db.migrations.migration_scripts.__path__
+        ):
+            if module_info.name.startswith("_"):
+                # Treat files starting with _ as private, do not import.
+                continue
+            module = importlib.import_module(
+                f"datacommons_db.migrations.migration_scripts.{module_info.name}"
+            )
+            for _, obj in inspect.getmembers(module, inspect.isclass):
+                if (
+                    issubclass(obj, SchemaMigration)
+                    and not inspect.isabstract(obj)
+                    and obj.__module__ == module.__name__  # Filter out imported objects
+                ):
+                    discovered.append(obj())
+
+        return cls.validate_migrations(discovered)
+
+    def get_applied_migrations(self) -> set[str]:
+        """Query the target database to determine the set of applied migration creation timestamps.
+
+        Returns:
+            Set of applied creation_timestamp strings, or empty set if SchemaMigrations table does not exist.
+
+        Raises:
+            RuntimeError: If querying SchemaMigrations fails unexpectedly.
+        """
+        if not self.spanner_client.table_exists("SchemaMigrations"):
+            return set()
+
+        get_applied_migrations_query = "SELECT CreationTimestamp FROM SchemaMigrations"
+        result = self.spanner_client.execute_query(get_applied_migrations_query)
+        if result.status != ExecutionStatus.SUCCESS:
+            raise RuntimeError(
+                f"Failed to query SchemaMigrations table: {result.error_message}"
+            )
+
+        applied_migrations: set[str] = set()
+        for row in result.rows:
+            # Check each row for validity.
+            if not isinstance(row, (list, tuple)) or len(row) == 0:
+                raise RuntimeError(
+                    f"Invalid or non-indexable row format in SchemaMigrations query results: {row!r}"
+                )
+            applied_migrations.add(str(row[0]))
+        return applied_migrations
+
+    def get_pending_migrations(self) -> list[SchemaMigration]:
+        """Get the list of pending migrations to be applied in chronological order.
+
+        Returns:
+            List of SchemaMigration instances that have not yet been applied.
+
+        Raises:
+            RuntimeError: If querying SchemaMigrations fails unexpectedly.
+        """
+        applied_migrations = self.get_applied_migrations()
+
+        known_timestamps = {m.creation_timestamp for m in self.migrations}
+        unknown_applied = applied_migrations - known_timestamps
+        if unknown_applied:
+            logger.warning(
+                "Database contains %d applied migration(s) not recognized by this codebase: %s. "
+                "The database schema may be newer than this application.",
+                len(unknown_applied),
+                sorted(unknown_applied),
+            )
+
+        return [
+            m for m in self.migrations if m.creation_timestamp not in applied_migrations
+        ]
+
+    def apply_migration(self, migration: SchemaMigration) -> MigrationResult:
+        """Apply a single migration and record its completion in SchemaMigrations.
+
+        Args:
+            migration: The SchemaMigration instance to apply.
+
+        Returns:
+            MigrationResult indicating successful execution.
+
+        Raises:
+            RuntimeError: If migration execution or recording in SchemaMigrations fails.
+        """
+        logger.info(
+            "Applying migration %s: %s",
+            migration.creation_timestamp,
+            migration.description,
+        )
+
+        # 1. Execute the migration's upgrade changes
+        migration.upgrade(self.spanner_client)
+
+        # 2. Record the applied migration in SchemaMigrations
+        self.record_applied_migration(
+            creation_timestamp=migration.creation_timestamp,
+            description=migration.description,
+        )
+
+        logger.info(
+            "Successfully applied migration %s",
+            migration.creation_timestamp,
+        )
+        return MigrationResult(
+            status=ExecutionStatus.SUCCESS,
+            creation_timestamp=migration.creation_timestamp,
+            description=migration.description,
+        )
+
+    def get_latest_applied_migration(self) -> str | None:
+        """Get the creation timestamp of the latest applied migration.
+
+        Returns:
+            The latest applied creation_timestamp string, or None if no migrations have been applied.
+
+        Raises:
+            RuntimeError: If querying SchemaMigrations fails unexpectedly.
+        """
+        applied = self.get_applied_migrations()
+        return max(applied) if applied else None
+
+    def record_applied_migration(
+        self, creation_timestamp: str, description: str
+    ) -> None:
+        """Record an applied migration in the SchemaMigrations table.
+
+        Args:
+            creation_timestamp: The migration creation_timestamp to record.
+            description: Description of the schema migration.
+
+        Raises:
+            RuntimeError: If inserting into SchemaMigrations fails.
+        """
+        insert_applied_migration_record_dml = (
+            "INSERT INTO SchemaMigrations (CreationTimestamp, AppliedTimestamp, Description) "
+            "VALUES (@creation_timestamp, PENDING_COMMIT_TIMESTAMP(), @description)"
+        )
+        params = {
+            "creation_timestamp": creation_timestamp,
+            "description": description,
+        }
+        param_types = {
+            "creation_timestamp": spanner.param_types.STRING,
+            "description": spanner.param_types.STRING,
+        }
+
+        dml_result = self.spanner_client.execute_dml(
+            insert_applied_migration_record_dml,
+            params=params,
+            param_types=param_types,
+        )
+        if dml_result.status != ExecutionStatus.SUCCESS:
+            raise RuntimeError(
+                f"Failed to record migration {creation_timestamp} in SchemaMigrations: "
+                f"{dml_result.error_message}"
+            )
+
+    def run_migrations(self) -> list[MigrationResult]:
+        """Execute all pending migrations in sequential chronological order.
+
+        Returns:
+            List of MigrationResult instances for all applied migrations.
+
+        Raises:
+            RuntimeError: If any migration fails during execution.
+        """
+        pending = self.get_pending_migrations()
+
+        if not pending:
+            logger.info("Database is already up-to-date. No migrations to apply.")
+            return []
+
+        results: list[MigrationResult] = []
+        for migration in pending:
+            try:
+                result = self.apply_migration(migration)
+                results.append(result)
+            except Exception as e:
+                logger.error(
+                    "Migration %s failed: %s. Halting migration sequence.",
+                    migration.creation_timestamp,
+                    e,
+                )
+                raise
+
+        return results