sqllocks-shape-sqlserver 0.9.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. sqllocks_shape_sqlserver-0.9.1/LICENSE +21 -0
  2. sqllocks_shape_sqlserver-0.9.1/PKG-INFO +51 -0
  3. sqllocks_shape_sqlserver-0.9.1/README.md +25 -0
  4. sqllocks_shape_sqlserver-0.9.1/pyproject.toml +43 -0
  5. sqllocks_shape_sqlserver-0.9.1/setup.cfg +4 -0
  6. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/__init__.py +28 -0
  7. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/auth.py +130 -0
  8. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/catalog.py +181 -0
  9. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/command.py +118 -0
  10. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/profiler.py +454 -0
  11. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/source.py +195 -0
  12. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/sql.py +320 -0
  13. sqllocks_shape_sqlserver-0.9.1/src/shape_sqlserver/testing.py +786 -0
  14. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/PKG-INFO +51 -0
  15. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/SOURCES.txt +25 -0
  16. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/dependency_links.txt +1 -0
  17. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/entry_points.txt +5 -0
  18. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/requires.txt +5 -0
  19. sqllocks_shape_sqlserver-0.9.1/src/sqllocks_shape_sqlserver.egg-info/top_level.txt +1 -0
  20. sqllocks_shape_sqlserver-0.9.1/tests/test_auth.py +193 -0
  21. sqllocks_shape_sqlserver-0.9.1/tests/test_emulator.py +345 -0
  22. sqllocks_shape_sqlserver-0.9.1/tests/test_portability.py +39 -0
  23. sqllocks_shape_sqlserver-0.9.1/tests/test_profiler.py +824 -0
  24. sqllocks_shape_sqlserver-0.9.1/tests/test_security_review.py +27 -0
  25. sqllocks_shape_sqlserver-0.9.1/tests/test_skeleton.py +22 -0
  26. sqllocks_shape_sqlserver-0.9.1/tests/test_source.py +327 -0
  27. sqllocks_shape_sqlserver-0.9.1/tests/test_sql.py +111 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025-2026 SQLLocks (Jonathan Stewart)
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,51 @@
1
+ Metadata-Version: 2.4
2
+ Name: sqllocks-shape-sqlserver
3
+ Version: 0.9.1
4
+ Summary: Shape plugin: SQL Server, Azure SQL and Fabric SQL database profiling and source
5
+ License-Expression: MIT
6
+ Project-URL: Homepage, https://github.com/sqllocks/shape
7
+ Project-URL: Repository, https://github.com/sqllocks/shape
8
+ Project-URL: Issues, https://github.com/sqllocks/shape/issues
9
+ Project-URL: Documentation, https://github.com/sqllocks/shape/tree/main/plugins/shape-sqlserver
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3 :: Only
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Programming Language :: Python :: 3.13
15
+ Classifier: Programming Language :: Python :: 3.14
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Topic :: Database
18
+ Requires-Python: >=3.11
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: sqllocks-shape==0.9.1
22
+ Requires-Dist: pyodbc>=5.0
23
+ Provides-Extra: entra
24
+ Requires-Dist: azure-identity>=1.16.1; extra == "entra"
25
+ Dynamic: license-file
26
+
27
+ # sqllocks-shape-sqlserver
28
+
29
+ Shape plugin for SQL Server, Azure SQL and Fabric SQL (warehouse and SQL database).
30
+
31
+ - **`shape profile-db`**: profile a database schema into a `.shape` profile. The catalog views
32
+ give exact primary keys, foreign keys, types and row counts; a sample of rows per table
33
+ (1000 by default) gives the distributions.
34
+ - **`mssql://` source**: read one table as Arrow record batches (`shape.sources`).
35
+ - **Writing** to a live SQL Server (the `sqlserver` sink, `mssql://` URIs, bulk insert) is in
36
+ `shape-fabric`, which builds on this plugin: see its README.
37
+ - **Shared helpers** for other SQL Server code (such as `shape-fabric`): `shape_sqlserver.sql`
38
+ (identifier quoting, connection strings and their redaction, catalog queries, type maps) and
39
+ `shape_sqlserver.auth` (SQL logins and Microsoft Entra tokens).
40
+
41
+ ```bash
42
+ pip install 'sqllocks-shape[sqlserver]' # also: sqllocks-shape-sqlserver[entra]
43
+ shape profile-db --server myserver.database.windows.net --database shop --auth cli \
44
+ --schema dbo -o shop.shape
45
+ ```
46
+
47
+ You also need the Microsoft ODBC Driver 18 for SQL Server (and unixODBC on Linux).
48
+ The full guide is `docs/plugins/sqlserver.md` in the repository.
49
+
50
+ Its version always equals core's (`sqllocks-shape`), and it is released together with core.
51
+ How plugins are written: `docs/plugins/authoring.md`.
@@ -0,0 +1,25 @@
1
+ # sqllocks-shape-sqlserver
2
+
3
+ Shape plugin for SQL Server, Azure SQL and Fabric SQL (warehouse and SQL database).
4
+
5
+ - **`shape profile-db`**: profile a database schema into a `.shape` profile. The catalog views
6
+ give exact primary keys, foreign keys, types and row counts; a sample of rows per table
7
+ (1000 by default) gives the distributions.
8
+ - **`mssql://` source**: read one table as Arrow record batches (`shape.sources`).
9
+ - **Writing** to a live SQL Server (the `sqlserver` sink, `mssql://` URIs, bulk insert) is in
10
+ `shape-fabric`, which builds on this plugin: see its README.
11
+ - **Shared helpers** for other SQL Server code (such as `shape-fabric`): `shape_sqlserver.sql`
12
+ (identifier quoting, connection strings and their redaction, catalog queries, type maps) and
13
+ `shape_sqlserver.auth` (SQL logins and Microsoft Entra tokens).
14
+
15
+ ```bash
16
+ pip install 'sqllocks-shape[sqlserver]' # also: sqllocks-shape-sqlserver[entra]
17
+ shape profile-db --server myserver.database.windows.net --database shop --auth cli \
18
+ --schema dbo -o shop.shape
19
+ ```
20
+
21
+ You also need the Microsoft ODBC Driver 18 for SQL Server (and unixODBC on Linux).
22
+ The full guide is `docs/plugins/sqlserver.md` in the repository.
23
+
24
+ Its version always equals core's (`sqllocks-shape`), and it is released together with core.
25
+ How plugins are written: `docs/plugins/authoring.md`.
@@ -0,0 +1,43 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "sqllocks-shape-sqlserver"
7
+ # Lockstep with core (T-09): scripts/check_plugin_skeletons.py fails when this drifts.
8
+ version = "0.9.1"
9
+ description = "Shape plugin: SQL Server, Azure SQL and Fabric SQL database profiling and source"
10
+ readme = "README.md"
11
+ requires-python = ">=3.11"
12
+ license = "MIT"
13
+ license-files = ["LICENSE"]
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "Programming Language :: Python :: 3 :: Only",
17
+ "Programming Language :: Python :: 3.11",
18
+ "Programming Language :: Python :: 3.12",
19
+ "Programming Language :: Python :: 3.13",
20
+ "Programming Language :: Python :: 3.14",
21
+ "Operating System :: OS Independent",
22
+ "Topic :: Database",
23
+ ]
24
+ dependencies = ["sqllocks-shape==0.9.1", "pyodbc>=5.0"]
25
+
26
+ [project.urls]
27
+ Homepage = "https://github.com/sqllocks/shape"
28
+ Repository = "https://github.com/sqllocks/shape"
29
+ Issues = "https://github.com/sqllocks/shape/issues"
30
+ Documentation = "https://github.com/sqllocks/shape/tree/main/plugins/shape-sqlserver"
31
+
32
+ [project.optional-dependencies]
33
+ # Microsoft Entra authentication (--auth cli | msi | spn).
34
+ entra = ["azure-identity>=1.16.1"]
35
+
36
+ [project.entry-points."shape.sources"]
37
+ mssql = "shape_sqlserver:SqlServerSource"
38
+
39
+ [project.entry-points."shape.commands"]
40
+ profile-db = "shape_sqlserver:ProfileDbCommand"
41
+
42
+ [tool.setuptools.packages.find]
43
+ where = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,28 @@
1
+ """Shape plugin: SQL Server, Azure SQL and Fabric SQL.
2
+
3
+ * ``shape.sources``: ``mssql`` reads one table as Arrow record batches.
4
+ * ``shape.commands``: ``shape profile-db`` profiles a database schema.
5
+ * helpers shared with ``shape-fabric``: :mod:`shape_sqlserver.sql` (identifier quoting,
6
+ connection strings, catalog queries, type maps) and :mod:`shape_sqlserver.auth`
7
+ (connections, Microsoft Entra tokens).
8
+
9
+ ``pyodbc`` and ``azure-identity`` load only when a connection is opened.
10
+ """
11
+
12
+ from .auth import Credentials, connect
13
+ from .command import ProfileDbCommand
14
+ from .profiler import profile_database
15
+ from .source import SqlServerSource
16
+ from .sql import SqlServerError
17
+
18
+ SHAPE_API = "1.0"
19
+
20
+ __all__ = [
21
+ "SHAPE_API",
22
+ "Credentials",
23
+ "ProfileDbCommand",
24
+ "SqlServerError",
25
+ "SqlServerSource",
26
+ "connect",
27
+ "profile_database",
28
+ ]
@@ -0,0 +1,130 @@
1
+ """Connections to SQL Server, Azure SQL and Fabric SQL: SQL logins and Microsoft Entra tokens.
2
+
3
+ ``pyodbc`` and ``azure-identity`` are imported only when a connection is made, so the plugin
4
+ loads (and ``shape plugins doctor`` is clean) on a machine that has neither.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import struct
10
+ from dataclasses import dataclass
11
+ from typing import Any
12
+
13
+ from .sql import SqlServerError, redact_connection_string, register_converters
14
+
15
+ AUTH_METHODS = ("sql", "cli", "msi", "spn", "fabric")
16
+ TOKEN_SCOPE = "https://database.windows.net/.default"
17
+ FABRIC_TOKEN_AUDIENCE = "https://database.windows.net/"
18
+ # ODBC connection attribute that carries an access token (SQL_COPT_SS_ACCESS_TOKEN).
19
+ ACCESS_TOKEN_ATTR = 1256
20
+
21
+
22
+ @dataclass(frozen=True, slots=True, repr=False)
23
+ class Credentials:
24
+ """How to authenticate.
25
+
26
+ ``sql``: the login is in the connection string. ``cli``: the Azure CLI's signed-in account.
27
+ ``msi``: managed identity (or the default Azure credential chain). ``spn``: a service
28
+ principal (``tenant_id``, ``client_id``, ``client_secret``). ``fabric``: the token of the
29
+ running Fabric notebook.
30
+ """
31
+
32
+ method: str = "cli"
33
+ tenant_id: str | None = None
34
+ client_id: str | None = None
35
+ client_secret: str | None = None
36
+
37
+ def __repr__(self) -> str: # never show the secret
38
+ return (
39
+ f"Credentials(method={self.method!r}, tenant_id={self.tenant_id!r}, "
40
+ f"client_id={self.client_id!r}, client_secret={'***' if self.client_secret else None})"
41
+ )
42
+
43
+ def validate(self) -> None:
44
+ if self.method not in AUTH_METHODS:
45
+ raise SqlServerError(
46
+ f"unsupported auth method {self.method!r}; use one of {', '.join(AUTH_METHODS)}"
47
+ )
48
+ if self.method == "spn" and not (self.tenant_id and self.client_id and self.client_secret):
49
+ raise SqlServerError("auth 'spn' needs a tenant id, a client id and a client secret")
50
+
51
+
52
+ def _credential(creds: Credentials) -> Any:
53
+ try:
54
+ from azure import identity # type: ignore[import-not-found,unused-ignore]
55
+ except ImportError as exc:
56
+ raise SqlServerError(
57
+ "Entra authentication needs azure-identity: "
58
+ "pip install 'sqllocks-shape-sqlserver[entra]'"
59
+ ) from exc
60
+ if creds.method == "cli":
61
+ return identity.AzureCliCredential()
62
+ if creds.method == "msi":
63
+ return identity.DefaultAzureCredential(exclude_managed_identity_credential=False)
64
+ return identity.ClientSecretCredential(
65
+ tenant_id=creds.tenant_id, client_id=creds.client_id, client_secret=creds.client_secret
66
+ )
67
+
68
+
69
+ def _fabric_token() -> str:
70
+ try:
71
+ try:
72
+ from notebookutils import mssparkutils # type: ignore[import-not-found,unused-ignore]
73
+ except ImportError:
74
+ import mssparkutils # type: ignore[import-not-found,no-redef,unused-ignore]
75
+ token: str = mssparkutils.credentials.getToken(FABRIC_TOKEN_AUDIENCE)
76
+ except ImportError as exc:
77
+ raise SqlServerError("auth 'fabric' only works inside a Fabric notebook") from exc
78
+ return token
79
+
80
+
81
+ def access_token(creds: Credentials) -> bytes:
82
+ """The Entra access token for ``creds``, UTF-16-LE encoded as the ODBC driver expects."""
83
+ creds.validate()
84
+ if creds.method == "sql":
85
+ raise SqlServerError("auth 'sql' uses the login in the connection string, not a token")
86
+ if creds.method == "fabric":
87
+ return _fabric_token().encode("utf-16-le")
88
+ token: str = _credential(creds).get_token(TOKEN_SCOPE).token
89
+ return token.encode("utf-16-le")
90
+
91
+
92
+ def token_struct(token: bytes) -> bytes:
93
+ """``token`` packed as the ODBC ``SQL_COPT_SS_ACCESS_TOKEN`` structure."""
94
+ return struct.pack(f"<I{len(token)}s", len(token), token)
95
+
96
+
97
+ def connect(connection_string: str, creds: Credentials | None = None, *, timeout: int = 30) -> Any:
98
+ """An open pyodbc connection.
99
+
100
+ The connection is in autocommit mode: Shape only reads, and an implicit transaction left
101
+ open on a pooled connection would block other sessions.
102
+
103
+ With ``sql`` the connection string carries the login. With any other method an Entra
104
+ token is fetched and passed to the driver; the connection string must then not hold a
105
+ ``UID``, ``PWD`` or ``Authentication`` key.
106
+ """
107
+ creds = creds or Credentials()
108
+ creds.validate()
109
+ try:
110
+ import pyodbc # type: ignore[import-not-found,unused-ignore]
111
+ except ImportError as exc:
112
+ raise SqlServerError(
113
+ "pyodbc is not available. Install it and the Microsoft ODBC Driver 18 for SQL Server "
114
+ "(on Linux also unixODBC)."
115
+ ) from exc
116
+ try:
117
+ if creds.method == "sql":
118
+ conn = pyodbc.connect(connection_string, timeout=timeout, autocommit=True)
119
+ else:
120
+ token = token_struct(access_token(creds))
121
+ conn = pyodbc.connect(
122
+ connection_string,
123
+ attrs_before={ACCESS_TOKEN_ATTR: token},
124
+ timeout=timeout,
125
+ autocommit=True,
126
+ )
127
+ except pyodbc.Error as exc:
128
+ raise SqlServerError(f"could not connect: {redact_connection_string(str(exc))}") from None
129
+ register_converters(conn)
130
+ return conn
@@ -0,0 +1,181 @@
1
+ """Schema walk: tables, columns, keys and row counts from the SQL Server catalog views."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import re
6
+ from collections.abc import Collection, Iterable
7
+ from dataclasses import dataclass, field
8
+ from typing import Any
9
+
10
+ from .sql import (
11
+ COLUMNS_QUERY,
12
+ FOREIGN_KEY_QUERY,
13
+ PRIMARY_KEY_QUERY,
14
+ ROW_COUNT_QUERY,
15
+ TABLES_QUERY,
16
+ fetch_dicts,
17
+ )
18
+
19
+
20
+ @dataclass(frozen=True, slots=True)
21
+ class ColumnInfo:
22
+ name: str
23
+ type_name: str
24
+ max_length: int
25
+ precision: int
26
+ scale: int
27
+ is_nullable: bool
28
+ is_identity: bool
29
+
30
+
31
+ @dataclass(slots=True)
32
+ class TableInfo:
33
+ schema: str
34
+ name: str
35
+ object_id: int
36
+ columns: list[ColumnInfo]
37
+ primary_key: list[str] # declared; empty when the catalog declares none
38
+ row_count: int
39
+
40
+
41
+ @dataclass(slots=True)
42
+ class ForeignKey:
43
+ name: str
44
+ child_table: str
45
+ parent_table: str
46
+ child_columns: list[str] = field(default_factory=list)
47
+ parent_columns: list[str] = field(default_factory=list)
48
+
49
+
50
+ @dataclass(slots=True)
51
+ class Catalog:
52
+ schema: str
53
+ tables: list[TableInfo]
54
+ foreign_keys: dict[str, ForeignKey]
55
+
56
+
57
+ def _execute(cursor: Any, sql: str, param: Any) -> list[dict[str, Any]]:
58
+ cursor.execute(sql, (param,))
59
+ return fetch_dicts(cursor)
60
+
61
+
62
+ def table_stem(table: str) -> str:
63
+ """A table's name without a leading warehouse prefix (``dimcustomer`` -> ``customer``,
64
+ ``fact_sales`` -> ``sales``); ``dim``/``fact`` elsewhere in a name are left alone."""
65
+ low = table.lower()
66
+ for prefix in ("dim", "fact"):
67
+ if low.startswith(prefix) and len(low) > len(prefix):
68
+ return low[len(prefix) :].removeprefix("_")
69
+ return low
70
+
71
+
72
+ _WORDS = re.compile(r"[A-Z]+(?![a-z])|[A-Z]?[a-z]+|[0-9]+")
73
+
74
+
75
+ def name_tokens(name: str) -> list[str]:
76
+ """A column name split into lower-case words, at ``_`` and at CamelCase boundaries
77
+ (``customer_id`` and ``CustomerId`` give ``customer, id``; ``customerID`` too;
78
+ ``paid`` is one word)."""
79
+ return [w.lower() for w in _WORDS.findall(name)]
80
+
81
+
82
+ def key_stem(name: str) -> str | None:
83
+ """The words before a trailing ``id`` or ``key`` word (``customer_id`` and ``CustomerKey``
84
+ give ``customer``, ``sales_region_id`` gives ``sales_region``); ``None`` when the name does
85
+ not end in a whole ``id``/``key`` word or has nothing before it. ``paid``, ``valid`` and
86
+ ``monkey`` do not end in one."""
87
+ words = name_tokens(name)
88
+ if len(words) < 2 or words[-1] not in ("id", "key"):
89
+ return None
90
+ return "_".join(words[:-1])
91
+
92
+
93
+ def _squash(name: str) -> str:
94
+ return "".join(name_tokens(name))
95
+
96
+
97
+ def _violates_key(values: list[Any] | None) -> bool:
98
+ """A sample that shows a null or a repeated value rules a column out as a key. No rows
99
+ (``None`` or empty) rule nothing out."""
100
+ if not values:
101
+ return False
102
+ present = [v for v in values if v is not None]
103
+ return len(present) != len(values) or len(set(present)) != len(present)
104
+
105
+
106
+ def guess_primary_key(
107
+ table: str,
108
+ columns: Iterable[ColumnInfo],
109
+ sample: dict[str, list[Any]] | None = None,
110
+ foreign_key_columns: Collection[str] = (),
111
+ ) -> list[str]:
112
+ """A primary key for a table whose catalog declares none (a Fabric warehouse enforces no
113
+ keys). Only a column that could be the key qualifies: it is not a foreign key (declared,
114
+ shown by the data or suggested by its name) and the sampled rows show no null or repeated
115
+ value in it (with no sampled rows nothing is ruled out). Of those, the identity column wins,
116
+ then a column named ``id`` or like the table: ``<table>_id``, ``<table>_key``, the same with
117
+ the table's singular (``orders`` gives ``order_id``). Otherwise the table is reported with no
118
+ primary key: guessing any other column would only look authoritative."""
119
+ # The name itself counts as well as its stem: ``factory`` is a table of factories, not a
120
+ # ``fact`` table of ``ory`` (and ``dimension`` is not a ``dim``ension of ``ension``).
121
+ names = {table_stem(table).replace("_", ""), table.lower().replace("_", "")}
122
+ stems = {n for name in names for n in (name, name.removesuffix("s"))} - {""}
123
+ wanted = {"id"} | {s + suffix for s in stems for suffix in ("id", "key")}
124
+ blocked = set(foreign_key_columns)
125
+ eligible = [
126
+ c
127
+ for c in columns
128
+ if c.name not in blocked and not _violates_key((sample or {}).get(c.name))
129
+ ]
130
+ for c in eligible:
131
+ if c.is_identity:
132
+ return [c.name]
133
+ for c in eligible:
134
+ if _squash(c.name) in wanted:
135
+ return [c.name]
136
+ return []
137
+
138
+
139
+ def read_catalog(cursor: Any, schema: str, tables: Iterable[str] | None = None) -> Catalog:
140
+ """Walk ``schema``: its tables (optionally only ``tables``), their columns, declared primary
141
+ keys (empty when the catalog has none: :func:`guess_primary_key` judges that from data),
142
+ declared foreign keys and row counts."""
143
+ keep = set(tables) if tables else None
144
+ table_rows = [
145
+ r for r in _execute(cursor, TABLES_QUERY, schema) if keep is None or r["table_name"] in keep
146
+ ]
147
+ counts = {
148
+ r["table_name"]: int(r["row_count"] or 0) for r in _execute(cursor, ROW_COUNT_QUERY, schema)
149
+ }
150
+
151
+ fks: dict[str, ForeignKey] = {}
152
+ for r in _execute(cursor, FOREIGN_KEY_QUERY, schema):
153
+ fk = fks.setdefault(
154
+ r["fk_name"], ForeignKey(r["fk_name"], r["child_table"], r["parent_table"])
155
+ )
156
+ fk.child_columns.append(r["child_column"])
157
+ fk.parent_columns.append(r["parent_column"])
158
+
159
+ infos: list[TableInfo] = []
160
+ for t in table_rows:
161
+ object_id = t["object_id"]
162
+ columns = [
163
+ ColumnInfo(
164
+ name=c["column_name"],
165
+ # an alias type is read, typed and profiled as its base system type
166
+ type_name=c.get("base_type_name") or c["type_name"],
167
+ max_length=int(c["max_length"] or 0),
168
+ precision=int(c["precision"] or 0),
169
+ scale=int(c["scale"] or 0),
170
+ is_nullable=bool(c["is_nullable"]),
171
+ is_identity=bool(c["is_identity"]),
172
+ )
173
+ for c in _execute(cursor, COLUMNS_QUERY, object_id)
174
+ ]
175
+ pk = [r["column_name"] for r in _execute(cursor, PRIMARY_KEY_QUERY, object_id)]
176
+ infos.append(
177
+ TableInfo(
178
+ schema, t["table_name"], object_id, columns, pk, counts.get(t["table_name"], 0)
179
+ )
180
+ )
181
+ return Catalog(schema, infos, fks)
@@ -0,0 +1,118 @@
1
+ """``shape profile-db``: profile a SQL Server, Azure SQL or Fabric SQL database."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import os
7
+ import sys
8
+ from typing import Any
9
+
10
+ from .auth import AUTH_METHODS, Credentials
11
+ from .sql import (
12
+ DEFAULT_SAMPLE_ROWS,
13
+ DEFAULT_SCHEMA,
14
+ SqlServerError,
15
+ build_connection_string,
16
+ redact_connection_string,
17
+ )
18
+
19
+ ENV_CONNECTION_STRING = "SHAPE_SQLSERVER_CONNECTION_STRING"
20
+ ENV_CLIENT_SECRET = "SHAPE_SQLSERVER_CLIENT_SECRET"
21
+
22
+
23
+ class ProfileDbCommand:
24
+ """``shape profile-db``: profile one schema of a database into a ``.shape`` artifact."""
25
+
26
+ name = "profile-db"
27
+ help = "profile a SQL Server / Azure SQL / Fabric SQL database schema"
28
+
29
+ def configure(self, parser: Any) -> None:
30
+ parser.add_argument(
31
+ "--connection-string",
32
+ metavar="ODBC",
33
+ help=f"ODBC connection string (or set {ENV_CONNECTION_STRING}; keeps secrets out of "
34
+ "shell history)",
35
+ )
36
+ parser.add_argument("--server", help="server name, instead of --connection-string")
37
+ parser.add_argument("--database", help="database name, with --server")
38
+ parser.add_argument(
39
+ "--auth",
40
+ choices=AUTH_METHODS,
41
+ default="cli",
42
+ help="sql: login in the connection string; cli, msi, spn or fabric: Microsoft Entra "
43
+ "(default: cli)",
44
+ )
45
+ parser.add_argument("--tenant-id")
46
+ parser.add_argument("--client-id")
47
+ parser.add_argument(
48
+ "--client-secret",
49
+ help=f"service principal secret (prefer the {ENV_CLIENT_SECRET} variable)",
50
+ )
51
+ parser.add_argument("--schema", default=DEFAULT_SCHEMA, help="schema to profile")
52
+ parser.add_argument(
53
+ "--tables", metavar="A,B", help="only these tables (default: every table in the schema)"
54
+ )
55
+ parser.add_argument(
56
+ "--sample-rows",
57
+ type=int,
58
+ default=DEFAULT_SAMPLE_ROWS,
59
+ help="rows sampled per table; 0 profiles the catalog only "
60
+ f"(default: {DEFAULT_SAMPLE_ROWS})",
61
+ )
62
+ parser.add_argument("-o", "--output", metavar="OUT.shape", help="write the profile here")
63
+ parser.add_argument("--json", metavar="SUMMARY.json", help="also write a JSON summary")
64
+
65
+ def _connection_string(self, args: Any) -> str:
66
+ text = args.connection_string or os.environ.get(ENV_CONNECTION_STRING)
67
+ if text:
68
+ return str(text)
69
+ if args.server:
70
+ return build_connection_string(args.server, args.database)
71
+ raise SqlServerError(
72
+ f"give --connection-string (or {ENV_CONNECTION_STRING}) or --server and --database"
73
+ )
74
+
75
+ def run(self, args: Any) -> int:
76
+ try:
77
+ return self._run(args)
78
+ except (SqlServerError, ValueError) as exc:
79
+ print(f"shape: error: {redact_connection_string(str(exc))}", file=sys.stderr)
80
+ return 2
81
+
82
+ def _run(self, args: Any) -> int:
83
+ if args.sample_rows < 0:
84
+ raise SqlServerError("--sample-rows cannot be negative")
85
+ if not args.output and not args.json:
86
+ raise SqlServerError("profile-db needs -o OUT.shape (and/or --json SUMMARY.json)")
87
+ import shape
88
+
89
+ from .profiler import profile_database
90
+
91
+ creds = Credentials(
92
+ method=args.auth,
93
+ tenant_id=args.tenant_id,
94
+ client_id=args.client_id,
95
+ client_secret=args.client_secret or os.environ.get(ENV_CLIENT_SECRET),
96
+ )
97
+ tables = [t.strip() for t in args.tables.split(",") if t.strip()] if args.tables else None
98
+ prof = profile_database(
99
+ self._connection_string(args),
100
+ credentials=creds,
101
+ schema=args.schema,
102
+ sample_rows=args.sample_rows,
103
+ tables=tables,
104
+ name=args.database or args.schema,
105
+ )
106
+ out: dict[str, Any] = {}
107
+ if args.output:
108
+ out["written"] = args.output
109
+ out["shape_content_id"] = shape.save(prof, args.output)
110
+ if args.json:
111
+ with open(args.json, "w", encoding="utf-8", newline="\n") as fh:
112
+ json.dump(prof.summary(), fh, indent=2, default=str)
113
+ fh.write("\n")
114
+ summary = prof.summary()
115
+ out["tables"] = len(summary["tables"])
116
+ out["relationships"] = len(summary["relationships"])
117
+ print(json.dumps(out))
118
+ return 0