svc-infra 0.1.589__py3-none-any.whl → 0.1.706__py3-none-any.whl

svc_infra/db/nosql/repository.py

@@ -1,8 +1,19 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, cast
 
-from bson import ObjectId
+try:
+    from bson import ObjectId
+
+    _HAS_BSON = True
+except ModuleNotFoundError:
+    # `bson` is provided by the optional `pymongo` dependency.
+    # Keep imports working for non-mongo users/tests; runtime Mongo usage still
+    # requires installing pymongo.
+    _HAS_BSON = False
+
+    class ObjectId:  # type: ignore[no-redef]
+        pass
 
 
 class NoSqlRepository:
@@ -78,11 +89,13 @@ class NoSqlRepository:
         if not parts:
             return {}
         if len(parts) == 1:
-            return parts[0]  # type: ignore[return-value]
+            return parts[0]
         return {"$and": parts}
 
     def _normalize_id_value(self, val: Any) -> Any:
         """If we use Mongo’s _id and a string is passed, coerce to ObjectId when possible."""
+        if not _HAS_BSON:
+            return val
         if self.id_field == "_id" and isinstance(val, str):
             try:
                 return ObjectId(val)
@@ -91,13 +104,14 @@ class NoSqlRepository:
         return val
 
     @staticmethod
-    def _public_doc(doc: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
-        if not doc:
-            return doc
+    def _public_doc(doc: Dict[str, Any]) -> Dict[str, Any]:
         d = dict(doc)
         if "_id" in d and "id" not in d:
             _id = d.pop("_id", None)
-            d["id"] = str(_id) if isinstance(_id, ObjectId) else _id
+            if _HAS_BSON and isinstance(_id, ObjectId):
+                d["id"] = str(_id)
+            else:
+                d["id"] = _id
         return d
 
     async def list(
@@ -108,7 +122,7 @@ class NoSqlRepository:
         offset: int,
         sort: Optional[List[Tuple[str, int]]] = None,
         filter: Optional[Dict[str, Any]] = None,
-    ) -> List[Dict]:
+    ) -> List[Dict[str, Any]]:
         filt = self._merge_and(self._alive_filter(), filter)
         cursor = db[self.collection_name].find(filt).skip(offset).limit(limit)
         if sort:
@@ -117,12 +131,14 @@ class NoSqlRepository:
 
     async def count(self, db, *, filter: Optional[Dict[str, Any]] = None) -> int:
         filt = self._merge_and(self._alive_filter(), filter)
-        return await db[self.collection_name].count_documents(filt or {})
+        return cast(int, await db[self.collection_name].count_documents(filt or {}))
 
     async def get(self, db, id_value: Any) -> Dict | None:
         id_value = self._normalize_id_value(id_value)
         filt = self._merge_and(self._alive_filter(), {self.id_field: id_value})
         doc = await db[self.collection_name].find_one(filt)
+        if doc is None:
+            return None
         return self._public_doc(doc)
 
     async def create(self, db, data: Dict[str, Any]) -> Dict[str, Any]:
@@ -155,10 +171,10 @@ class NoSqlRepository:
             res = await db[self.collection_name].update_one(
                 {self.id_field: id_value}, {"$set": set_ops}
             )
-            return res.modified_count > 0
+            return cast(int, res.modified_count) > 0
 
         res = await db[self.collection_name].delete_one({self.id_field: id_value})
-        return res.deleted_count > 0
+        return cast(int, res.deleted_count) > 0
 
     async def search(
         self,
@@ -169,11 +185,13 @@ class NoSqlRepository:
         limit: int,
         offset: int,
         sort: Optional[List[Tuple[str, int]]] = None,
-    ) -> List[Dict]:
+    ) -> List[Dict[str, Any]]:
         regex = {"$regex": q, "$options": "i"}
         or_filter = [{"$or": [{f: regex} for f in fields]}] if fields else []
         filt = (
-            self._merge_and(self._alive_filter(), *or_filter) if or_filter else self._alive_filter()
+            self._merge_and(self._alive_filter(), *or_filter)
+            if or_filter
+            else self._alive_filter()
         )
         cursor = db[self.collection_name].find(filt).skip(offset).limit(limit)
         if sort:
