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

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

@@ -0,0 +1,313 @@
+from __future__ import annotations
+
+from al_data_utils.logging_utils import warning, info
+from al_data_utils.sql_utils.sql_utils import UnifiedSQLClient
+from ..models import SqlObjectDef
+
+
+_SQL_STORED_PROCEDURE_NAMES = r"""
+select
+    s.name as schema_name,
+    p.name as object_name
+from sys.procedures p
+inner join sys.schemas s
+  on s.schema_id = p.schema_id
+where p.is_ms_shipped = 0
+order by s.name, p.name;
+"""
+
+
+_SQL_STORED_PROCEDURES_BULK = r"""
+select
+        s.name as schema_name,
+        p.name as object_name,
+        convert(nvarchar(max), m.definition) as definition
+from sys.procedures p
+inner join sys.schemas s
+    on s.schema_id = p.schema_id
+inner join sys.sql_modules m
+    on m.object_id = p.object_id
+where p.is_ms_shipped = 0
+order by s.name, p.name;
+"""
+
+
+_SQL_STORED_PROCEDURE_DEFINITION = r"""
+select
+    convert(nvarchar(max), object_definition(object_id(quotename(?) + '.' + quotename(?)))) as definition
+"""
+
+
+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, name = v.split('.', 1)
+        schema = schema.strip().lower()
+        name = name.strip().lower()
+        if schema and name:
+            pairs.append((schema, name))
+    return pairs
+
+
+def _build_filter_sql_and_params(
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_objects: set[str] | None = None,
+    exclude_objects: set[str] | None = None,
+) -> tuple[str, dict[str, str]]:
+    inc_schemas = sorted(_normalized_set(include_schemas))
+    exc_schemas = sorted(_normalized_set(exclude_schemas))
+    inc_objects = _split_qualified(_normalized_set(include_objects))
+    exc_objects = _split_qualified(_normalized_set(exclude_objects))
+
+    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(s.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(s.name) not in ({", ".join(names)})')
+
+    if inc_objects:
+        parts: list[str] = []
+        for i, (schema, name) in enumerate(inc_objects):
+            sk = f'inc_obj_schema_{i}'
+            nk = f'inc_obj_name_{i}'
+            params[sk] = schema
+            params[nk] = name
+            parts.append(f'(lower(s.name) = :{sk} and lower(p.name) = :{nk})')
+        clauses.append(f'({" or ".join(parts)})')
+
+    if exc_objects:
+        parts = []
+        for i, (schema, name) in enumerate(exc_objects):
+            sk = f'exc_obj_schema_{i}'
+            nk = f'exc_obj_name_{i}'
+            params[sk] = schema
+            params[nk] = name
+            parts.append(f'(lower(s.name) = :{sk} and lower(p.name) = :{nk})')
+        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 p.is_ms_shipped = 0', f'where p.is_ms_shipped = 0{filters_sql}'
+    )
+
+
+def _execute_direct(
+    sql_client: UnifiedSQLClient, query: str, params: tuple | None = None
+) -> list[dict[str, object]]:
+    """Execute a query via raw DBAPI cursor and return rows as dicts.
+
+    This path avoids SQLAlchemy transaction state issues after low-level TDS errors.
+    """
+    conn = sql_client.engine.raw_connection()
+    cursor = conn.cursor()
+    try:
+        if params:
+            cursor.execute(query, params)
+        else:
+            cursor.execute(query)
+
+        columns = [c[0] for c in cursor.description or []]
+        rows = cursor.fetchall() if columns else []
+        return [dict(zip(columns, row)) for row in rows]
+    finally:
+        try:
+            cursor.close()
+        except Exception:
+            pass
+        try:
+            conn.close()
+        except Exception:
+            pass
+
+
+def get_stored_procedures(
+    sql_client: UnifiedSQLClient,
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_objects: set[str] | None = None,
+    exclude_objects: set[str] | None = None,
+    verbose: bool = True,
+) -> list[SqlObjectDef]:
+    _info = info if verbose else (lambda _: None)
+
+    filters_sql, params = _build_filter_sql_and_params(
+        include_schemas=include_schemas,
+        exclude_schemas=exclude_schemas,
+        include_objects=include_objects,
+        exclude_objects=exclude_objects,
+    )
+
+    include_pairs = _split_qualified(_normalized_set(include_objects))
+
+    # Exact include-list mode: fetch only requested procedure definitions.
+    if include_pairs:
+        total = len(include_pairs)
+        _info(
+            '[schema_snapshot..] stored procedures discovered: '
+            f'{total} (targeted include mode)'
+        )
+        result: list[SqlObjectDef] = []
+        for i, (schema_name, object_name) in enumerate(include_pairs, start=1):
+            _info(
+                f'[schema_snapshot..] fetching stored procedure definition [{i}/{total}]: '
+                f'{schema_name}.{object_name}'
+            )
+
+            try:
+                def_rows = _execute_direct(
+                    sql_client,
+                    _SQL_STORED_PROCEDURE_DEFINITION,
+                    params=(schema_name, object_name),
+                )
+                definition = ''
+                if def_rows:
+                    definition = str(def_rows[0].get('definition') or '')
+
+                if not definition:
+                    warning(
+                        '[schema_snapshot..] stored procedure definition is empty '
+                        f'for {schema_name}.{object_name} (not found or encrypted); skipping'
+                    )
+                    continue
+
+                result.append(
+                    SqlObjectDef(
+                        schema_name=schema_name,
+                        object_name=object_name,
+                        definition=definition,
+                    )
+                )
+            except Exception as one_err:
+                warning(
+                    '[schema_snapshot..] failed to extract stored procedure '
+                    f'[{i}/{total}] {schema_name}.{object_name}; skipping ({one_err})'
+                )
+
+        return result
+
+    # Fast path: single bulk query for all procedure definitions.
+    # Falls back to per-procedure extraction only if the bulk query fails.
+    try:
+        rows = sql_client.execute_command_with_result_to_dicts(
+            _inject_filters(_SQL_STORED_PROCEDURES_BULK, filters_sql), params
+        )
+        total = len(rows)
+        _info(f'[schema_snapshot..] stored procedures discovered: {total} (bulk mode)')
+
+        result: list[SqlObjectDef] = []
+        for i, r in enumerate(rows, start=1):
+            schema_name = str(r['schema_name'])
+            object_name = str(r['object_name'])
+            definition = str(r.get('definition') or '')
+
+            if not definition:
+                warning(
+                    '[schema_snapshot..] stored procedure definition is empty '
+                    f'for {schema_name}.{object_name} (possibly encrypted); skipping'
+                )
+                continue
+
+            result.append(
+                SqlObjectDef(
+                    schema_name=schema_name,
+                    object_name=object_name,
+                    definition=definition,
+                )
+            )
+
+            _info(
+                f'[schema_snapshot..] extracted stored procedure [{i}/{total}]: '
+                f'{schema_name}.{object_name}'
+            )
+
+        return result
+    except Exception as bulk_err:
+        warning(
+            '[schema_snapshot..] bulk stored procedure extraction failed; '
+            f'falling back to per-procedure mode ({bulk_err})'
+        )
+
+    # Fetch names first, then definition one by one.
+    # This avoids failing the whole snapshot because of one problematic procedure.
+    result: list[SqlObjectDef] = []
+    name_rows = sql_client.execute_command_with_result_to_dicts(
+        _inject_filters(_SQL_STORED_PROCEDURE_NAMES, filters_sql), params
+    )
+    total = len(name_rows)
+    _info(
+        f'[schema_snapshot..] stored procedures discovered: {total} (per-procedure mode)'
+    )
+    for i, r in enumerate(name_rows, start=1):
+        schema_name = str(r['schema_name'])
+        object_name = str(r['object_name'])
+        _info(
+            f'[schema_snapshot..] fetching stored procedure definition [{i}/{total}]: '
+            f'{schema_name}.{object_name}'
+        )
+        try:
+            def_rows = _execute_direct(
+                sql_client,
+                _SQL_STORED_PROCEDURE_DEFINITION,
+                params=(schema_name, object_name),
+            )
+            definition = ''
+            if def_rows:
+                definition = str(def_rows[0].get('definition') or '')
+
+            if not definition:
+                warning(
+                    '[schema_snapshot..] stored procedure definition is empty '
+                    f'for {schema_name}.{object_name} (possibly encrypted); skipping'
+                )
+                continue
+
+            result.append(
+                SqlObjectDef(
+                    schema_name=schema_name,
+                    object_name=object_name,
+                    definition=definition,
+                )
+            )
+            _info(
+                f'[schema_snapshot..] extracted stored procedure [{i}/{total}]: '
+                f'{schema_name}.{object_name}'
+            )
+        except Exception as one_err:
+            warning(
+                '[schema_snapshot..] failed to extract stored procedure '
+                f'[{i}/{total}] {schema_name}.{object_name}; skipping ({one_err})'
+            )
+
+    return result

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

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from al_data_utils.sql_utils.sql_utils import UnifiedSQLClient
+from ..models import SqlObjectDef
+
+
+_SQL_SYNONYMS = r"""
+select
+    s.name as schema_name,
+    sn.name as object_name,
+    sn.base_object_name
+from sys.synonyms sn
+inner join sys.schemas s
+  on s.schema_id = sn.schema_id
+order by s.name, sn.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, name = v.split('.', 1)
+        schema = schema.strip().lower()
+        name = name.strip().lower()
+        if schema and name:
+            pairs.append((schema, name))
+    return pairs
+
+
+def _build_filter_sql_and_params(
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_objects: set[str] | None = None,
+    exclude_objects: set[str] | None = None,
+) -> tuple[str, dict[str, str]]:
+    inc_schemas = sorted(_normalized_set(include_schemas))
+    exc_schemas = sorted(_normalized_set(exclude_schemas))
+    inc_objects = _split_qualified(_normalized_set(include_objects))
+    exc_objects = _split_qualified(_normalized_set(exclude_objects))
+
+    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(s.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(s.name) not in ({", ".join(names)})')
+
+    if inc_objects:
+        parts: list[str] = []
+        for i, (schema, name) in enumerate(inc_objects):
+            sk = f'inc_obj_schema_{i}'
+            nk = f'inc_obj_name_{i}'
+            params[sk] = schema
+            params[nk] = name
+            parts.append(f'(lower(s.name) = :{sk} and lower(sn.name) = :{nk})')
+        clauses.append(f'({" or ".join(parts)})')
+
+    if exc_objects:
+        parts = []
+        for i, (schema, name) in enumerate(exc_objects):
+            sk = f'exc_obj_schema_{i}'
+            nk = f'exc_obj_name_{i}'
+            params[sk] = schema
+            params[nk] = name
+            parts.append(f'(lower(s.name) = :{sk} and lower(sn.name) = :{nk})')
+        clauses.append(f'not ({" or ".join(parts)})')
+
+    if not clauses:
+        return '', {}
+
+    return '\nwhere ' + '\n  and '.join(clauses), params
+
+
+def _inject_filters(sql: str, filters_sql: str) -> str:
+    if not filters_sql:
+        return sql
+    return sql.replace(
+        'order by s.name, sn.name;', f'{filters_sql}\norder by s.name, sn.name;'
+    )
+
+
+def get_synonyms(
+    sql_client: UnifiedSQLClient,
+    *,
+    include_schemas: set[str] | None = None,
+    exclude_schemas: set[str] | None = None,
+    include_objects: set[str] | None = None,
+    exclude_objects: set[str] | None = None,
+) -> list[SqlObjectDef]:
+    filters_sql, params = _build_filter_sql_and_params(
+        include_schemas=include_schemas,
+        exclude_schemas=exclude_schemas,
+        include_objects=include_objects,
+        exclude_objects=exclude_objects,
+    )
+
+    rows = sql_client.execute_command_with_result_to_dicts(
+        _inject_filters(_SQL_SYNONYMS, filters_sql), params
+    )
+
+    result: list[SqlObjectDef] = []
+    for r in rows:
+        schema_name = str(r['schema_name'])
+        object_name = str(r['object_name'])
+        base_object_name = str(r.get('base_object_name') or '')
+        definition = (
+            f'create synonym {schema_name}.{object_name} for {base_object_name}'
+        )
+        result.append(
+            SqlObjectDef(
+                schema_name=schema_name, object_name=object_name, definition=definition
+            )
+        )
+    return result

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

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+from collections import defaultdict
+
+from al_data_utils.sql_utils.sql_utils import UnifiedSQLClient
+from ..models import TableColumnDef, TableDef, PrimaryKeyDef
+
+
+_SQL_TABLE_COLUMNS = r"""
+select
+    s.name as schema_name,
+    t.name as table_name,
+    c.name as column_name,
+    c.column_id as ordinal_position,
+    ty.name as data_type,
+    cast(c.max_length as int) as max_length,
+    cast(c.precision as int) as precision,
+    cast(c.scale as int) as scale,
+    cast(c.is_nullable as bit) as is_nullable,
+    cast(c.is_identity as bit) as is_identity,
+    cast(ic.seed_value as int) as identity_seed,
+    cast(ic.increment_value as int) as identity_increment,
+    cast(c.is_computed as bit) as is_computed,
+    cc.definition as computed_definition,
+    c.collation_name,
+    dc.name as default_constraint_name,
+    dc.definition as default_definition
+from sys.tables t
+inner join sys.schemas s
+  on s.schema_id = t.schema_id
+inner join sys.columns c
+  on c.object_id = t.object_id
+inner join sys.types ty
+  on ty.user_type_id = c.user_type_id
+left outer join sys.identity_columns ic
+  on ic.object_id = c.object_id
+ and ic.column_id = c.column_id
+left outer join sys.computed_columns cc
+  on cc.object_id = c.object_id
+ and cc.column_id = c.column_id
+left outer join sys.default_constraints dc
+  on dc.parent_object_id = c.object_id
+ and dc.parent_column_id = c.column_id
+where t.is_ms_shipped = 0
+order by s.name, t.name, c.column_id;
+"""
+
+
+_SQL_PRIMARY_KEYS = r"""
+select
+    s.name as schema_name,
+    t.name as table_name,
+    kc.name as constraint_name,
+    i.type_desc,
+    stuff((
+        select
+            ', ' + quotename(c.name) + case when ic2.is_descending_key = 1 then ' desc' else ' asc' end
+        from sys.index_columns ic2
+      inner join sys.columns c
+          on c.object_id = ic2.object_id and c.column_id = ic2.column_id
+        where ic2.object_id = i.object_id
+          and ic2.index_id = i.index_id
+          and ic2.key_ordinal > 0
+        order by ic2.key_ordinal
+        for xml path(''), type
+    ).value('.', 'nvarchar(max)'), 1, 2, '') as key_columns
+from sys.key_constraints kc
+inner join sys.tables t
+  on t.object_id = kc.parent_object_id
+inner join sys.schemas s
+  on s.schema_id = t.schema_id
+inner join sys.indexes i
+  on i.object_id = kc.parent_object_id
+ and i.index_id = kc.unique_index_id
+where kc.type = 'PK'
+order by s.name, t.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_table_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(s.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(s.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(s.name) = :{sk} and lower(t.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(s.name) = :{sk} and lower(t.name) = :{tk})')
+        clauses.append(f'not ({" or ".join(parts)})')
+
+    if not clauses:
+        return '', {}
+
+    return '\n  and ' + '\n  and '.join(clauses), params
+
+
+def _inject_table_filters(sql: str, filters_sql: str) -> str:
+    if not filters_sql:
+        return sql
+    if 'where t.is_ms_shipped = 0' in sql:
+        return sql.replace(
+            'where t.is_ms_shipped = 0', f'where t.is_ms_shipped = 0{filters_sql}'
+        )
+    if "where kc.type = 'PK'" in sql:
+        return sql.replace("where kc.type = 'PK'", f"where kc.type = 'PK'{filters_sql}")
+    return sql
+
+
+def get_tables(
+    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[TableDef]:
+    filters_sql, params = _build_table_filter_sql_and_params(
+        include_schemas=include_schemas,
+        exclude_schemas=exclude_schemas,
+        include_tables=include_tables,
+        exclude_tables=exclude_tables,
+    )
+
+    col_rows = sql_client.execute_command_with_result_to_dicts(
+        _inject_table_filters(_SQL_TABLE_COLUMNS, filters_sql), params
+    )
+    pk_rows = sql_client.execute_command_with_result_to_dicts(
+        _inject_table_filters(_SQL_PRIMARY_KEYS, filters_sql), params
+    )
+
+    columns_by_table: dict[tuple[str, str], list[TableColumnDef]] = defaultdict(list)
+    for r in col_rows:
+        col = TableColumnDef(
+            schema_name=str(r['schema_name']),
+            table_name=str(r['table_name']),
+            column_name=str(r['column_name']),
+            ordinal_position=int(r['ordinal_position']),
+            data_type=str(r['data_type']),
+            max_length=None if r['max_length'] is None else int(r['max_length']),
+            precision=None if r['precision'] is None else int(r['precision']),
+            scale=None if r['scale'] is None else int(r['scale']),
+            is_nullable=bool(r['is_nullable']),
+            is_identity=bool(r['is_identity']),
+            identity_seed=None
+            if r['identity_seed'] is None
+            else int(r['identity_seed']),
+            identity_increment=None
+            if r['identity_increment'] is None
+            else int(r['identity_increment']),
+            is_computed=bool(r['is_computed']),
+            computed_definition=None
+            if r['computed_definition'] is None
+            else str(r['computed_definition']),
+            collation_name=None
+            if r['collation_name'] is None
+            else str(r['collation_name']),
+            default_constraint_name=None
+            if r.get('default_constraint_name') is None
+            else str(r.get('default_constraint_name')),
+            default_definition=None
+            if r['default_definition'] is None
+            else str(r['default_definition']),
+        )
+        columns_by_table[(col.schema_name, col.table_name)].append(col)
+
+    pk_by_table: dict[tuple[str, str], PrimaryKeyDef] = {}
+    for r in pk_rows:
+        pk = PrimaryKeyDef(
+            schema_name=str(r['schema_name']),
+            table_name=str(r['table_name']),
+            constraint_name=str(r['constraint_name']),
+            key_columns=str(r['key_columns'] or ''),
+            type_desc=str(r['type_desc'] or ''),
+        )
+        pk_by_table[(pk.schema_name, pk.table_name)] = pk
+
+    result: list[TableDef] = []
+    for key in sorted(
+        columns_by_table.keys(), key=lambda x: (x[0].lower(), x[1].lower())
+    ):
+        schema_name, table_name = key
+        cols = sorted(columns_by_table[key], key=lambda c: c.ordinal_position)
+        result.append(
+            TableDef(
+                schema_name=schema_name,
+                table_name=table_name,
+                columns=cols,
+                primary_key=pk_by_table.get((schema_name, table_name)),
+            )
+        )
+
+    return result