weaverstack 0.1.1__py3-none-any.whl

weaver/sql/__init__.py

@@ -0,0 +1,40 @@
+"""Driver-neutral Warehouse SQL execution backed by ``mssql-python``."""
+
+from __future__ import annotations
+
+from .authentication import AccessTokenAuthentication, SqlAuthentication
+from .connection import SqlEndpoint, build_connection_string, connect
+from .errors import (
+    SqlConnectionError,
+    SqlError,
+    SqlExecutionError,
+    SqlPoolClosedError,
+)
+from .execution import PooledSqlExecutor, SqlExecutor, SqlRow
+from .pool import (
+    DEFAULT_MAX_CONNECTIONS,
+    SqlConnectionLease,
+    SqlConnectionPool,
+    SqlPoolRegistry,
+)
+from .wipe import generate_warehouse_wipe_sql
+
+__all__ = [
+    "AccessTokenAuthentication",
+    "SqlAuthentication",
+    "SqlEndpoint",
+    "build_connection_string",
+    "connect",
+    "SqlError",
+    "SqlConnectionError",
+    "SqlExecutionError",
+    "SqlPoolClosedError",
+    "SqlExecutor",
+    "SqlRow",
+    "PooledSqlExecutor",
+    "SqlConnectionLease",
+    "SqlConnectionPool",
+    "SqlPoolRegistry",
+    "DEFAULT_MAX_CONNECTIONS",
+    "generate_warehouse_wipe_sql",
+]

weaver/sql/authentication.py

@@ -0,0 +1,38 @@
+"""Authentication records consumed by the common SQL connector.
+
+Identity policy stays at the caller boundary.  This module only converts a
+fresh access token into the connection argument understood by
+``mssql-python``.
+"""
+
+from __future__ import annotations
+
+import struct
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Protocol
+
+# SQL_COPT_SS_ACCESS_TOKEN from msodbcsql.h.  mssql-python accepts the same
+# attrs_before shape as pyodbc and exposes this value as
+# ConstantsDDBC.SQL_COPT_SS_ACCESS_TOKEN.
+SQL_ACCESS_TOKEN_ATTRIBUTE = 1256
+
+
+class SqlAuthentication(Protocol):
+    """Authentication material for one new physical connection."""
+
+    def connection_arguments(self) -> Mapping[str, object]:
+        """Return current driver keyword arguments."""
+
+
+@dataclass(frozen=True)
+class AccessTokenAuthentication:
+    """Acquire and encode a fresh Entra access token per connection."""
+
+    token_provider: Callable[[], str]
+
+    def connection_arguments(self) -> Mapping[str, object]:
+        token = self.token_provider()
+        token_bytes = token.encode("utf-16-le")
+        packed = struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
+        return {"attrs_before": {SQL_ACCESS_TOKEN_ATTRIBUTE: packed}}

weaver/sql/connection.py

