al-data-utils 0.3.1__py3-none-any.whl

al_data_utils/backup_utils/mssql/schema_snapshot/config.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class SnapshotConfig:
+    """Configuration for MSSQL schema snapshot.
+
+    Notes:
+    - output_root: final destination folder.
+    - staging_root: if None, staging folder is created as sibling: f"{output_root}__staging".
+    - snapshot writes into staging first, then syncs staging -> output_root using files_sync_utils.
+    """
+
+    output_root: Path
+    staging_root: Path | None = None
+    cleanup_staging: bool = True
+    staging_database_name: str | None = None
+    staging_group_by_database: bool = True
+    staging_include_run_guid: bool = True
+    sync_log_folder: Path | None = None
+    sync_log_retention_days: int | None = 30
+    verbose: bool = True
+
+    include_indexes: bool = True
+    include_foreign_keys: bool = True
+    include_table_definitions: bool = True
+    include_views: bool = True
+    include_stored_procedures: bool = True
+    include_functions: bool = True
+    include_sequences: bool = True
+    include_synonyms: bool = True
+
+    root_folder_name: str = 'schemas'
+
+    batch_separator: str = 'go'
+    header_comment: str = ''
+    section_comments: bool = False
+    newline: str = '\n'
+    encoding: str = 'utf-8'
+
+    datagrip_templates_root: Path | None = None
+    prefer_template_files: bool = True
+
+    include_schemas: set[str] | None = None
+    exclude_schemas: set[str] | None = None
+
+    include_tables: set[str] | None = None
+    exclude_tables: set[str] | None = None
+
+    include_views_objects: set[str] | None = None
+    exclude_views_objects: set[str] | None = None
+
+    include_stored_procedures_objects: set[str] | None = None
+    exclude_stored_procedures_objects: set[str] | None = None
+
+    include_functions_objects: set[str] | None = None
+    exclude_functions_objects: set[str] | None = None
+
+    include_sequences_objects: set[str] | None = None
+    exclude_sequences_objects: set[str] | None = None
+
+    include_synonyms_objects: set[str] | None = None
+    exclude_synonyms_objects: set[str] | None = None
+
+    def is_schema_allowed(self, schema: str) -> bool:
+        s = schema.lower()
+        if self.include_schemas is not None and s not in {
+            x.lower() for x in self.include_schemas
+        }:
+            return False
+        if self.exclude_schemas is not None and s in {
+            x.lower() for x in self.exclude_schemas
+        }:
+            return False
+        return True
+
+    def is_table_allowed(self, schema: str, table: str) -> bool:
+        key = f'{schema}.{table}'.lower()
+        if self.include_tables is not None and key not in {
+            x.lower() for x in self.include_tables
+        }:
+            return False
+        if self.exclude_tables is not None and key in {
+            x.lower() for x in self.exclude_tables
+        }:
+            return False
+        return True
+
+    def is_view_allowed(self, schema: str, view_name: str) -> bool:
+        key = f'{schema}.{view_name}'.lower()
+        if self.include_views_objects is not None and key not in {
+            x.lower() for x in self.include_views_objects
+        }:
+            return False
+        if self.exclude_views_objects is not None and key in {
+            x.lower() for x in self.exclude_views_objects
+        }:
+            return False
+        return True
+
+    def is_stored_procedure_allowed(self, schema: str, procedure_name: str) -> bool:
+        key = f'{schema}.{procedure_name}'.lower()
+        if self.include_stored_procedures_objects is not None and key not in {
+            x.lower() for x in self.include_stored_procedures_objects
+        }:
+            return False
+        if self.exclude_stored_procedures_objects is not None and key in {
+            x.lower() for x in self.exclude_stored_procedures_objects
+        }:
+            return False
+        return True
+
+    def is_function_allowed(self, schema: str, function_name: str) -> bool:
+        key = f'{schema}.{function_name}'.lower()
+        if self.include_functions_objects is not None and key not in {
+            x.lower() for x in self.include_functions_objects
+        }:
+            return False
+        if self.exclude_functions_objects is not None and key in {
+            x.lower() for x in self.exclude_functions_objects
+        }:
+            return False
+        return True
+
+    def is_sequence_allowed(self, schema: str, sequence_name: str) -> bool:
+        key = f'{schema}.{sequence_name}'.lower()
+        if self.include_sequences_objects is not None and key not in {
+            x.lower() for x in self.include_sequences_objects
+        }:
+            return False
+        if self.exclude_sequences_objects is not None and key in {
+            x.lower() for x in self.exclude_sequences_objects
+        }:
+            return False
+        return True
+
+    def is_synonym_allowed(self, schema: str, synonym_name: str) -> bool:
+        key = f'{schema}.{synonym_name}'.lower()
+        if self.include_synonyms_objects is not None and key not in {
+            x.lower() for x in self.include_synonyms_objects
+        }:
+            return False
+        if self.exclude_synonyms_objects is not None and key in {
+            x.lower() for x in self.exclude_synonyms_objects
+        }:
+            return False
+        return True

