sqllocks-shape-sqlserver 0.9.1__py3-none-any.whl

shape_sqlserver/__init__.py

@@ -0,0 +1,28 @@
+"""Shape plugin: SQL Server, Azure SQL and Fabric SQL.
+
+* ``shape.sources``: ``mssql`` reads one table as Arrow record batches.
+* ``shape.commands``: ``shape profile-db`` profiles a database schema.
+* helpers shared with ``shape-fabric``: :mod:`shape_sqlserver.sql` (identifier quoting,
+  connection strings, catalog queries, type maps) and :mod:`shape_sqlserver.auth`
+  (connections, Microsoft Entra tokens).
+
+``pyodbc`` and ``azure-identity`` load only when a connection is opened.
+"""
+
+from .auth import Credentials, connect
+from .command import ProfileDbCommand
+from .profiler import profile_database
+from .source import SqlServerSource
+from .sql import SqlServerError
+
+SHAPE_API = "1.0"
+
+__all__ = [
+    "SHAPE_API",
+    "Credentials",
+    "ProfileDbCommand",
+    "SqlServerError",
+    "SqlServerSource",
+    "connect",
+    "profile_database",
+]

shape_sqlserver/auth.py

@@ -0,0 +1,130 @@
+"""Connections to SQL Server, Azure SQL and Fabric SQL: SQL logins and Microsoft Entra tokens.
+
+``pyodbc`` and ``azure-identity`` are imported only when a connection is made, so the plugin
+loads (and ``shape plugins doctor`` is clean) on a machine that has neither.
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+from typing import Any
+
+from .sql import SqlServerError, redact_connection_string, register_converters
+
+AUTH_METHODS = ("sql", "cli", "msi", "spn", "fabric")
+TOKEN_SCOPE = "https://database.windows.net/.default"
+FABRIC_TOKEN_AUDIENCE = "https://database.windows.net/"
+# ODBC connection attribute that carries an access token (SQL_COPT_SS_ACCESS_TOKEN).
+ACCESS_TOKEN_ATTR = 1256
+
+
+@dataclass(frozen=True, slots=True, repr=False)
+class Credentials:
+    """How to authenticate.
+
+    ``sql``: the login is in the connection string. ``cli``: the Azure CLI's signed-in account.
+    ``msi``: managed identity (or the default Azure credential chain). ``spn``: a service
+    principal (``tenant_id``, ``client_id``, ``client_secret``). ``fabric``: the token of the
+    running Fabric notebook.
+    """
+
+    method: str = "cli"
+    tenant_id: str | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
+
+    def __repr__(self) -> str:  # never show the secret
+        return (
+            f"Credentials(method={self.method!r}, tenant_id={self.tenant_id!r}, "
+            f"client_id={self.client_id!r}, client_secret={'***' if self.client_secret else None})"
+        )
+
+    def validate(self) -> None:
+        if self.method not in AUTH_METHODS:
+            raise SqlServerError(
+                f"unsupported auth method {self.method!r}; use one of {', '.join(AUTH_METHODS)}"
+            )
+        if self.method == "spn" and not (self.tenant_id and self.client_id and self.client_secret):
+            raise SqlServerError("auth 'spn' needs a tenant id, a client id and a client secret")
+
+
+def _credential(creds: Credentials) -> Any:
+    try:
+        from azure import identity  # type: ignore[import-not-found,unused-ignore]
+    except ImportError as exc:
+        raise SqlServerError(
+            "Entra authentication needs azure-identity: "
+            "pip install 'sqllocks-shape-sqlserver[entra]'"
+        ) from exc
+    if creds.method == "cli":
+        return identity.AzureCliCredential()
+    if creds.method == "msi":
+        return identity.DefaultAzureCredential(exclude_managed_identity_credential=False)
+    return identity.ClientSecretCredential(
+        tenant_id=creds.tenant_id, client_id=creds.client_id, client_secret=creds.client_secret
+    )
+
+
+def _fabric_token() -> str:
+    try:
+        try:
+            from notebookutils import mssparkutils  # type: ignore[import-not-found,unused-ignore]
+        except ImportError:
+            import mssparkutils  # type: ignore[import-not-found,no-redef,unused-ignore]
+        token: str = mssparkutils.credentials.getToken(FABRIC_TOKEN_AUDIENCE)
+    except ImportError as exc:
+        raise SqlServerError("auth 'fabric' only works inside a Fabric notebook") from exc
+    return token
+
+
+def access_token(creds: Credentials) -> bytes:
+    """The Entra access token for ``creds``, UTF-16-LE encoded as the ODBC driver expects."""
+    creds.validate()
+    if creds.method == "sql":
+        raise SqlServerError("auth 'sql' uses the login in the connection string, not a token")
+    if creds.method == "fabric":
+        return _fabric_token().encode("utf-16-le")
+    token: str = _credential(creds).get_token(TOKEN_SCOPE).token
+    return token.encode("utf-16-le")
+
+
+def token_struct(token: bytes) -> bytes:
+    """``token`` packed as the ODBC ``SQL_COPT_SS_ACCESS_TOKEN`` structure."""
+    return struct.pack(f"<I{len(token)}s", len(token), token)
+
+
+def connect(connection_string: str, creds: Credentials | None = None, *, timeout: int = 30) -> Any:
+    """An open pyodbc connection.
+
+    The connection is in autocommit mode: Shape only reads, and an implicit transaction left
+    open on a pooled connection would block other sessions.
+
+    With ``sql`` the connection string carries the login. With any other method an Entra
+    token is fetched and passed to the driver; the connection string must then not hold a
+    ``UID``, ``PWD`` or ``Authentication`` key.
+    """
+    creds = creds or Credentials()
+    creds.validate()
+    try:
+        import pyodbc  # type: ignore[import-not-found,unused-ignore]
+    except ImportError as exc:
+        raise SqlServerError(
+            "pyodbc is not available. Install it and the Microsoft ODBC Driver 18 for SQL Server "
+            "(on Linux also unixODBC)."
+        ) from exc
+    try:
+        if creds.method == "sql":
+            conn = pyodbc.connect(connection_string, timeout=timeout, autocommit=True)
+        else:
+            token = token_struct(access_token(creds))
+            conn = pyodbc.connect(
+                connection_string,
+                attrs_before={ACCESS_TOKEN_ATTR: token},
+                timeout=timeout,
+                autocommit=True,
+            )
+    except pyodbc.Error as exc:
+        raise SqlServerError(f"could not connect: {redact_connection_string(str(exc))}") from None
+    register_converters(conn)
+    return conn

shape_sqlserver/catalog.py

@@ -0,0 +1,181 @@
+"""Schema walk: tables, columns, keys and row counts from the SQL Server catalog views."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Collection, Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+from .sql import (
+    COLUMNS_QUERY,
+    FOREIGN_KEY_QUERY,
+    PRIMARY_KEY_QUERY,
+    ROW_COUNT_QUERY,
+    TABLES_QUERY,
+    fetch_dicts,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ColumnInfo:
+    name: str
+    type_name: str
+    max_length: int
+    precision: int
+    scale: int
+    is_nullable: bool
+    is_identity: bool
+
+
+@dataclass(slots=True)
+class TableInfo:
+    schema: str
+    name: str
+    object_id: int
+    columns: list[ColumnInfo]
+    primary_key: list[str]  # declared; empty when the catalog declares none
+    row_count: int
+
+
+@dataclass(slots=True)
+class ForeignKey:
+    name: str
+    child_table: str
+    parent_table: str
+    child_columns: list[str] = field(default_factory=list)
+    parent_columns: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class Catalog:
+    schema: str
+    tables: list[TableInfo]
+    foreign_keys: dict[str, ForeignKey]
+
+
+def _execute(cursor: Any, sql: str, param: Any) -> list[dict[str, Any]]:
+    cursor.execute(sql, (param,))
+    return fetch_dicts(cursor)
+
+
+def table_stem(table: str) -> str:
+    """A table's name without a leading warehouse prefix (``dimcustomer`` -> ``customer``,
+    ``fact_sales`` -> ``sales``); ``dim``/``fact`` elsewhere in a name are left alone."""
+    low = table.lower()
+    for prefix in ("dim", "fact"):
+        if low.startswith(prefix) and len(low) > len(prefix):
+            return low[len(prefix) :].removeprefix("_")
+    return low
+
+
+_WORDS = re.compile(r"[A-Z]+(?![a-z])|[A-Z]?[a-z]+|[0-9]+")
+
+
+def name_tokens(name: str) -> list[str]:
+    """A column name split into lower-case words, at ``_`` and at CamelCase boundaries
+    (``customer_id`` and ``CustomerId`` give ``customer, id``; ``customerID`` too;
+    ``paid`` is one word)."""
+    return [w.lower() for w in _WORDS.findall(name)]
+
+
+def key_stem(name: str) -> str | None:
+    """The words before a trailing ``id`` or ``key`` word (``customer_id`` and ``CustomerKey``
+    give ``customer``, ``sales_region_id`` gives ``sales_region``); ``None`` when the name does
+    not end in a whole ``id``/``key`` word or has nothing before it. ``paid``, ``valid`` and
+    ``monkey`` do not end in one."""
+    words = name_tokens(name)
+    if len(words) < 2 or words[-1] not in ("id", "key"):
+        return None
+    return "_".join(words[:-1])
+
+
+def _squash(name: str) -> str:
+    return "".join(name_tokens(name))
+
+
+def _violates_key(values: list[Any] | None) -> bool:
+    """A sample that shows a null or a repeated value rules a column out as a key. No rows
+    (``None`` or empty) rule nothing out."""
+    if not values:
+        return False
+    present = [v for v in values if v is not None]
+    return len(present) != len(values) or len(set(present)) != len(present)
+
+
+def guess_primary_key(
+    table: str,
+    columns: Iterable[ColumnInfo],
+    sample: dict[str, list[Any]] | None = None,
+    foreign_key_columns: Collection[str] = (),
+) -> list[str]:
+    """A primary key for a table whose catalog declares none (a Fabric warehouse enforces no
+    keys). Only a column that could be the key qualifies: it is not a foreign key (declared,
+    shown by the data or suggested by its name) and the sampled rows show no null or repeated
+    value in it (with no sampled rows nothing is ruled out). Of those, the identity column wins,
+    then a column named ``id`` or like the table: ``<table>_id``, ``<table>_key``, the same with
+    the table's singular (``orders`` gives ``order_id``). Otherwise the table is reported with no
+    primary key: guessing any other column would only look authoritative."""
+    # The name itself counts as well as its stem: ``factory`` is a table of factories, not a
+    # ``fact`` table of ``ory`` (and ``dimension`` is not a ``dim``ension of ``ension``).
+    names = {table_stem(table).replace("_", ""), table.lower().replace("_", "")}
+    stems = {n for name in names for n in (name, name.removesuffix("s"))} - {""}
+    wanted = {"id"} | {s + suffix for s in stems for suffix in ("id", "key")}
+    blocked = set(foreign_key_columns)
+    eligible = [
+        c
+        for c in columns
+        if c.name not in blocked and not _violates_key((sample or {}).get(c.name))
+    ]
+    for c in eligible:
+        if c.is_identity:
+            return [c.name]
+    for c in eligible:
+        if _squash(c.name) in wanted:
+            return [c.name]
+    return []
+
+
+def read_catalog(cursor: Any, schema: str, tables: Iterable[str] | None = None) -> Catalog:
+    """Walk ``schema``: its tables (optionally only ``tables``), their columns, declared primary
+    keys (empty when the catalog has none: :func:`guess_primary_key` judges that from data),
+    declared foreign keys and row counts."""
+    keep = set(tables) if tables else None
+    table_rows = [
+        r for r in _execute(cursor, TABLES_QUERY, schema) if keep is None or r["table_name"] in keep
+    ]
+    counts = {
+        r["table_name"]: int(r["row_count"] or 0) for r in _execute(cursor, ROW_COUNT_QUERY, schema)
+    }
+
+    fks: dict[str, ForeignKey] = {}
+    for r in _execute(cursor, FOREIGN_KEY_QUERY, schema):
+        fk = fks.setdefault(
+            r["fk_name"], ForeignKey(r["fk_name"], r["child_table"], r["parent_table"])
+        )
+        fk.child_columns.append(r["child_column"])
+        fk.parent_columns.append(r["parent_column"])
+
+    infos: list[TableInfo] = []
+    for t in table_rows:
+        object_id = t["object_id"]
+        columns = [
+            ColumnInfo(
+                name=c["column_name"],
+                # an alias type is read, typed and profiled as its base system type
+                type_name=c.get("base_type_name") or c["type_name"],
+                max_length=int(c["max_length"] or 0),
+                precision=int(c["precision"] or 0),
+                scale=int(c["scale"] or 0),
+                is_nullable=bool(c["is_nullable"]),
+                is_identity=bool(c["is_identity"]),
+            )
+            for c in _execute(cursor, COLUMNS_QUERY, object_id)
+        ]
+        pk = [r["column_name"] for r in _execute(cursor, PRIMARY_KEY_QUERY, object_id)]
+        infos.append(
+            TableInfo(
+                schema, t["table_name"], object_id, columns, pk, counts.get(t["table_name"], 0)
+            )
+        )
+    return Catalog(schema, infos, fks)