@@ -0,0 +1,90 @@
+"""Resolved SQL endpoints and common ``mssql-python`` connection creation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .authentication import SqlAuthentication
+from .errors import SqlConnectionError
+
+DEFAULT_SQL_PORT = 1433
+DEFAULT_CONNECT_TIMEOUT = 60
+
+
+@dataclass(frozen=True)
+class SqlEndpoint:
+    """One resolved Warehouse SQL endpoint.
+
+    The optional Fabric identities make the pool key stable without treating a
+    display name or a credential-bearing connection string as identity.
+    """
+
+    server: str
+    database: str
+    port: int = DEFAULT_SQL_PORT
+    workspace_id: str | None = None
+    warehouse_id: str | None = None
+    warehouse_name: str | None = None
+
+    def __post_init__(self) -> None:
+        if not self.server.strip():
+            raise ValueError("SQL server must not be empty")
+        if not self.database.strip():
+            raise ValueError("SQL database must not be empty")
+
+    @property
+    def pool_key(self) -> tuple[str, str, str, str, int]:
+        return (
+            self.workspace_id or "",
+            self.warehouse_id or "",
+            self.server.lower(),
+            self.database.lower(),
+            self.port,
+        )
+
+    def __str__(self) -> str:
+        return f"{self.server}:{self.port}/{self.database}"
+
+
+def build_connection_string(endpoint: SqlEndpoint) -> str:
+    """The validated Fabric Warehouse connection-string shape."""
+
+    return (
+        f"Server={endpoint.server},{endpoint.port};"
+        f"Database={endpoint.database};"
+        "Encrypt=yes;"
+        "TrustServerCertificate=no;"
+    )
+
+
+DriverConnector = Callable[..., Any]
+
+
+def connect(
+    endpoint: SqlEndpoint,
+    authentication: SqlAuthentication,
+    *,
+    timeout: int = DEFAULT_CONNECT_TIMEOUT,
+    connector: DriverConnector | None = None,
+):
+    """Open one physical connection using current authentication material."""
+
+    if connector is None:
+        try:
+            from mssql_python import connect as connector
+        except ImportError as exc:  # pragma: no cover - installation dependent
+            raise SqlConnectionError(
+                "mssql-python is required for Warehouse SQL execution"
+            ) from exc
+
+    try:
+        return connector(
+            build_connection_string(endpoint),
+            timeout=timeout,
+            **dict(authentication.connection_arguments()),
+        )
+    except SqlConnectionError:
+        raise
+    except Exception as exc:
+        raise SqlConnectionError(f"failed to connect to {endpoint}: {exc}") from exc

weaver/sql/errors.py

@@ -0,0 +1,25 @@
+"""Weaver-facing SQL errors.
+
+Driver exceptions do not escape this package.  The normalised errors retain the
+selected endpoint and chain the original exception for diagnosis.
+"""
+
+from __future__ import annotations
+
+from ..errors import WeaverError
+
+
+class SqlError(WeaverError):
+    """Base class for SQL connection and execution failures."""
+
+
+class SqlConnectionError(SqlError):
+    """Raised when a physical SQL connection cannot be opened."""
+
+
+class SqlExecutionError(SqlError):
+    """Raised when a SQL statement, script, or query fails."""
+
+
+class SqlPoolClosedError(SqlError):
+    """Raised when a caller tries to lease from a closed pool."""

weaver/sql/execution.py

@@ -0,0 +1,123 @@
+"""Shared SQL statement, script, query, and transaction handling."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+from .errors import SqlError, SqlExecutionError
+from .pool import SqlConnectionPool
+
+SqlRow = dict[str, Any]
+
+
+class SqlExecutor(Protocol):
+    """The SQL surface used by Weaver operations."""
+
+    def execute(
+        self, statement: str, parameters: Sequence[object] | None = None
+    ) -> None: ...
+
+    def execute_script(self, script: str) -> None: ...
+
+    def query(
+        self, statement: str, parameters: Sequence[object] | None = None
+    ) -> Sequence[SqlRow]: ...
+
+
+class PooledSqlExecutor:
+    """Execute through one owned or injected bounded connection pool."""
+
+    def __init__(self, pool: SqlConnectionPool, *, owns_pool: bool = False) -> None:
+        self.pool = pool
+        self.owns_pool = owns_pool
+
+    def execute(
+        self, statement: str, parameters: Sequence[object] | None = None
+    ) -> None:
+        self._run(statement, parameters=parameters, query=False, drain=True)
+
+    def execute_script(self, script: str) -> None:
+        self._run(script, parameters=None, query=False, drain=True)
+
+    def query(
+        self, statement: str, parameters: Sequence[object] | None = None
+    ) -> Sequence[SqlRow]:
+        return self._run(statement, parameters=parameters, query=True, drain=False)
+
+    def _run(
+        self,
+        statement: str,
+        *,
+        parameters: Sequence[object] | None,
+        query: bool,
+        drain: bool,
+    ):
+        with self.pool.lease() as lease:
+            connection = lease.connection
+            cursor = None
+            try:
+                cursor = connection.cursor()
+                if parameters is None:
+                    cursor.execute(statement)
+                else:
+                    cursor.execute(statement, tuple(parameters))
+
+                if query:
+                    if cursor.description is None:
+                        rows = []
+                    else:
+                        columns = [column[0] for column in cursor.description]
+                        rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
+                    connection.commit()
+                    return rows
+
+                if drain:
+                    _drain(cursor)
+                connection.commit()
+                return None
+            except SqlError:
+                lease.discard()
+                _rollback(connection)
+                raise
+            except Exception as exc:
+                lease.discard()
+                _rollback(connection)
+                operation = "query" if query else "SQL execution"
+                raise SqlExecutionError(
+                    f"{operation} failed on {self.pool.endpoint}: {exc}"
+                ) from exc
+            finally:
+                if cursor is not None:
+                    try:
+                        cursor.close()
+                    except Exception:
+                        lease.discard()
+
+    def close(self) -> None:
+        if self.owns_pool:
+            self.pool.close()
+
+    def __enter__(self) -> "PooledSqlExecutor":
+        return self
+
+    def __exit__(self, *exc) -> bool:
+        self.close()
+        return False
+
+
+def _drain(cursor) -> None:
+    """Consume all result sets so multi-statement T-SQL can commit reliably."""
+
+    while True:
+        if cursor.description is not None:
+            cursor.fetchall()
+        if not cursor.nextset():
+            return
+
+
+def _rollback(connection) -> None:
+    try:
+        connection.rollback()
+    except Exception:
+        pass