al_data_utils/backup_utils/mssql/schema_snapshot/emit/__init__.py

@@ -0,0 +1,11 @@
+from .indexes import emit_index_ddl
+from .foreign_keys import emit_foreign_key_ddl
+from .tables import emit_table_ddl
+from .sql_objects import emit_sql_object_ddl
+
+__all__ = [
+    'emit_index_ddl',
+    'emit_foreign_key_ddl',
+    'emit_table_ddl',
+    'emit_sql_object_ddl',
+]

al_data_utils/backup_utils/mssql/schema_snapshot/emit/foreign_keys.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from ..models import ForeignKeyDef
+
+
+def _emit_action(action_desc: str) -> str:
+    a = (action_desc or "").upper()
+    if a in ("NO_ACTION", ""):
+        return ""
+    if a == "CASCADE":
+        return "cascade"
+    if a == "SET_NULL":
+        return "set null"
+    if a == "SET_DEFAULT":
+        return "set default"
+    return action_desc.lower()
+
+
+def emit_foreign_key_ddl(fk: ForeignKeyDef) -> str:
+    parent = f"{fk.parent_schema}.{fk.parent_table}"
+    ref = f"{fk.referenced_schema}.{fk.referenced_table}"
+
+    ddl = (
+        f"alter table {parent}\n"
+        f"    add constraint {fk.fk_name}\n"
+        f"        foreign key ({fk.parent_columns})\n"
+        f"            references {ref} ({fk.referenced_columns})"
+    )
+
+    on_delete = _emit_action(fk.on_delete)
+    if on_delete:
+        ddl += f"\n            on delete {on_delete}"
+
+    on_update = _emit_action(fk.on_update)
+    if on_update:
+        ddl += f"\n            on update {on_update}"
+
+    return ddl

al_data_utils/backup_utils/mssql/schema_snapshot/emit/indexes.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import re
+
+from ..models import IndexDef
+
+
+def _strip_brackets(expr: str) -> str:
+    return re.sub(r'\[([^\]]+)\]', r'\1', expr or '')
+
+
+def _create_index_prefix(idx: IndexDef) -> str:
+    t = (idx.type_desc or '').strip().lower()
+    if idx.is_unique and t and t != 'nonclustered':
+        return f'create unique {t} index'
+    if idx.is_unique:
+        return 'create unique index'
+    if t and t != 'nonclustered':
+        return f'create {t} index'
+    return 'create index'
+
+
+def emit_index_ddl(idx: IndexDef) -> str:
+    schema_table = f'{idx.schema_name}.{idx.table_name}'
+    key_columns = _strip_brackets(idx.key_columns)
+    include_columns = _strip_brackets(idx.include_columns)
+
+    ddl = f'{_create_index_prefix(idx)} {idx.index_name}\n'
+    ddl += f'    on {schema_table} ({key_columns})'
+
+    if include_columns:
+        ddl += f' include ({include_columns})'
+
+    if idx.has_filter and idx.filter_definition:
+        ddl += f'\n    where {idx.filter_definition}'
+
+    return ddl