shape_sqlserver/command.py

@@ -0,0 +1,118 @@
+"""``shape profile-db``: profile a SQL Server, Azure SQL or Fabric SQL database."""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any
+
+from .auth import AUTH_METHODS, Credentials
+from .sql import (
+    DEFAULT_SAMPLE_ROWS,
+    DEFAULT_SCHEMA,
+    SqlServerError,
+    build_connection_string,
+    redact_connection_string,
+)
+
+ENV_CONNECTION_STRING = "SHAPE_SQLSERVER_CONNECTION_STRING"
+ENV_CLIENT_SECRET = "SHAPE_SQLSERVER_CLIENT_SECRET"
+
+
+class ProfileDbCommand:
+    """``shape profile-db``: profile one schema of a database into a ``.shape`` artifact."""
+
+    name = "profile-db"
+    help = "profile a SQL Server / Azure SQL / Fabric SQL database schema"
+
+    def configure(self, parser: Any) -> None:
+        parser.add_argument(
+            "--connection-string",
+            metavar="ODBC",
+            help=f"ODBC connection string (or set {ENV_CONNECTION_STRING}; keeps secrets out of "
+            "shell history)",
+        )
+        parser.add_argument("--server", help="server name, instead of --connection-string")
+        parser.add_argument("--database", help="database name, with --server")
+        parser.add_argument(
+            "--auth",
+            choices=AUTH_METHODS,
+            default="cli",
+            help="sql: login in the connection string; cli, msi, spn or fabric: Microsoft Entra "
+            "(default: cli)",
+        )
+        parser.add_argument("--tenant-id")
+        parser.add_argument("--client-id")
+        parser.add_argument(
+            "--client-secret",
+            help=f"service principal secret (prefer the {ENV_CLIENT_SECRET} variable)",
+        )
+        parser.add_argument("--schema", default=DEFAULT_SCHEMA, help="schema to profile")
+        parser.add_argument(
+            "--tables", metavar="A,B", help="only these tables (default: every table in the schema)"
+        )
+        parser.add_argument(
+            "--sample-rows",
+            type=int,
+            default=DEFAULT_SAMPLE_ROWS,
+            help="rows sampled per table; 0 profiles the catalog only "
+            f"(default: {DEFAULT_SAMPLE_ROWS})",
+        )
+        parser.add_argument("-o", "--output", metavar="OUT.shape", help="write the profile here")
+        parser.add_argument("--json", metavar="SUMMARY.json", help="also write a JSON summary")
+
+    def _connection_string(self, args: Any) -> str:
+        text = args.connection_string or os.environ.get(ENV_CONNECTION_STRING)
+        if text:
+            return str(text)
+        if args.server:
+            return build_connection_string(args.server, args.database)
+        raise SqlServerError(
+            f"give --connection-string (or {ENV_CONNECTION_STRING}) or --server and --database"
+        )
+
+    def run(self, args: Any) -> int:
+        try:
+            return self._run(args)
+        except (SqlServerError, ValueError) as exc:
+            print(f"shape: error: {redact_connection_string(str(exc))}", file=sys.stderr)
+            return 2
+
+    def _run(self, args: Any) -> int:
+        if args.sample_rows < 0:
+            raise SqlServerError("--sample-rows cannot be negative")
+        if not args.output and not args.json:
+            raise SqlServerError("profile-db needs -o OUT.shape (and/or --json SUMMARY.json)")
+        import shape
+
+        from .profiler import profile_database
+
+        creds = Credentials(
+            method=args.auth,
+            tenant_id=args.tenant_id,
+            client_id=args.client_id,
+            client_secret=args.client_secret or os.environ.get(ENV_CLIENT_SECRET),
+        )
+        tables = [t.strip() for t in args.tables.split(",") if t.strip()] if args.tables else None
+        prof = profile_database(
+            self._connection_string(args),
+            credentials=creds,
+            schema=args.schema,
+            sample_rows=args.sample_rows,
+            tables=tables,
+            name=args.database or args.schema,
+        )
+        out: dict[str, Any] = {}
+        if args.output:
+            out["written"] = args.output
+            out["shape_content_id"] = shape.save(prof, args.output)
+        if args.json:
+            with open(args.json, "w", encoding="utf-8", newline="\n") as fh:
+                json.dump(prof.summary(), fh, indent=2, default=str)
+                fh.write("\n")
+        summary = prof.summary()
+        out["tables"] = len(summary["tables"])
+        out["relationships"] = len(summary["relationships"])
+        print(json.dumps(out))
+        return 0