weaver/sql/pool.py

@@ -0,0 +1,174 @@
+"""Small, bounded, thread-safe SQL connection pools."""
+
+from __future__ import annotations
+
+import threading
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator
+
+from .authentication import SqlAuthentication
+from .connection import SqlEndpoint, connect
+from .errors import SqlPoolClosedError
+
+DEFAULT_MAX_CONNECTIONS = 4
+ConnectionFactory = Callable[[SqlEndpoint, SqlAuthentication], Any]
+
+
+@dataclass
+class SqlConnectionLease:
+    """One exclusively leased physical connection."""
+
+    connection: Any
+    usable: bool = True
+
+    def discard(self) -> None:
+        """Prevent this connection from returning to the idle pool."""
+
+        self.usable = False
+
+
+class SqlConnectionPool:
+    """Bounded connection reuse owned by one workflow or Fabric session."""
+
+    def __init__(
+        self,
+        endpoint: SqlEndpoint,
+        authentication: SqlAuthentication,
+        *,
+        max_connections: int = DEFAULT_MAX_CONNECTIONS,
+        connection_factory: ConnectionFactory = connect,
+    ) -> None:
+        if max_connections < 1:
+            raise ValueError("max_connections must be at least one")
+        self.endpoint = endpoint
+        self.authentication = authentication
+        self.max_connections = max_connections
+        self._connection_factory = connection_factory
+        self._condition = threading.Condition()
+        self._idle: list[Any] = []
+        self._physical_count = 0
+        self._closed = False
+
+    @contextmanager
+    def lease(self) -> Iterator[SqlConnectionLease]:
+        """Lease one connection exclusively, waiting when the pool is full."""
+
+        connection = self._acquire()
+        lease = SqlConnectionLease(connection)
+        try:
+            yield lease
+        finally:
+            self._release(lease)
+
+    def _acquire(self):
+        create = False
+        with self._condition:
+            while True:
+                if self._closed:
+                    raise SqlPoolClosedError(f"SQL pool for {self.endpoint} is closed")
+                if self._idle:
+                    return self._idle.pop()
+                if self._physical_count < self.max_connections:
+                    # Reserve the slot before opening outside the lock.
+                    self._physical_count += 1
+                    create = True
+                    break
+                self._condition.wait()
+
+        if create:
+            try:
+                return self._connection_factory(self.endpoint, self.authentication)
+            except Exception:
+                with self._condition:
+                    self._physical_count -= 1
+                    self._condition.notify()
+                raise
+        raise AssertionError("unreachable")
+
+    def _release(self, lease: SqlConnectionLease) -> None:
+        close = False
+        with self._condition:
+            if self._closed or not lease.usable:
+                self._physical_count -= 1
+                close = True
+            else:
+                self._idle.append(lease.connection)
+            self._condition.notify()
+        if close:
+            _close(lease.connection)
+
+    def close(self) -> None:
+        """Close idle connections; leased connections close when returned."""
+
+        with self._condition:
+            if self._closed:
+                return
+            self._closed = True
+            idle, self._idle = self._idle, []
+            self._physical_count -= len(idle)
+            self._condition.notify_all()
+        for connection in idle:
+            _close(connection)
+
+    def __enter__(self) -> "SqlConnectionPool":
+        return self
+
+    def __exit__(self, *exc) -> bool:
+        self.close()
+        return False
+
+
+class SqlPoolRegistry:
+    """One endpoint-keyed pool registry for an owning execution context."""
+
+    def __init__(
+        self,
+        *,
+        max_connections: int = DEFAULT_MAX_CONNECTIONS,
+        connection_factory: ConnectionFactory = connect,
+    ) -> None:
+        self.max_connections = max_connections
+        self.connection_factory = connection_factory
+        self._lock = threading.Lock()
+        self._pools: dict[tuple[str, str, str, str, int], SqlConnectionPool] = {}
+        self._closed = False
+
+    def pool(
+        self, endpoint: SqlEndpoint, authentication: SqlAuthentication
+    ) -> SqlConnectionPool:
+        with self._lock:
+            if self._closed:
+                raise SqlPoolClosedError("SQL pool registry is closed")
+            key = endpoint.pool_key
+            if key not in self._pools:
+                self._pools[key] = SqlConnectionPool(
+                    endpoint,
+                    authentication,
+                    max_connections=self.max_connections,
+                    connection_factory=self.connection_factory,
+                )
+            return self._pools[key]
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            pools, self._pools = tuple(self._pools.values()), {}
+        for pool in pools:
+            pool.close()
+
+    def __enter__(self) -> "SqlPoolRegistry":
+        return self
+
+    def __exit__(self, *exc) -> bool:
+        self.close()
+        return False
+
+
+def _close(connection: Any) -> None:
+    try:
+        connection.close()
+    except Exception:
+        pass

