databricks-schema 0.5.1__tar.gz → 0.6.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 (30) hide show
  1. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/.gitignore +4 -1
  2. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/PKG-INFO +20 -2
  3. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/README.md +19 -1
  4. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/__init__.py +6 -1
  5. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/cli.py +135 -3
  6. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/diff.py +48 -0
  7. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/extractor.py +5 -1
  8. databricks_schema-0.6.1/databricks_schema/validate.py +136 -0
  9. databricks_schema-0.6.1/examples/01_extract_schema.py +19 -0
  10. databricks_schema-0.6.1/examples/02_save_load_yaml.py +30 -0
  11. databricks_schema-0.6.1/examples/03_diff_schemas.py +44 -0
  12. databricks_schema-0.6.1/examples/04_generate_sql.py +41 -0
  13. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/pyproject.toml +1 -1
  14. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/test_diff.py +82 -0
  15. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/test_extractor.py +27 -0
  16. databricks_schema-0.6.1/tests/test_validate.py +146 -0
  17. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/uv.lock +1 -1
  18. databricks_schema-0.5.1/examples/full_example.py +0 -49
  19. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/.github/workflows/ci.yml +0 -0
  20. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/.github/workflows/release.yml +0 -0
  21. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/.python-version +0 -0
  22. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/CLAUDE.md +0 -0
  23. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/LICENSE +0 -0
  24. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/models.py +0 -0
  25. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/sql_gen.py +0 -0
  26. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/databricks_schema/yaml_io.py +0 -0
  27. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/__init__.py +0 -0
  28. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/test_models.py +0 -0
  29. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/test_sql_gen.py +0 -0
  30. {databricks_schema-0.5.1 → databricks_schema-0.6.1}/tests/test_yaml_io.py +0 -0
@@ -10,4 +10,7 @@ wheels/
10
10
  .venv
11
11
 
12
12
  # Tests
13
- .test_output/
13
+ .test_output/
14
+
15
+ # Claude Code
16
+ .claude/
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: databricks-schema
3
- Version: 0.5.1
3
+ Version: 0.6.1
4
4
  Summary: Databricks Unity Catalog schema extractor
5
5
  Project-URL: Repository, https://github.com/frizzleqq/databricks-schema
6
6
  Project-URL: Issues, https://github.com/frizzleqq/databricks-schema/issues