@@ -184,9 +202,11 @@ class NoSqlRepository:
         regex = {"$regex": q, "$options": "i"}
         or_filter = {"$or": [{f: regex} for f in fields]} if fields else {}
         filt = self._merge_and(self._alive_filter(), or_filter)
-        return await db[self.collection_name].count_documents(filt or {})
+        return cast(int, await db[self.collection_name].count_documents(filt or {}))
 
     async def exists(self, db, *, where: Iterable[Dict[str, Any]]) -> bool:
         filt = self._merge_and(self._alive_filter(), *list(where))
-        doc = await db[self.collection_name].find_one(filt, projection={self.id_field: 1})
+        doc = await db[self.collection_name].find_one(
+            filt, projection={self.id_field: 1}
+        )
         return doc is not None

svc_infra/db/nosql/resource.py

@@ -1,9 +1,26 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, Optional, Sequence, Type, Union
-
-from pymongo import IndexModel
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Optional,
+    Sequence,
+    TYPE_CHECKING,
+    Type,
+    Union,
+)
+
+if TYPE_CHECKING:
+    from pymongo import IndexModel
+else:
+    try:
+        from pymongo import IndexModel
+    except ModuleNotFoundError:
+        # Minimal runtime stub so importing svc_infra works without optional Mongo deps.
+        class IndexModel:  # type: ignore[no-redef]
+            pass
 
 
 def _snake(name: str) -> str:

svc_infra/db/nosql/scaffold.py

@@ -6,7 +6,9 @@ from typing import Any, Dict, Literal, Optional
 from svc_infra.db.utils import normalize_dir, pascal, plural_snake, snake
 from svc_infra.utils import ensure_init_py, render_template, write
 
-_INIT_CONTENT_PAIRED = 'from . import documents, schemas\n\n__all__ = ["documents", "schemas"]\n'
+_INIT_CONTENT_PAIRED = (
+    'from . import documents, schemas\n\n__all__ = ["documents", "schemas"]\n'
+)
 _INIT_CONTENT_MINIMAL = "# package marker; add explicit exports here if desired\n"
 
 