weaver/sql/wipe.py

@@ -0,0 +1,156 @@
+"""Pure generation of the proven Fabric Warehouse wipe batch."""
+
+from __future__ import annotations
+
+
+def generate_warehouse_wipe_sql() -> str:
+    """Return deterministic SQL that removes supported user-created objects.
+
+    This is ported from the legacy Weaver Warehouse implementation:
+    ``src/weaver_runtime/dbrep/sql/templates/admin/wipe.sql`` for object
+    enumeration and ordering, plus ``_drop_user_schemas`` in
+    ``src/weaver_runtime/dbrep/sql/backend.py`` for schema cleanup.
+    """
+
+    return _WAREHOUSE_WIPE_SQL
+
+
+_WAREHOUSE_WIPE_SQL = """\
+set nocount on;
+
+declare @weaver_sql nvarchar(max);
+
+/* Foreign keys first so dependent tables can be dropped. */
+select
+    @weaver_sql = string_agg(
+        convert(
+            nvarchar(max),
+            N'alter table '
+            + quotename(object_schema_name(fk.parent_object_id))
+            + N'.'
+            + quotename(object_name(fk.parent_object_id))
+            + N' drop constraint '
+            + quotename(fk.name)
+            + N';'
+        ),
+        char(10)
+    ) within group (order by fk.object_id)
+from sys.foreign_keys as fk
+where lower(schema_name(schema_id)) not in
+    (N'guest', N'information_schema', N'sys', N'queryinsights');
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+
+/* Views before the objects and tables they can depend on. */
+set @weaver_sql = null;
+select
+    @weaver_sql = string_agg(
+        convert(
+            nvarchar(max),
+            N'drop view '
+            + quotename(schema_name(schema_id))
+            + N'.'
+            + quotename(name)
+            + N';'
+        ),
+        char(10)
+    ) within group (order by object_id)
+from sys.views
+where lower(schema_name(schema_id)) not in
+    (N'guest', N'information_schema', N'sys', N'queryinsights');
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+
+/* Stored procedures. */
+set @weaver_sql = null;
+select
+    @weaver_sql = string_agg(
+        convert(
+            nvarchar(max),
+            N'drop procedure '
+            + quotename(schema_name(schema_id))
+            + N'.'
+            + quotename(name)
+            + N';'
+        ),
+        char(10)
+    ) within group (order by object_id)
+from sys.procedures
+where lower(schema_name(schema_id)) not in
+    (N'guest', N'information_schema', N'sys', N'queryinsights');
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+
+/* Functions (scalar, inline table-valued, multi-statement table-valued). */
+set @weaver_sql = null;
+select
+    @weaver_sql = string_agg(
+        convert(
+            nvarchar(max),
+            N'drop function '
+            + quotename(schema_name(schema_id))
+            + N'.'
+            + quotename(name)
+            + N';'
+        ),
+        char(10)
+    ) within group (order by object_id)
+from sys.objects
+where type in (N'FN', N'IF', N'TF', N'FS', N'FT')
+    and lower(schema_name(schema_id)) not in
+        (N'guest', N'information_schema', N'sys', N'queryinsights');
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+
+/* Tables. */
+set @weaver_sql = null;
+select
+    @weaver_sql = string_agg(
+        convert(
+            nvarchar(max),
+            N'drop table '
+            + quotename(schema_name(schema_id))
+            + N'.'
+            + quotename(name)
+            + N';'
+        ),
+        char(10)
+    ) within group (order by object_id)
+from sys.tables
+where lower(schema_name(schema_id)) not in
+    (N'guest', N'information_schema', N'sys', N'queryinsights');
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+
+/* User schemas last; the built-in and Fabric-owned schemas survive. */
+set @weaver_sql = null;
+select
+    @weaver_sql = string_agg(
+        convert(nvarchar(max), N'drop schema ' + quotename(name) + N';'),
+        char(10)
+    ) within group (order by schema_id)
+from sys.schemas
+where lower(name) not in
+    (N'dbo', N'guest', N'information_schema', N'sys', N'queryinsights', N'_rsc')
+    and schema_id < 16384;
+
+if @weaver_sql is not null
+begin
+    exec sys.sp_executesql @weaver_sql;
+end;
+"""