@@ -103,7 +103,7 @@ uv sync --all-groups
103
103
  The tool uses the [Databricks SDK](https://github.com/databricks/databricks-sdk-py) for auth. Configure it via environment variables:
104
104
 
105
105
  ```bash
106
- export DATABRICKS_HOST=https://<workspace>.cloud.databricks.com
106
+ export DATABRICKS_HOST=https://<databricks-url>
107
107
  export DATABRICKS_TOKEN=<your-personal-access-token>
108
108
  ```
109
109
 
@@ -235,6 +235,22 @@ Include additional metadata (`owner`, `storage_location`) in the comparison:
235
235
  databricks-schema generate-sql <catalog> ./schemas/ --include-metadata
236
236
  ```
237
237
 
238
+ ### `validate`
239
+
240
+ Validate local schema files for structural integrity (no Databricks connection needed):
241
+
242
+ ```bash
243
+ databricks-schema validate ./schemas/
244
+ ```
245
+
246
+ ### `diff-files`
247
+
248
+ Compare two local directories of schema files (no Databricks connection needed):
249
+
250
+ ```bash
251
+ databricks-schema diff-files ./schemas-prod/ ./schemas-test/
252
+ ```
253
+
238
254
  ### `list-catalogs`
239
255
 
240
256
  List all accessible catalogs:
@@ -253,6 +269,8 @@ databricks-schema list-schemas <catalog>
253
269
 
254
270
  ## Python Library Usage
255
271
 
272
+ See also the [`examples/`](examples/) directory for standalone scripts.
273
+
256
274
  ```python
257
275
  from pathlib import Path
258
276
  from databricks_schema import CatalogExtractor, catalog_to_yaml, schema_from_yaml
@@ -81,7 +81,7 @@ uv sync --all-groups
81
81
  The tool uses the [Databricks SDK](https://github.com/databricks/databricks-sdk-py) for auth. Configure it via environment variables:
82
82
 
83
83
  ```bash
84
- export DATABRICKS_HOST=https://<workspace>.cloud.databricks.com
84
+ export DATABRICKS_HOST=https://<databricks-url>
85
85
  export DATABRICKS_TOKEN=<your-personal-access-token>
86
86
  ```
87
87
 
@@ -213,6 +213,22 @@ Include additional metadata (`owner`, `storage_location`) in the comparison:
213
213
  databricks-schema generate-sql <catalog> ./schemas/ --include-metadata
214
214
  ```
215
215
 
216
+ ### `validate`
217
+
218
+ Validate local schema files for structural integrity (no Databricks connection needed):
219
+
220
+ ```bash
221
+ databricks-schema validate ./schemas/
222
+ ```
223
+
224
+ ### `diff-files`
225
+
226
+ Compare two local directories of schema files (no Databricks connection needed):
227
+
228
+ ```bash
229
+ databricks-schema diff-files ./schemas-prod/ ./schemas-test/
230
+ ```
231
+
216
232
  ### `list-catalogs`
217
233
 
218
234
  List all accessible catalogs:
@@ -231,6 +247,8 @@ databricks-schema list-schemas <catalog>
231
247
 
232
248
  ## Python Library Usage
233
249
 
250
+ See also the [`examples/`](examples/) directory for standalone scripts.
251
+
234
252
  ```python
235
253
  from pathlib import Path
236
254
  from databricks_schema import CatalogExtractor, catalog_to_yaml, schema_from_yaml
@@ -7,11 +7,13 @@ from databricks_schema.diff import (
7
7
  SchemaDiff,
8
8
  TableDiff,
9
9
  diff_catalog_with_dir,
10
+ diff_schema_dirs,
10
11
  diff_schemas,
11
12
  )
12
13
  from databricks_schema.extractor import CatalogExtractor
13
14
  from databricks_schema.models import Catalog, Column, ForeignKey, PrimaryKey, Schema, Table
14
15
  from databricks_schema.sql_gen import schema_diff_to_sql
16
+ from databricks_schema.validate import ValidationIssue, ValidationResult, validate_schemas
15
17
  from databricks_schema.yaml_io import (
16
18
  catalog_from_json,
17
19
  catalog_from_yaml,
@@ -37,16 +39,19 @@ __all__ = [
37
39
  "Table",
38
40
  "TableDiff",
39
41
  "TableType",
42
+ "ValidationIssue",
43
+ "ValidationResult",
40
44
  "catalog_from_json",
41
45
  "catalog_from_yaml",
42
46
  "catalog_to_json",
43
47
  "catalog_to_yaml",
44
48
  "diff_catalog_with_dir",
49
+ "diff_schema_dirs",
45
50
  "diff_schemas",
46
- "schema_from_json",
47
51
  "schema_diff_to_sql",
48
52
  "schema_from_json",
49
53
  "schema_from_yaml",
50
54
  "schema_to_json",
51
55
  "schema_to_yaml",
56
+ "validate_schemas",
52
57
  ]
@@ -7,11 +7,19 @@ import sys
7
7
  from pathlib import Path
8
8
 
9
9
  from databricks.sdk import WorkspaceClient
10
-
11
- from databricks_schema.diff import CatalogDiff, SchemaDiff, diff_catalog_with_dir, diff_schemas
10
+ from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied, Unauthenticated
11
+
12
+ from databricks_schema.diff import (
13
+ CatalogDiff,
14
+ SchemaDiff,
15
+ diff_catalog_with_dir,
16
+ diff_schema_dirs,
17
+ diff_schemas,
18
+ )
12
19
  from databricks_schema.extractor import CatalogExtractor
13
20
  from databricks_schema.models import Schema
14
21
  from databricks_schema.sql_gen import schema_diff_to_sql
22
+ from databricks_schema.validate import validate_schemas
15
23
  from databricks_schema.yaml_io import (
16
24
  schema_from_json,
17
25
  schema_from_yaml,
@@ -237,6 +245,83 @@ def _cmd_generate_sql(args: argparse.Namespace) -> None:
237
245
  print()
238
246
 
239
247
 
248
+ def _detect_fmt(schema_dir: Path) -> str:
249
+ """Auto-detect YAML or JSON format from files in a directory. Exits 2 on mixed or empty."""
250
+ yaml_files = list(schema_dir.glob("*.yaml"))
251
+ json_files = list(schema_dir.glob("*.json"))
252
+ if yaml_files and json_files:
253
+ print(f"Error: mixed YAML and JSON files in {schema_dir}.", file=sys.stderr)
254
+ sys.exit(2)
255
+ elif not yaml_files and not json_files:
256
+ print(f"No YAML or JSON files found in {schema_dir}.", file=sys.stderr)
257
+ sys.exit(2)
258
+ return "json" if json_files else "yaml"
259
+
260
+
261
+ def _cmd_validate(args: argparse.Namespace) -> None:
262
+ """Validate local schema files for structural integrity."""
263
+ schema_dir: Path = args.schema_dir
264
+ if not schema_dir.is_dir():
265
+ print(f"Error: {schema_dir} is not a directory.", file=sys.stderr)
266
+ sys.exit(2)
267
+
268
+ fmt = _detect_fmt(schema_dir)
269
+ ext = ".json" if fmt == "json" else ".yaml"
270
+ loader = schema_from_json if fmt == "json" else schema_from_yaml
271
+
272
+ schema_filter_set: frozenset[str] | None = frozenset(args.schema) if args.schema else None
273
+ schemas: dict[str, Schema] = {}
274
+ for schema_file in sorted(schema_dir.glob(f"*{ext}")):
275
+ if schema_filter_set is not None and schema_file.stem not in schema_filter_set:
276
+ continue
277
+ schema = loader(schema_file.read_text(encoding="utf-8"))
278
+ schemas[schema.name] = schema
279
+
280
+ if not schemas:
281
+ print("No schemas to validate.", file=sys.stderr)
282
+ sys.exit(2)
283
+
284
+ result = validate_schemas(schemas)
285
+ if not result.has_errors:
286
+ print(f"OK — {len(schemas)} schema(s) validated, no issues found.")
287
+ return
288
+
289
+ for issue in result.issues:
290
+ print(f"ERROR: {issue}")
291
+ print(f"\n{len(result.issues)} issue(s) found in {len(schemas)} schema(s).", file=sys.stderr)
292
+ sys.exit(1)
293
+
294
+
295
+ def _cmd_diff_files(args: argparse.Namespace) -> None:
296
+ """Compare two local directories of schema files."""
297
+ dir1: Path = args.dir1
298
+ dir2: Path = args.dir2
299
+ for d in (dir1, dir2):
300
+ if not d.is_dir():
301
+ print(f"Error: {d} is not a directory.", file=sys.stderr)
302
+ sys.exit(2)
303
+
304
+ fmt1 = _detect_fmt(dir1)
305
+ fmt2 = _detect_fmt(dir2)
306
+
307
+ print(f"Comparing {dir1} against {dir2}...", file=sys.stderr)
308
+ result = diff_schema_dirs(
309
+ dir1,
310
+ dir2,
311
+ fmt1=fmt1,
312
+ fmt2=fmt2,
313
+ schema_names=frozenset(args.schema) if args.schema else None,
314
+ include_metadata=args.include_metadata,
315
+ )
316
+
317
+ if not result.has_changes:
318
+ print("No differences found.")
319
+ return
320
+
321
+ _print_diff(result)
322
+ sys.exit(1)
323
+
324
+
240
325
  def _cmd_list_catalogs(args: argparse.Namespace) -> None:
241
326
  """List all accessible catalogs."""
242
327
  client = _make_client(args.host, args.token)
@@ -396,6 +481,45 @@ def _build_parser() -> argparse.ArgumentParser:
396
481
  _add_connection_args(gen_sql_p)
397
482
  gen_sql_p.set_defaults(func=_cmd_generate_sql)
398
483
 
484
+ # validate
485
+ validate_p = subparsers.add_parser(
486
+ "validate",
487
+ help="Validate local schema files for structural integrity (no Databricks connection).",
488
+ )
489
+ validate_p.add_argument(
490
+ "schema_dir", type=Path, help="Directory containing per-schema YAML or JSON files"
491
+ )
492
+ validate_p.add_argument(
493
+ "--schema",
494
+ "-s",
495
+ action="append",
496
+ metavar="SCHEMA",
497
+ help="Schema filter (repeatable)",
498
+ )
499
+ validate_p.set_defaults(func=_cmd_validate)
500
+
501
+ # diff-files
502
+ diff_files_p = subparsers.add_parser(
503
+ "diff-files",
504
+ help="Compare two local directories of schema files (no Databricks connection needed).",
505
+ )
506
+ diff_files_p.add_argument("dir1", type=Path, help="Baseline directory (old/stored state)")
507
+ diff_files_p.add_argument("dir2", type=Path, help="Current directory (new/live state)")
508
+ diff_files_p.add_argument(
509
+ "--schema",
510
+ "-s",
511
+ action="append",
512
+ metavar="SCHEMA",
513
+ help="Schema filter (repeatable)",
514
+ )
515
+ diff_files_p.add_argument(
516
+ "--include-metadata",
517
+ action="store_true",
518
+ dest="include_metadata",
519
+ help="Include owner in comparison",
520
+ )
521
+ diff_files_p.set_defaults(func=_cmd_diff_files)
522
+
399
523
  # list-catalogs
400
524
  list_catalogs_p = subparsers.add_parser("list-catalogs", help="List all accessible catalogs.")
401
525
  _add_connection_args(list_catalogs_p)
@@ -416,4 +540,12 @@ def main() -> None:
416
540
  parser.print_help()
417
541
  sys.exit(0)
418
542
  args = parser.parse_args()
419
- args.func(args)
543
+ try:
544
+ args.func(args)
545
+ except NotFound as e:
546
+ print(f"Error: {type(e).__name__}: {e}", file=sys.stderr)
547
+ print("Check that the catalog and schema names are correct.", file=sys.stderr)
548
+ sys.exit(2)
549
+ except (Unauthenticated, PermissionDenied, DatabricksError) as e:
550
+ print(f"Error: {type(e).__name__}: {e}", file=sys.stderr)
551
+ sys.exit(2)
@@ -128,6 +128,54 @@ def diff_schemas(live: Schema, stored: Schema, include_metadata: bool = False) -
128
128
  return SchemaDiff(name=live.name, status=status, changes=changes, tables=table_diffs)
129
129
 
130
130
 
131
+ def diff_schema_dirs(
132
+ dir1: Path,
133
+ dir2: Path,
134
+ fmt1: Literal["yaml", "json"] = "yaml",
135
+ fmt2: Literal["yaml", "json"] = "yaml",
136
+ schema_names: frozenset[str] | None = None,
137
+ include_metadata: bool = False,
138
+ ) -> CatalogDiff:
139
+ """Compare two local directories of schema files.
140
+
141
+ dir1 is treated as the stored/baseline state (old).
142
+ dir2 is treated as the live/current state (new).
143
+
144
+ Schemas present in dir1 but not dir2 are reported as "removed".
145
+ Schemas present in dir2 but not dir1 are reported as "added".
146
+ Schemas present in both are compared with diff_schemas.
147
+
148
+ schema_names: if set, only files whose stem is in this set are loaded from either directory.
149
+ fmt1 / fmt2: file format for dir1 and dir2 respectively ("yaml" or "json").
150
+ """
151
+
152
+ def _load_dir(path: Path, fmt: Literal["yaml", "json"]) -> dict[str, Schema]:
153
+ ext = ".json" if fmt == "json" else ".yaml"
154
+ loader = schema_from_json if fmt == "json" else schema_from_yaml
155
+ schemas: dict[str, Schema] = {}
156
+ for f in sorted(path.glob(f"*{ext}")):
157
+ if schema_names is not None and f.stem not in schema_names:
158
+ continue
159
+ schema = loader(f.read_text(encoding="utf-8"))
160
+ schemas[schema.name] = schema
161
+ return schemas
162
+
163
+ stored = _load_dir(dir1, fmt1)
164
+ live = _load_dir(dir2, fmt2)
165
+
166
+ schema_diffs: list[SchemaDiff] = []
167
+ for name, stored_schema in stored.items():
168
+ if name not in live:
169
+ schema_diffs.append(SchemaDiff(name=name, status="removed"))
170
+ else:
171
+ schema_diffs.append(diff_schemas(live[name], stored_schema, include_metadata))
172
+ for name in live:
173
+ if name not in stored:
174
+ schema_diffs.append(SchemaDiff(name=name, status="added"))
175
+
176
+ return CatalogDiff(schemas=schema_diffs)
177
+
178
+
131
179
  def diff_catalog_with_dir(
132
180
  catalog: Catalog,
133
181
  schema_dir: Path,
@@ -3,6 +3,7 @@ from __future__ import annotations
3
3
  import logging
4
4
  from collections.abc import Iterator
5
5
  from concurrent.futures import ThreadPoolExecutor
6
+ from urllib.parse import quote
6
7
 
7
8
  from databricks.sdk import WorkspaceClient
8
9
  from databricks.sdk.errors import NotFound
@@ -22,7 +23,10 @@ class CatalogExtractor:
22
23
  def _fetch_tags(self, entity_type: str, entity_name: str) -> dict[str, str]:
23
24
  tags: dict[str, str] = {}
24
25
  try:
25
- for assignment in self.client.entity_tag_assignments.list(entity_type, entity_name):
26
+ # Percent-encode entity_name so special characters (e.g. slashes in column names)
27
+ # don't break URL path routing. Dots are safe — they're the UC name separators.
28
+ encoded_name = quote(entity_name, safe=".")
29
+ for assignment in self.client.entity_tag_assignments.list(entity_type, encoded_name):
26
30
  key = getattr(assignment, "tag_key", None)
27
31
  value = getattr(assignment, "tag_value", None)
28
32
  if key is not None:
@@ -0,0 +1,136 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass, field
4
+
5
+ from databricks_schema.models import Schema, Table
6
+
7
+
8
+ @dataclass
9
+ class ValidationIssue:
10
+ schema: str
11
+ table: str | None
12
+ message: str
13
+
14
+ def __str__(self) -> str:
15
+ location = f"{self.schema}.{self.table}" if self.table else self.schema
16
+ return f"{location}: {self.message}"
17
+
18
+
19
+ @dataclass
20
+ class ValidationResult:
21
+ issues: list[ValidationIssue] = field(default_factory=list)
22
+
23
+ @property
24
+ def has_errors(self) -> bool:
25
+ return bool(self.issues)
26
+
27
+
28
+ def _validate_table(
29
+ schema_name: str, table: Table, all_schemas: dict[str, Schema]
30
+ ) -> list[ValidationIssue]:
31
+ issues: list[ValidationIssue] = []
32
+ col_names = [c.name for c in table.columns]
33
+ col_set = set(col_names)
34
+
35
+ # Duplicate column names
36
+ seen: set[str] = set()
37
+ for name in col_names:
38
+ if name in seen:
39
+ issues.append(
40
+ ValidationIssue(
41
+ schema=schema_name,
42
+ table=table.name,
43
+ message=f"duplicate column name: '{name}'",
44
+ )
45
+ )
46
+ seen.add(name)
47
+
48
+ # PK columns exist in table
49
+ if table.primary_key:
50
+ for col in table.primary_key.columns:
51
+ if col not in col_set:
52
+ issues.append(
53
+ ValidationIssue(
54
+ schema=schema_name,
55
+ table=table.name,
56
+ message=f"primary key references unknown column: '{col}'",
57
+ )
58
+ )
59
+
60
+ # FK checks
61
+ for fk in table.foreign_keys:
62
+ # FK source columns exist in this table
63
+ for col in fk.columns:
64
+ if col not in col_set:
65
+ issues.append(
66
+ ValidationIssue(
67
+ schema=schema_name,
68
+ table=table.name,
69
+ message=f"foreign key references unknown source column: '{col}'",
70
+ )
71
+ )
72
+
73
+ # FK ref_schema + ref_table exist in loaded schemas
74
+ ref_schema = all_schemas.get(fk.ref_schema)
75
+ if ref_schema is None:
76
+ issues.append(
77
+ ValidationIssue(
78
+ schema=schema_name,
79
+ table=table.name,
80
+ message=f"foreign key references unknown schema: '{fk.ref_schema}'",
81
+ )
82
+ )
83
+ else:
84
+ ref_table_map = {t.name: t for t in ref_schema.tables}
85
+ ref_table = ref_table_map.get(fk.ref_table)
86
+ if ref_table is None:
87
+ issues.append(
88
+ ValidationIssue(
89
+ schema=schema_name,
90
+ table=table.name,
91
+ message=(
92
+ f"foreign key references unknown table:"
93
+ f" '{fk.ref_schema}.{fk.ref_table}'"
94
+ ),
95
+ )
96
+ )
97
+ else:
98
+ # FK ref_columns exist in referenced table
99
+ ref_col_set = {c.name for c in ref_table.columns}
100
+ for col in fk.ref_columns:
101
+ if col not in ref_col_set:
102
+ issues.append(
103
+ ValidationIssue(
104
+ schema=schema_name,
105
+ table=table.name,
106
+ message=(
107
+ f"foreign key references unknown column '{col}'"
108
+ f" in '{fk.ref_schema}.{fk.ref_table}'"
109
+ ),
110
+ )
111
+ )
112
+
113
+ return issues
114
+
115
+
116
+ def validate_schemas(schemas: dict[str, Schema]) -> ValidationResult:
117
+ """Validate a collection of schemas for structural integrity.
118
+
119
+ Checks performed (no SDK or I/O required):
120
+ - Duplicate column names within a table
121
+ - PK columns that reference columns not present in the table
122
+ - FK source columns that reference columns not present in the table
123
+ - FK ref_schema + ref_table combinations not found in the loaded schemas
124
+ - FK ref_columns not found in the referenced table
125
+
126
+ Args:
127
+ schemas: mapping of schema name → Schema, as loaded from a directory.
128
+
129
+ Returns:
130
+ ValidationResult with a list of ValidationIssue for every problem found.
131
+ """
132
+ issues: list[ValidationIssue] = []
133
+ for schema_name, schema in schemas.items():
134
+ for table in schema.tables:
135
+ issues.extend(_validate_table(schema_name, table, schemas))
136
+ return ValidationResult(issues=issues)
@@ -0,0 +1,19 @@
1
+ from __future__ import annotations
2
+
3
+ from databricks_schema import CatalogExtractor
4
+
5
+ # Extract a single schema and inspect its contents.
6
+ # Uses default Databricks auth (env vars, ~/.databrickscfg, etc.)
7
+
8
+ CATALOG = "lake_prod"
9
+ SCHEMA = ["gold", "silver"]
10
+
11
+ extractor = CatalogExtractor()
12
+ catalog = extractor.extract_catalog(CATALOG, schema_filter=SCHEMA)
13
+
14
+ schema = catalog.schemas[0]
15
+ print(f"Schema: {schema.name} ({len(schema.tables)} tables)")
16
+
17
+ for table in schema.tables:
18
+ col_names = [c.name for c in table.columns]
19
+ print(f" {table.name} [{table.table_type}]: {', '.join(col_names)}")
@@ -0,0 +1,30 @@
1
+ from __future__ import annotations
2
+
3
+ from pathlib import Path
4
+
5
+ from databricks_schema import CatalogExtractor, schema_from_yaml, schema_to_yaml
6
+
7
+ # Extract schemas, save each to a YAML file, then reload and inspect the result.
8
+ # Uses default Databricks auth (env vars, ~/.databrickscfg, etc.)
9
+
10
+ CATALOG = "lake_prod"
11
+ SCHEMAS = ["gold", "silver"]
12
+ OUTPUT_DIR = Path(".test_output/02_save_load_yaml")
13
+
14
+ extractor = CatalogExtractor()
15
+ catalog = extractor.extract_catalog(CATALOG, schema_filter=SCHEMAS)
16
+
17
+ OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
18
+
19
+ for schema in catalog.schemas:
20
+ path = OUTPUT_DIR / f"{schema.name}.yaml"
21
+ path.write_text(schema_to_yaml(schema), encoding="utf-8")
22
+ print(f"Saved {path}")
23
+
24
+ # Reload and inspect
25
+ for schema in catalog.schemas:
26
+ path = OUTPUT_DIR / f"{schema.name}.yaml"
27
+ loaded = schema_from_yaml(path.read_text(encoding="utf-8"))
28
+ print(f"Reloaded '{loaded.name}': {len(loaded.tables)} tables")
29
+ for table in loaded.tables:
30
+ print(f" {table.name} ({len(table.columns)} columns)")
@@ -0,0 +1,44 @@
1
+ from __future__ import annotations
2
+
3
+ from pathlib import Path
4
+
5
+ from databricks_schema import CatalogExtractor, diff_catalog_with_dir, schema_to_yaml
6
+
7
+ # Diff live catalog schemas against local YAML files.
8
+ # Saves a baseline on first run, then diffs live state against it.
9
+ # Uses default Databricks auth (env vars, ~/.databrickscfg, etc.)
10
+
11
+ CATALOG = "lake_prod"
12
+ SCHEMAS = ["gold", "silver"]
13
+ SCHEMA_DIR = Path(".test_output/03_diff_schemas")
14
+
15
+ extractor = CatalogExtractor()
16
+ catalog = extractor.extract_catalog(CATALOG, schema_filter=SCHEMAS)
17
+
18
+ # Save baseline if not already present
19
+ SCHEMA_DIR.mkdir(parents=True, exist_ok=True)
20
+ for schema in catalog.schemas:
21
+ path = SCHEMA_DIR / f"{schema.name}.yaml"
22
+ if not path.exists():
23
+ path.write_text(schema_to_yaml(schema), encoding="utf-8")
24
+ print(f"Saved baseline {path}")
25
+
26
+ result = diff_catalog_with_dir(catalog, SCHEMA_DIR, schema_names=frozenset(SCHEMAS))
27
+
28
+ if not result.has_changes:
29
+ print("No changes — live catalog matches local files.")
30
+ else:
31
+ for schema_diff in result.schemas:
32
+ if not schema_diff.has_changes:
33
+ continue
34
+ print(f"\nSchema '{schema_diff.name}': {schema_diff.status}")
35
+ for fc in schema_diff.changes:
36
+ print(f" field '{fc.field}': {fc.old!r} -> {fc.new!r}")
37
+ for table_diff in schema_diff.tables:
38
+ print(f" table '{table_diff.name}': {table_diff.status}")
39
+ for fc in table_diff.changes:
40
+ print(f" field '{fc.field}': {fc.old!r} -> {fc.new!r}")
41
+ for col_diff in table_diff.columns:
42
+ print(f" column '{col_diff.name}': {col_diff.status}")
43
+ for fc in col_diff.changes:
44
+ print(f" {fc.field}: {fc.old!r} -> {fc.new!r}")
@@ -0,0 +1,41 @@
1
+ from __future__ import annotations
2
+
3
+ from pathlib import Path
4
+
5
+ from databricks_schema import (
6
+ CatalogExtractor,
7
+ diff_schemas,
8
+ schema_diff_to_sql,
9
+ schema_from_yaml,
10
+ schema_to_yaml,
11
+ )
12
+
13
+ # Generate SQL to reconcile a live schema with a stored YAML baseline.
14
+ # Simulates drift by modifying the stored YAML, then generates SQL to fix it.
15
+ # Uses default Databricks auth (env vars, ~/.databrickscfg, etc.)
16
+
17
+ CATALOG = "lake_prod"
18
+ SCHEMA = "gold"
19
+ OUTPUT_DIR = Path(".test_output/04_generate_sql")
20
+
21
+ extractor = CatalogExtractor()
22
+ catalog = extractor.extract_catalog(CATALOG, schema_filter=[SCHEMA])
23
+ live_schema = catalog.schemas[0]
24
+
25
+ # Save current state as the baseline
26
+ OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
27
+ baseline_path = OUTPUT_DIR / f"{SCHEMA}.yaml"
28
+ baseline_path.write_text(schema_to_yaml(live_schema), encoding="utf-8")
29
+ print(f"Saved baseline {baseline_path}")
30
+
31
+ # Load stored schema (this is your "desired state")
32
+ stored_schema = schema_from_yaml(baseline_path.read_text(encoding="utf-8"))
33
+
34
+ # Diff live vs stored
35
+ schema_diff = diff_schemas(live=live_schema, stored=stored_schema)
36
+
37
+ if not schema_diff.has_changes:
38
+ print("No drift detected — nothing to generate.")
39
+ else:
40
+ sql = schema_diff_to_sql(CATALOG, schema_diff, stored_schema, allow_drop=False)
41
+ print(sql)
@@ -7,7 +7,7 @@ packages = ["databricks_schema"]
7
7
 
8
8
  [project]
9
9
  name = "databricks-schema"
10
- version = "0.5.1"
10
+ version = "0.6.1"
11
11
  description = "Databricks Unity Catalog schema extractor"
12
12
  readme = "README.md"
13
13
  requires-python = ">=3.11"
@@ -4,6 +4,7 @@ from pathlib import Path
4
4
 
5
5
  from databricks_schema.diff import (
6
6
  diff_catalog_with_dir,
7
+ diff_schema_dirs,
7
8
  diff_schemas,
8
9
  )
9
10
  from databricks_schema.models import Catalog, Column, ForeignKey, PrimaryKey, Schema, Table
@@ -246,3 +247,84 @@ class TestDiffCatalogWithDir:
246
247
  result = diff_catalog_with_dir(catalog, tmp_path, fmt="json")
247
248
  assert not result.has_changes
248
249
  assert result.schemas[0].status == "unchanged"
250
+
251
+
252
+ class TestDiffSchemaDirs:
253
+ def test_no_changes(self, tmp_path: Path):
254
+ d1 = tmp_path / "d1"
255
+ d2 = tmp_path / "d2"
256
+ d1.mkdir()
257
+ d2.mkdir()
258
+ schema = _schema("main", tables=[_table("users", columns=[_col("id")])])
259
+ (d1 / "main.yaml").write_text(schema_to_yaml(schema))
260
+ (d2 / "main.yaml").write_text(schema_to_yaml(schema))
261
+ result = diff_schema_dirs(d1, d2)
262
+ assert not result.has_changes
263
+ assert result.schemas[0].status == "unchanged"
264
+
265
+ def test_schema_added_in_dir2(self, tmp_path: Path):
266
+ d1 = tmp_path / "d1"
267
+ d2 = tmp_path / "d2"
268
+ d1.mkdir()
269
+ d2.mkdir()
270
+ schema = _schema("main")
271
+ (d1 / "main.yaml").write_text(schema_to_yaml(schema))
272
+ (d2 / "main.yaml").write_text(schema_to_yaml(schema))
273
+ extra = _schema("extra")
274
+ (d2 / "extra.yaml").write_text(schema_to_yaml(extra))
275
+ result = diff_schema_dirs(d1, d2)
276
+ added = next(s for s in result.schemas if s.name == "extra")
277
+ assert added.status == "added"
278
+
279
+ def test_schema_removed_from_dir2(self, tmp_path: Path):
280
+ d1 = tmp_path / "d1"
281
+ d2 = tmp_path / "d2"
282
+ d1.mkdir()
283
+ d2.mkdir()
284
+ schema = _schema("main")
285
+ (d1 / "main.yaml").write_text(schema_to_yaml(schema))
286
+ result = diff_schema_dirs(d1, d2)
287
+ assert result.schemas[0].name == "main"
288
+ assert result.schemas[0].status == "removed"
289
+
290
+ def test_field_modified(self, tmp_path: Path):
291
+ d1 = tmp_path / "d1"
292
+ d2 = tmp_path / "d2"
293
+ d1.mkdir()
294
+ d2.mkdir()
295
+ old = _schema("main", comment="old")
296
+ new = _schema("main", comment="new")
297
+ (d1 / "main.yaml").write_text(schema_to_yaml(old))
298
+ (d2 / "main.yaml").write_text(schema_to_yaml(new))
299
+ result = diff_schema_dirs(d1, d2)
300
+ s = result.schemas[0]
301
+ assert s.status == "modified"
302
+ assert s.changes[0].field == "comment"
303
+ assert s.changes[0].old == "old"
304
+ assert s.changes[0].new == "new"
305
+
306
+ def test_json_format(self, tmp_path: Path):
307
+ d1 = tmp_path / "d1"
308
+ d2 = tmp_path / "d2"
309
+ d1.mkdir()
310
+ d2.mkdir()
311
+ schema = _schema("main", tables=[_table("t", columns=[_col("id")])])
312
+ (d1 / "main.json").write_text(schema_to_json(schema))
313
+ (d2 / "main.json").write_text(schema_to_json(schema))
314
+ result = diff_schema_dirs(d1, d2, fmt1="json", fmt2="json")
315
+ assert not result.has_changes
316
+
317
+ def test_schema_names_filter(self, tmp_path: Path):
318
+ d1 = tmp_path / "d1"
319
+ d2 = tmp_path / "d2"
320
+ d1.mkdir()
321
+ d2.mkdir()
322
+ main = _schema("main")
323
+ raw_old = _schema("raw", comment="old")
324
+ raw_new = _schema("raw", comment="new")
325
+ (d1 / "main.yaml").write_text(schema_to_yaml(main))
326
+ (d1 / "raw.yaml").write_text(schema_to_yaml(raw_old))
327
+ (d2 / "main.yaml").write_text(schema_to_yaml(main))
328
+ (d2 / "raw.yaml").write_text(schema_to_yaml(raw_new))
329
+ result = diff_schema_dirs(d1, d2, schema_names=frozenset({"main"}))
330
+ assert not result.has_changes
@@ -398,6 +398,33 @@ class TestCatalogExtractor:
398
398
 
399
399
  client.tables.get.assert_not_called()
400
400
 
401
+ def test_column_tags_slash_in_name(self):
402
+ """Column names containing slashes must be percent-encoded before tag lookup."""
403
+ extractor, client = self._extractor()
404
+ sdk_catalog = MagicMock()
405
+ sdk_catalog.comment = None
406
+ client.catalogs.get.return_value = sdk_catalog
407
+
408
+ s = MagicMock()
409
+ s.name = "main"
410
+ s.comment = None
411
+ s.owner = None
412
+ client.schemas.list.return_value = [s]
413
+
414
+ col = _make_col("/example/column", position=1, type_text="STRING")
415
+ sdk_table = _make_sdk_table("orders", columns=[col])
416
+ client.tables.list.return_value = [sdk_table]
417
+
418
+ def mock_list(entity_type, entity_name):
419
+ if entity_type == "columns" and entity_name == "mycat.main.orders.%2Fexample%2Fcolumn":
420
+ return [_make_tag("sensitivity", "high")]
421
+ return []
422
+
423
+ client.entity_tag_assignments.list.side_effect = mock_list
424
+
425
+ catalog = extractor.extract_catalog("mycat")
426
+ assert catalog.schemas[0].tables[0].columns[0].tags == {"sensitivity": "high"}
427
+
401
428
  def test_parallel_extraction(self):
402
429
  """With max_workers > 1, all tables are still extracted correctly."""
403
430
  extractor, client = self._extractor(max_workers=4)
@@ -0,0 +1,146 @@
1
+ from __future__ import annotations
2
+
3
+ from databricks_schema.models import Column, ForeignKey, PrimaryKey, Schema, Table
4
+ from databricks_schema.validate import validate_schemas
5
+
6
+
7
+ def _col(name: str, data_type: str = "STRING") -> Column:
8
+ return Column(name=name, data_type=data_type)
9
+
10
+
11
+ def _table(name: str, columns: list[Column] | None = None, **kwargs) -> Table:
12
+ return Table(name=name, columns=columns or [], **kwargs)
13
+
14
+
15
+ def _schema(name: str, tables: list[Table] | None = None) -> Schema:
16
+ return Schema(name=name, tables=tables or [])
17
+
18
+
19
+ def _schemas(*schemas: Schema) -> dict[str, Schema]:
20
+ return {s.name: s for s in schemas}
21
+
22
+
23
+ class TestValidateClean:
24
+ def test_empty_schemas(self):
25
+ result = validate_schemas({})
26
+ assert not result.has_errors
27
+ assert result.issues == []
28
+
29
+ def test_valid_schema_no_constraints(self):
30
+ schema = _schema("main", [_table("users", [_col("id"), _col("email")])])
31
+ result = validate_schemas(_schemas(schema))
32
+ assert not result.has_errors
33
+
34
+ def test_valid_pk(self):
35
+ table = _table("users", [_col("id")], primary_key=PrimaryKey(columns=["id"]))
36
+ result = validate_schemas(_schemas(_schema("main", [table])))
37
+ assert not result.has_errors
38
+
39
+ def test_valid_fk_cross_schema(self):
40
+ orgs = _schema("orgs", [_table("org", [_col("id")])])
41
+ fk = ForeignKey(columns=["org_id"], ref_schema="orgs", ref_table="org", ref_columns=["id"])
42
+ users = _schema("main", [_table("users", [_col("id"), _col("org_id")], foreign_keys=[fk])])
43
+ result = validate_schemas(_schemas(orgs, users))
44
+ assert not result.has_errors
45
+
46
+
47
+ class TestDuplicateColumns:
48
+ def test_duplicate_column_name(self):
49
+ table = _table("t", [_col("id"), _col("id")])
50
+ result = validate_schemas(_schemas(_schema("main", [table])))
51
+ assert result.has_errors
52
+ assert len(result.issues) == 1
53
+ assert "duplicate column name" in result.issues[0].message
54
+ assert "'id'" in result.issues[0].message
55
+
56
+ def test_duplicate_reported_once_per_name(self):
57
+ # three columns named "x" — only one issue (detected on second occurrence)
58
+ table = _table("t", [_col("x"), _col("x"), _col("x")])
59
+ result = validate_schemas(_schemas(_schema("main", [table])))
60
+ assert result.has_errors
61
+ # second and third "x" both trigger, so two issues
62
+ assert len(result.issues) == 2
63
+
64
+
65
+ class TestPrimaryKeyValidation:
66
+ def test_pk_column_missing(self):
67
+ table = _table("t", [_col("id")], primary_key=PrimaryKey(columns=["nonexistent"]))
68
+ result = validate_schemas(_schemas(_schema("main", [table])))
69
+ assert result.has_errors
70
+ assert any(
71
+ "primary key" in i.message and "'nonexistent'" in i.message for i in result.issues
72
+ )
73
+
74
+ def test_pk_multiple_columns_one_missing(self):
75
+ table = _table("t", [_col("a"), _col("b")], primary_key=PrimaryKey(columns=["a", "c"]))
76
+ result = validate_schemas(_schemas(_schema("main", [table])))
77
+ assert result.has_errors
78
+ assert len(result.issues) == 1
79
+ assert "'c'" in result.issues[0].message
80
+
81
+
82
+ class TestForeignKeyValidation:
83
+ def test_fk_source_column_missing(self):
84
+ fk = ForeignKey(columns=["bad_col"], ref_schema="other", ref_table="t", ref_columns=["id"])
85
+ other = _schema("other", [_table("t", [_col("id")])])
86
+ main = _schema("main", [_table("users", [_col("id")], foreign_keys=[fk])])
87
+ result = validate_schemas(_schemas(other, main))
88
+ assert result.has_errors
89
+ assert any(
90
+ "unknown source column" in i.message and "'bad_col'" in i.message for i in result.issues
91
+ )
92
+
93
+ def test_fk_ref_schema_missing(self):
94
+ fk = ForeignKey(columns=["org_id"], ref_schema="ghost", ref_table="org", ref_columns=["id"])
95
+ main = _schema("main", [_table("users", [_col("id"), _col("org_id")], foreign_keys=[fk])])
96
+ result = validate_schemas(_schemas(main))
97
+ assert result.has_errors
98
+ assert any("unknown schema" in i.message and "'ghost'" in i.message for i in result.issues)
99
+
100
+ def test_fk_ref_table_missing(self):
101
+ fk = ForeignKey(
102
+ columns=["org_id"], ref_schema="other", ref_table="ghost", ref_columns=["id"]
103
+ )
104
+ other = _schema("other", [_table("real_table", [_col("id")])])
105
+ main = _schema("main", [_table("users", [_col("id"), _col("org_id")], foreign_keys=[fk])])
106
+ result = validate_schemas(_schemas(other, main))
107
+ assert result.has_errors
108
+ assert any(
109
+ "unknown table" in i.message and "other.ghost" in i.message for i in result.issues
110
+ )
111
+
112
+ def test_fk_ref_column_missing(self):
113
+ fk = ForeignKey(
114
+ columns=["org_id"], ref_schema="other", ref_table="org", ref_columns=["bad"]
115
+ )
116
+ other = _schema("other", [_table("org", [_col("id")])])
117
+ main = _schema("main", [_table("users", [_col("id"), _col("org_id")], foreign_keys=[fk])])
118
+ result = validate_schemas(_schemas(other, main))
119
+ assert result.has_errors
120
+ assert any("'bad'" in i.message and "other.org" in i.message for i in result.issues)
121
+
122
+ def test_multiple_fk_issues_reported(self):
123
+ fk = ForeignKey(
124
+ columns=["bad_src"],
125
+ ref_schema="ghost",
126
+ ref_table="t",
127
+ ref_columns=["id"],
128
+ )
129
+ main = _schema("main", [_table("t", [_col("id")], foreign_keys=[fk])])
130
+ result = validate_schemas(_schemas(main))
131
+ # bad_src missing + ghost schema missing = 2 issues
132
+ assert len(result.issues) == 2
133
+
134
+
135
+ class TestIssueStr:
136
+ def test_str_with_table(self):
137
+ from databricks_schema.validate import ValidationIssue
138
+
139
+ issue = ValidationIssue(schema="main", table="users", message="some problem")
140
+ assert str(issue) == "main.users: some problem"
141
+
142
+ def test_str_without_table(self):
143
+ from databricks_schema.validate import ValidationIssue
144
+
145
+ issue = ValidationIssue(schema="main", table=None, message="schema-level problem")
146
+ assert str(issue) == "main: schema-level problem"
@@ -233,7 +233,7 @@ wheels = [
233
233
 
234
234
  [[package]]
235
235
  name = "databricks-schema"
236
- version = "0.5.1"
236
+ version = "0.6.1"
237
237
  source = { editable = "." }
238
238
  dependencies = [
239
239
  { name = "databricks-sdk" },
@@ -1,49 +0,0 @@
1
- import sys
2
- from pathlib import Path
3
-
4
- from databricks_schema import (
5
- CatalogExtractor,
6
- catalog_to_yaml,
7
- diff_catalog_with_dir,
8
- diff_schemas,
9
- schema_from_yaml,
10
- schema_to_yaml,
11
- )
12
-
13
- # Configuration
14
- CATALOG_NAME = "lake_prod"
15
- SCHEMA_FILTER = ["gold", "silver"]
16
- SINGLE_SCHEMA = "gold"
17
- OUTPUT_DIR = Path(".test_output/example")
18
-
19
- # Extract using configured auth
20
- extractor = CatalogExtractor()
21
- catalog = extractor.extract_catalog(CATALOG_NAME, schema_filter=SCHEMA_FILTER)
22
-
23
- # Serialise to YAML String
24
- yaml_text = catalog_to_yaml(catalog)
25
-
26
- # Serialise to YAML file(s)
27
- OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
28
-
29
- for s in catalog.schemas:
30
- out_file = OUTPUT_DIR / f"{s.name}.yaml"
31
- out_file.write_text(schema_to_yaml(s), encoding="utf-8")
32
- print(f" Wrote {out_file}", file=sys.stderr)
33
-
34
- print(f"Done — {len(catalog.schemas)} schema(s) written to {OUTPUT_DIR}", file=sys.stderr)
35
-
36
- # Deserialise from YAML
37
- schema = schema_from_yaml(open(OUTPUT_DIR / f"{SINGLE_SCHEMA}.yaml").read())
38
- print(schema.tables[0].columns)
39
-
40
- # Compare live catalog against local YAML files
41
- result = diff_catalog_with_dir(catalog, OUTPUT_DIR)
42
- if result.has_changes:
43
- for schema_diff in result.schemas:
44
- print(schema_diff.name, schema_diff.status)
45
-
46
- # Compare two Schema objects directly
47
- stored = schema_from_yaml(open(OUTPUT_DIR / f"{SINGLE_SCHEMA}.yaml").read())
48
- live = extractor._extract_schema(CATALOG_NAME, schema)
49
- diff = diff_schemas(live=live, stored=stored)