al_data_utils/backup_utils/mssql/schema_snapshot/emit/sql_objects.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from ..models import SqlObjectDef
+
+
+def emit_sql_object_ddl(obj: SqlObjectDef) -> str:
+    blocks: list[str] = []
+    definition = (obj.definition or '').strip()
+    if definition:
+        blocks.append(definition)
+
+    for stmt in obj.permissions_sql or []:
+        s = (stmt or '').strip()
+        if s:
+            blocks.append(s)
+
+    return '\n\ngo\n\n'.join(blocks).strip()

al_data_utils/backup_utils/mssql/schema_snapshot/emit/tables.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import re
+
+from ..models import TableDef, TableColumnDef, ForeignKeyDef
+
+
+def _q(name: str) -> str:
+    n = name.strip()
+    if re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', n):
+        return n
+    return '[' + n.replace(']', ']]') + ']'
+
+
+def _strip_brackets(expr: str) -> str:
+    return re.sub(r'\[([^\]]+)\]', r'\1', expr or '')
+
+
+def _sql_type(col: TableColumnDef) -> str:
+    t = col.data_type.lower()
+
+    if t in {'nvarchar', 'nchar'}:
+        if col.max_length is None:
+            return t
+        if col.max_length == -1:
+            return f'{t}(max)'
+        return f'{t}({int(col.max_length / 2)})'
+
+    if t in {'varchar', 'char', 'varbinary', 'binary'}:
+        if col.max_length is None:
+            return t
+        if col.max_length == -1:
+            return f'{t}(max)'
+        return f'{t}({col.max_length})'
+
+    if t in {'decimal', 'numeric'}:
+        p = col.precision if col.precision is not None else 18
+        s = col.scale if col.scale is not None else 0
+        return f'{t}({p},{s})'
+
+    if t in {'datetime2', 'datetimeoffset', 'time'}:
+        if col.scale is None:
+            return t
+        return f'{t}({col.scale})'
+
+    return t
+
+
+def _emit_column(col: TableColumnDef) -> str:
+    name = _q(col.column_name)
+
+    if col.is_computed:
+        expr = col.computed_definition or ''
+        return f'{name} as {expr}'
+
+    line = f'{name} {_sql_type(col)}'
+
+    if col.collation_name and col.data_type.lower() in {
+        'char',
+        'varchar',
+        'nchar',
+        'nvarchar',
+    }:
+        line += f' collate {col.collation_name}'
+
+    if col.is_identity:
+        seed = col.identity_seed if col.identity_seed is not None else 1
+        inc = col.identity_increment if col.identity_increment is not None else 1
+        line += f' identity({seed},{inc})'
+
+    if col.default_definition:
+        if col.default_constraint_name:
+            line += f' constraint {_q(col.default_constraint_name)} default {col.default_definition}'
+        else:
+            line += f' default {col.default_definition}'
+
+    if not col.is_nullable:
+        line += ' not null'
+    return line
+
+
+def _emit_fk_action(action_desc: str) -> str:
+    a = (action_desc or '').upper()
+    if a in {'', 'NO_ACTION'}:
+        return ''
+    if a == 'CASCADE':
+        return 'cascade'
+    if a == 'SET_NULL':
+        return 'set null'
+    if a == 'SET_DEFAULT':
+        return 'set default'
+    return a.lower()
+
+
+def _emit_inline_fk_constraint(fk: ForeignKeyDef) -> str:
+    ref_table = f'{_q(fk.referenced_schema)}.{_q(fk.referenced_table)}'
+    parent_cols = _strip_brackets(fk.parent_columns)
+    referenced_cols = _strip_brackets(fk.referenced_columns)
+    line = (
+        f'CONSTRAINT {_q(fk.fk_name)} FOREIGN KEY ({parent_cols}) '
+        f'REFERENCES {ref_table} ({referenced_cols})'
+    )
+
+    on_delete = _emit_fk_action(fk.on_delete)
+    if on_delete:
+        line += f' ON DELETE {on_delete}'
+
+    on_update = _emit_fk_action(fk.on_update)
+    if on_update:
+        line += f' ON UPDATE {on_update}'
+
+    return line
+
+
+def emit_table_ddl(
+    table: TableDef, *, foreign_keys: list[ForeignKeyDef] | None = None
+) -> str:
+    full_name = f'{_q(table.schema_name)}.{_q(table.table_name)}'
+
+    lines: list[str] = [f'create table {full_name} (']
+
+    body: list[str] = [f'    {_emit_column(col)}' for col in table.columns]
+
+    if table.primary_key and table.primary_key.key_columns:
+        pk_type = table.primary_key.type_desc.lower().replace('_', ' ')
+        pk_cols = _strip_brackets(table.primary_key.key_columns)
+        pk_cols = re.sub(r'\s+asc\b', '', pk_cols, flags=re.IGNORECASE)
+        body.append(
+            f'    CONSTRAINT {_q(table.primary_key.constraint_name)} PRIMARY KEY {pk_type} ({pk_cols})'
+        )
+
+    if foreign_keys:
+        for fk in sorted(foreign_keys, key=lambda x: x.fk_name.lower()):
+            body.append(f'    {_emit_inline_fk_constraint(fk)}')
+
+    lines.append(',\n'.join(body))
+    lines.append(')')
+
+    return '\n'.join(lines)