@@ -46,7 +48,9 @@ def scaffold_core(
     schemas_txt = render_template(
         tmpl_dir="svc_infra.db.nosql.mongo.templates",
         name="schemas.py.tmpl",
-        subs={"Entity": ent},  # (only if your schemas.tmpl doesn't need collection_name)
+        subs={
+            "Entity": ent
+        },  # (only if your schemas.tmpl doesn't need collection_name)
     )
 
     if same_dir:
@@ -103,7 +107,9 @@ def scaffold_schemas_core(
     dest = normalize_dir(dest_dir)
     ent = pascal(entity_name)
     txt = render_template(
-        tmpl_dir="svc_infra.db.nosql.mongo.templates", name="schemas.py.tmpl", subs={"Entity": ent}
+        tmpl_dir="svc_infra.db.nosql.mongo.templates",
+        name="schemas.py.tmpl",
+        subs={"Entity": ent},
     )
     filename = schemas_filename or f"{snake(entity_name)}.py"
     res = write(dest / filename, txt, overwrite)

svc_infra/db/nosql/service.py

@@ -43,7 +43,9 @@ class NoSqlService:
     async def search(
         self, db, *, q: str, fields: Sequence[str], limit: int, offset: int, sort=None
     ):
-        return await self.repo.search(db, q=q, fields=fields, limit=limit, offset=offset, sort=sort)
+        return await self.repo.search(
+            db, q=q, fields=fields, limit=limit, offset=offset, sort=sort
+        )
 
     async def count_filtered(self, db, *, q: str, fields: Sequence[str]) -> int:
         return await self.repo.count_filtered(db, q=q, fields=fields)

svc_infra/db/nosql/types.py

@@ -11,7 +11,9 @@ class PyObjectId(ObjectId):
     """Pydantic v2-compatible ObjectId type."""
 
     @classmethod
-    def __get_pydantic_core_schema__(cls, _source_type: Any, _handler: GetCoreSchemaHandler):
+    def __get_pydantic_core_schema__(
+        cls, _source_type: Any, _handler: GetCoreSchemaHandler
+    ):
         def validate(v: Any) -> ObjectId:
             if isinstance(v, ObjectId):
                 return v
@@ -22,4 +24,6 @@ class PyObjectId(ObjectId):
                     raise ValueError(f"Invalid ObjectId: {v}") from e
             raise ValueError("ObjectId required")
 
-        return core_schema.no_info_after_validator_function(validate, core_schema.any_schema())
+        return core_schema.no_info_after_validator_function(
+            validate, core_schema.any_schema()
+        )

svc_infra/db/ops.py

@@ -0,0 +1,384 @@
+"""Database operations utilities for one-off administrative tasks.
+
+This module provides synchronous database utilities for operations that
+don't fit the normal async SQLAlchemy workflow, such as:
+- Waiting for database readiness at startup
+- Executing maintenance SQL
+- Dropping tables with lock handling
+- Terminating blocking queries
+
+These utilities use psycopg2 directly for maximum reliability in
+edge cases where the ORM might not be available or appropriate.
+
+Example:
+    >>> from svc_infra.db.ops import wait_for_database, run_sync_sql
+    >>>
+    >>> # Wait for database before app starts
+    >>> wait_for_database(timeout=30)
+    >>>
+    >>> # Run maintenance query
+    >>> run_sync_sql("VACUUM ANALYZE my_table")
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import time
+from typing import Any, Optional, Sequence, cast
+
+from .sql.utils import get_database_url_from_env
+
+logger = logging.getLogger(__name__)
+
+# Timeout for individual database operations (seconds)
+DEFAULT_STATEMENT_TIMEOUT = 30
+
+# Default wait-for-database settings
+DEFAULT_WAIT_TIMEOUT = 30
+DEFAULT_WAIT_INTERVAL = 1.0
+
+
+def _flush() -> None:
+    """Force flush stdout/stderr for containerized log visibility."""
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+
+def _get_connection(url: Optional[str] = None, connect_timeout: int = 10) -> Any:
+    """
+    Get a psycopg2 connection.
+
+    Args:
+        url: Database URL. If None, resolved from environment.
+        connect_timeout: Connection timeout in seconds.
+
+    Returns:
+        psycopg2 connection object
+
+    Raises:
+        ImportError: If psycopg2 is not installed
+        RuntimeError: If no database URL is available
+    """
+    try:
+        import psycopg2
+    except ImportError as e:
+        raise ImportError(
+            "psycopg2 is required for db.ops utilities. "
+            "Install with: pip install psycopg2-binary"
+        ) from e
+
+    if url is None:
+        url = get_database_url_from_env(required=True)
+
+    # Add connect_timeout to connection options
+    return psycopg2.connect(url, connect_timeout=connect_timeout)
+
+
+def wait_for_database(
+    url: Optional[str] = None,
+    timeout: float = DEFAULT_WAIT_TIMEOUT,
+    interval: float = DEFAULT_WAIT_INTERVAL,
+    verbose: bool = True,
+) -> bool:
+    """
+    Wait for database to be ready, with retries.
+
+    Useful for container startup scripts where the database may not
+    be immediately available.
+
+    Args:
+        url: Database URL. If None, resolved from environment.
+        timeout: Maximum time to wait in seconds (default: 30)
+        interval: Time between retry attempts in seconds (default: 1.0)
+        verbose: If True, log progress messages
+
+    Returns:
+        True if database is ready, False if timeout reached
+
+    Example:
+        >>> # In container startup script
+        >>> if not wait_for_database(timeout=60):
+        ...     sys.exit(1)
+        >>> # Database is ready, continue with app startup
+    """
+    if url is None:
+        url = get_database_url_from_env(required=True)
+
+    start = time.monotonic()
+    attempt = 0
+
+    while True:
+        attempt += 1
+        elapsed = time.monotonic() - start
+
+        if elapsed >= timeout:
+            if verbose:
+                logger.error(
+                    f"Database not ready after {timeout}s ({attempt} attempts)"
+                )
+                _flush()
+            return False
+
+        try:
+            conn = _get_connection(url, connect_timeout=min(5, int(timeout - elapsed)))
+            conn.close()
+            if verbose:
+                logger.info(f"Database ready after {elapsed:.1f}s ({attempt} attempts)")
+                _flush()
+            return True
+        except Exception as e:
+            if verbose:
+                remaining = timeout - elapsed
+                logger.debug(
+                    f"Database not ready ({e}), retrying... ({remaining:.0f}s remaining)"
+                )
+                _flush()
+            time.sleep(interval)
+
+
+def run_sync_sql(
+    sql: str,
+    params: Optional[Sequence[Any]] = None,
+    url: Optional[str] = None,
+    timeout: int = DEFAULT_STATEMENT_TIMEOUT,
+    fetch: bool = False,
+) -> Optional[list[tuple[Any, ...]]]:
+    """
+    Execute SQL synchronously with a statement timeout.
+
+    This is useful for one-off administrative queries that don't fit
+    the normal async SQLAlchemy workflow.
+
+    Args:
+        sql: SQL statement to execute
+        params: Optional parameters for parameterized queries
+        url: Database URL. If None, resolved from environment.
+        timeout: Statement timeout in seconds (default: 30)
+        fetch: If True, return fetched rows; if False, return None
+
+    Returns:
+        List of tuples if fetch=True, otherwise None
+
+    Raises:
+        psycopg2.Error: On database errors
+        TimeoutError: If statement exceeds timeout
+
+    Example:
+        >>> # Run a maintenance query
+        >>> run_sync_sql("VACUUM ANALYZE users")
+        >>>
+        >>> # Fetch data with timeout
+        >>> rows = run_sync_sql(
+        ...     "SELECT id, name FROM users WHERE active = %s",
+        ...     params=(True,),
+        ...     fetch=True,
+        ...     timeout=10
+        ... )
+    """
+    conn = _get_connection(url)
+    try:
+        with conn.cursor() as cur:
+            # Set statement timeout (PostgreSQL-specific)
+            cur.execute(f"SET statement_timeout = '{timeout}s'")
+
+            if params:
+                cur.execute(sql, params)
+            else:
+                cur.execute(sql)
+
+            if fetch:
+                return cast(list[tuple[Any, ...]], cur.fetchall())
+
+            conn.commit()
+            return None
+    finally:
+        conn.close()
+
+
+def kill_blocking_queries(
+    table_name: str,
+    url: Optional[str] = None,
+    timeout: int = DEFAULT_STATEMENT_TIMEOUT,
+    dry_run: bool = False,
+) -> list[dict[str, Any]]:
+    """
+    Terminate queries blocking operations on a specific table.
+
+    This is useful before DROP TABLE or ALTER TABLE operations that
+    might be blocked by long-running queries or idle transactions.
+
+    Args:
+        table_name: Name of the table (can include schema as 'schema.table')
+        url: Database URL. If None, resolved from environment.
+        timeout: Statement timeout in seconds (default: 30)
+        dry_run: If True, only report blocking queries without terminating
+
+    Returns:
+        List of dicts with info about terminated (or found) queries:
+        [{"pid": 123, "query": "SELECT...", "state": "active", "terminated": True}]
+
+    Example:
+        >>> # Check what would be terminated
+        >>> blocking = kill_blocking_queries("embeddings", dry_run=True)
+        >>> print(f"Found {len(blocking)} blocking queries")
+        >>>
+        >>> # Actually terminate them
+        >>> kill_blocking_queries("embeddings")
+    """
+    # Query to find blocking queries on a table
+    find_blocking_sql = """
+        SELECT pid, state, query, age(clock_timestamp(), query_start) as duration
+        FROM pg_stat_activity
+        WHERE pid != pg_backend_pid()
+          AND state != 'idle'
+          AND (
+              query ILIKE %s
+              OR query ILIKE %s
+              OR query ILIKE %s
+          )
+        ORDER BY query_start;
+    """
+
+    # Patterns to match queries involving the table
+    patterns = (
+        f"%{table_name}%",
+        f"%{table_name.split('.')[-1]}%",  # Just table name without schema
+        f"%{table_name.replace('.', '%')}%",  # Handle schema.table pattern
+    )
+
+    conn = _get_connection(url)
+    terminated: list[dict[str, Any]] = []
+
+    try:
+        with conn.cursor() as cur:
+            cur.execute(f"SET statement_timeout = '{timeout}s'")
+            cur.execute(find_blocking_sql, patterns)
+            rows = cur.fetchall()
+
+            for pid, state, query, duration in rows:
+                info = {
+                    "pid": pid,
+                    "state": state,
+                    "query": query[:200] + "..." if len(query) > 200 else query,
+                    "duration": str(duration),
+                    "terminated": False,
+                }
+
+                if not dry_run:
+                    try:
+                        cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
+                        info["terminated"] = True
+                        logger.info(f"Terminated query PID {pid}: {query[:100]}...")
+                    except Exception as e:
+                        logger.warning(f"Failed to terminate PID {pid}: {e}")
+                        info["error"] = str(e)
+
+                terminated.append(info)
+
+            conn.commit()
+    finally:
+        conn.close()
+
+    _flush()
+    return terminated
+
+
+def drop_table_safe(
+    table_name: str,
+    url: Optional[str] = None,
+    timeout: int = DEFAULT_STATEMENT_TIMEOUT,
+    kill_blocking: bool = True,
+    if_exists: bool = True,
+    cascade: bool = False,
+) -> bool:
+    """
+    Drop a table safely with lock handling.
+
+    Handles common issues with DROP TABLE:
+    - Terminates blocking queries first (optional)
+    - Uses statement timeout to avoid hanging
+    - Handles 'table does not exist' gracefully
+
+    Args:
+        table_name: Name of table to drop (can include schema)
+        url: Database URL. If None, resolved from environment.
+        timeout: Statement timeout in seconds (default: 30)
+        kill_blocking: If True, terminate blocking queries first (default: True)
+        if_exists: If True, don't error if table doesn't exist (default: True)
+        cascade: If True, drop dependent objects (default: False)
+
+    Returns:
+        True if table was dropped (or didn't exist), False on error
+
+    Example:
+        >>> # Drop table, killing any blocking queries first
+        >>> drop_table_safe("embeddings", cascade=True)
+        True
+        >>>
+        >>> # Safe to call even if table doesn't exist
+        >>> drop_table_safe("nonexistent_table")
+        True
+    """
+    if url is None:
+        url = get_database_url_from_env(required=True)
+
+    # Kill blocking queries first if requested
+    if kill_blocking:
+        blocked = kill_blocking_queries(table_name, url=url, timeout=timeout)
+        if blocked:
+            logger.info(f"Terminated {len(blocked)} blocking queries before DROP")
+            # Brief pause to let connections clean up
+            time.sleep(0.5)
+
+    # Build DROP statement
+    drop_sql = "DROP TABLE"
+    if if_exists:
+        drop_sql += " IF EXISTS"
+    drop_sql += f" {table_name}"
+    if cascade:
+        drop_sql += " CASCADE"
+
+    try:
+        run_sync_sql(drop_sql, url=url, timeout=timeout)
+        logger.info(f"Dropped table: {table_name}")
+        _flush()
+        return True
+    except Exception as e:
+        logger.error(f"Failed to drop table {table_name}: {e}")
+        _flush()
+        return False
+
+
+def get_database_url(
+    required: bool = True,
+    normalize: bool = True,
+) -> Optional[str]:
+    """
+    Convenience wrapper for get_database_url_from_env().
+
+    This is the recommended way to get the database URL, as it
+    handles all common environment variable names and normalizations.
+
+    Args:
+        required: If True, raise RuntimeError when no URL is found
+        normalize: If True, convert postgres:// to postgresql://
+
+    Returns:
+        Database URL string, or None if not found and not required
+
+    Example:
+        >>> url = get_database_url()
+        >>> print(url)
+        'postgresql://user:pass@host:5432/db'
+    """
+    return get_database_url_from_env(required=required, normalize=normalize)
+
+
+__all__ = [
+    "wait_for_database",
+    "run_sync_sql",
+    "kill_blocking_queries",
+    "drop_table_safe",
+    "get_database_url",
+]