al_data_utils/backup_utils/mssql/schema_snapshot/introspect/__init__.py

@@ -0,0 +1,19 @@
+from .indexes import get_indexes
+from .foreign_keys import get_foreign_keys
+from .tables import get_tables
+from .views import get_views
+from .stored_procedures import get_stored_procedures
+from .functions import get_functions
+from .sequences import get_sequences
+from .synonyms import get_synonyms
+
+__all__ = [
+    'get_indexes',
+    'get_foreign_keys',
+    'get_tables',
+    'get_views',
+    'get_stored_procedures',
+    'get_functions',
+    'get_sequences',
+    'get_synonyms',
+]

al_data_utils/backup_utils/mssql/schema_snapshot/introspect/foreign_keys.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from al_data_utils.sql_utils.sql_utils import UnifiedSQLClient
+from ..models import ForeignKeyDef
+
+
+_SQL_FOREIGN_KEYS = r"""
+select
+    ps.name as parent_schema,
+    pt.name as parent_table,
+    fk.name as fk_name,
+    rs.name as referenced_schema,
+    rt.name as referenced_table,
+    fk.delete_referential_action_desc as on_delete,
+    fk.update_referential_action_desc as on_update,
+    cast(fk.is_not_for_replication as bit) as is_not_for_replication,
+    cast(fk.is_disabled as bit) as is_disabled,
+    cast(fk.is_not_trusted as bit) as is_not_trusted,
+
+    stuff((
+        select
+            ', ' + quotename(pc.name)
+        from sys.foreign_key_columns fkc
+        inner join sys.columns pc
+          on pc.object_id = fkc.parent_object_id
+         and pc.column_id = fkc.parent_column_id
+        where fkc.constraint_object_id = fk.object_id
+        order by fkc.constraint_column_id
+        for xml path(''), type
+    ).value('.', 'nvarchar(max)'), 1, 2, '') as parent_columns,
+
+    stuff((
+        select
+            ', ' + quotename(rc.name)
+        from sys.foreign_key_columns fkc
+        inner join sys.columns rc
+          on rc.object_id = fkc.referenced_object_id
+         and rc.column_id = fkc.referenced_column_id
+        where fkc.constraint_object_id = fk.object_id
+        order by fkc.constraint_column_id
+        for xml path(''), type
+    ).value('.', 'nvarchar(max)'), 1, 2, '') as referenced_columns
+
+from sys.foreign_keys fk
+inner join sys.tables pt on pt.object_id = fk.parent_object_id
+inner join sys.schemas ps on ps.schema_id = pt.schema_id
+inner join sys.tables rt on rt.object_id = fk.referenced_object_id
+inner join sys.schemas rs on rs.schema_id = rt.schema_id
+where pt.is_ms_shipped = 0
+order by ps.name, pt.name, fk.name;
+"""
+
+
+def _normalized_set(values: set[str] | None) -> set[str]:
+    if not values:
+        return set()
+    return {str(v).strip().lower() for v in values if str(v).strip()}
+
+
+def _split_qualified(values: set[str]) -> list[tuple[str, str]]:
+    pairs: list[tuple[str, str]] = []
+    for v in sorted(values):
+        if '.' not in v:
+            continue
+        schema, table = v.split('.', 1)
+        schema = schema.strip().lower()
+        table = table.strip().lower()
+        if schema and table:
+            pairs.append((schema, table))
+    return pairs
+
+
+def _build_filter_sql_and_params(
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_tables: set[str] | None = None,
+    exclude_tables: set[str] | None = None,
+) -> tuple[str, dict[str, str]]:
+    inc_schemas = sorted(_normalized_set(include_schemas))
+    exc_schemas = sorted(_normalized_set(exclude_schemas))
+    inc_tables = _split_qualified(_normalized_set(include_tables))
+    exc_tables = _split_qualified(_normalized_set(exclude_tables))
+
+    params: dict[str, str] = {}
+    clauses: list[str] = []
+
+    if inc_schemas:
+        names: list[str] = []
+        for i, schema in enumerate(inc_schemas):
+            key = f'inc_schema_{i}'
+            params[key] = schema
+            names.append(f':{key}')
+        clauses.append(f'lower(ps.name) in ({", ".join(names)})')
+
+    if exc_schemas:
+        names = []
+        for i, schema in enumerate(exc_schemas):
+            key = f'exc_schema_{i}'
+            params[key] = schema
+            names.append(f':{key}')
+        clauses.append(f'lower(ps.name) not in ({", ".join(names)})')
+
+    if inc_tables:
+        parts: list[str] = []
+        for i, (schema, table) in enumerate(inc_tables):
+            sk = f'inc_tbl_schema_{i}'
+            tk = f'inc_tbl_name_{i}'
+            params[sk] = schema
+            params[tk] = table
+            parts.append(f'(lower(ps.name) = :{sk} and lower(pt.name) = :{tk})')
+        clauses.append(f'({" or ".join(parts)})')
+
+    if exc_tables:
+        parts = []
+        for i, (schema, table) in enumerate(exc_tables):
+            sk = f'exc_tbl_schema_{i}'
+            tk = f'exc_tbl_name_{i}'
+            params[sk] = schema
+            params[tk] = table
+            parts.append(f'(lower(ps.name) = :{sk} and lower(pt.name) = :{tk})')
+        clauses.append(f'not ({" or ".join(parts)})')
+
+    if not clauses:
+        return '', {}
+
+    return '\n  and ' + '\n  and '.join(clauses), params
+
+
+def _inject_filters(sql: str, filters_sql: str) -> str:
+    if not filters_sql:
+        return sql
+    return sql.replace(
+        'where pt.is_ms_shipped = 0', f'where pt.is_ms_shipped = 0{filters_sql}'
+    )
+
+
+def get_foreign_keys(
+    sql_client: UnifiedSQLClient,
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_tables: set[str] | None = None,
+    exclude_tables: set[str] | None = None,
+) -> list[ForeignKeyDef]:
+    filters_sql, params = _build_filter_sql_and_params(
+        include_schemas=include_schemas,
+        exclude_schemas=exclude_schemas,
+        include_tables=include_tables,
+        exclude_tables=exclude_tables,
+    )
+    rows = sql_client.execute_command_with_result_to_dicts(
+        _inject_filters(_SQL_FOREIGN_KEYS, filters_sql), params
+    )
+    result: list[ForeignKeyDef] = []
+    for r in rows:
+        result.append(
+            ForeignKeyDef(
+                parent_schema=str(r['parent_schema']),
+                parent_table=str(r['parent_table']),
+                fk_name=str(r['fk_name']),
+                referenced_schema=str(r['referenced_schema']),
+                referenced_table=str(r['referenced_table']),
+                on_delete=str(r['on_delete']),
+                on_update=str(r['on_update']),
+                is_not_for_replication=bool(r['is_not_for_replication']),
+                is_disabled=bool(r['is_disabled']),
+                is_not_trusted=bool(r['is_not_trusted']),
+                parent_columns=str(r['parent_columns'] or ''),
+                referenced_columns=str(r['referenced_columns'] or ''),
+            )
+        )
+    return result