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

al_data_utils/azure_as_utils/models.py

@@ -0,0 +1,265 @@
+"""Value objects describing a refresh request and its observed status.
+
+The AAS REST API speaks JSON with PascalCase keys on the way in (``Type``,
+``CommitMode``, ``Objects`` ...) and camelCase keys on the way out
+(``status``, ``startTime`` ...). The dataclasses here give both sides a typed
+shape while still accepting the plain dicts callers have always passed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+REFRESH_TYPES: frozenset[str] = frozenset(
+    {'full', 'clearvalues', 'calculate', 'dataonly', 'automatic', 'defragment'}
+)
+COMMIT_MODES: frozenset[str] = frozenset({'transactional', 'partialbatch'})
+
+#: Statuses reported while a refresh has not finished yet.
+ACTIVE_STATUSES: frozenset[str] = frozenset({'notstarted', 'inprogress'})
+SUCCEEDED_STATUS = 'succeeded'
+
+
+@dataclass(frozen=True)
+class RefreshObject:
+    """One unit of work inside a refresh: a whole table or a single partition."""
+
+    table: str
+    partition: str | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.table, str) or not self.table.strip():
+            raise ValueError('refresh object needs a non-empty "table" name')
+        if self.partition is not None and not isinstance(self.partition, str):
+            raise ValueError('"partition" must be a string when given')
+        if self.partition is not None and not self.partition.strip():
+            object.__setattr__(self, 'partition', None)
+
+    @classmethod
+    def coerce(cls, value: RefreshObject | Mapping[str, Any]) -> RefreshObject:
+        """Accept an existing instance or a ``{'table': .., 'partition': ..}`` dict."""
+        if isinstance(value, RefreshObject):
+            return value
+        if not isinstance(value, Mapping):
+            raise TypeError(f'refresh object must be a mapping, got {type(value)!r}')
+        lowered = {str(key).lower(): val for key, val in value.items()}
+        unknown = sorted(set(lowered) - {'table', 'partition'})
+        if unknown:
+            raise ValueError(f'unsupported refresh object keys: {unknown}')
+        if 'table' not in lowered:
+            raise ValueError(f'refresh object is missing "table": {dict(value)!r}')
+        return cls(table=lowered['table'], partition=lowered.get('partition'))
+
+    @property
+    def label(self) -> str:
+        """Short human-readable name used in log lines."""
+        if self.partition is None:
+            return self.table
+        return f'{self.table} / {self.partition}'
+
+    def table_only(self) -> RefreshObject:
+        """Same table with the partition dropped."""
+        return RefreshObject(self.table)
+
+    def to_payload(self) -> dict[str, str]:
+        payload = {'table': self.table}
+        if self.partition is not None:
+            payload['partition'] = self.partition
+        return payload
+
+
+_REQUEST_KEYS: dict[str, str] = {
+    'type': 'refresh_type',
+    'commitmode': 'commit_mode',
+    'maxparallelism': 'max_parallelism',
+    'retrycount': 'retry_count',
+    'objects': 'objects',
+}
+
+
+@dataclass(frozen=True)
+class RefreshRequest:
+    """Body of ``POST .../models/{model}/refreshes``.
+
+    ``objects`` empty means the whole model. Keys the class does not model
+    (for example ``ApplyRefreshPolicy``) survive in ``extra`` and are sent
+    untouched.
+    """
+
+    refresh_type: str = 'Full'
+    commit_mode: str = 'transactional'
+    max_parallelism: int = 10
+    retry_count: int = 0
+    objects: tuple[RefreshObject, ...] = ()
+    extra: Mapping[str, Any] = field(default_factory=dict, hash=False)
+
+    def __post_init__(self) -> None:
+        if (
+            not isinstance(self.refresh_type, str)
+            or self.refresh_type.lower() not in REFRESH_TYPES
+        ):
+            raise ValueError(
+                f'unsupported refresh Type {self.refresh_type!r}; '
+                f'expected one of {sorted(REFRESH_TYPES)}'
+            )
+        if (
+            not isinstance(self.commit_mode, str)
+            or self.commit_mode.lower() not in COMMIT_MODES
+        ):
+            raise ValueError(
+                f'unsupported CommitMode {self.commit_mode!r}; '
+                f'expected one of {sorted(COMMIT_MODES)}'
+            )
+        if isinstance(self.max_parallelism, bool) or not isinstance(
+            self.max_parallelism, int
+        ):
+            raise TypeError('MaxParallelism must be an integer')
+        if self.max_parallelism < 1:
+            raise ValueError('MaxParallelism must be at least 1')
+        if isinstance(self.retry_count, bool) or not isinstance(self.retry_count, int):
+            raise TypeError('RetryCount must be an integer')
+        if self.retry_count < 0:
+            raise ValueError('RetryCount cannot be negative')
+        if self.objects is None:
+            object.__setattr__(self, 'objects', ())
+        else:
+            object.__setattr__(
+                self, 'objects', tuple(RefreshObject.coerce(o) for o in self.objects)
+            )
+        object.__setattr__(self, 'extra', dict(self.extra or {}))
+
+    @classmethod
+    def coerce(cls, value: RefreshRequest | Mapping[str, Any]) -> RefreshRequest:
+        """Accept an existing instance or the raw AAS JSON body as a dict."""
+        if isinstance(value, RefreshRequest):
+            return value
+        if not isinstance(value, Mapping):
+            raise TypeError(f'refresh request must be a mapping, got {type(value)!r}')
+        kwargs: dict[str, Any] = {}
+        extra: dict[str, Any] = {}
+        for key, val in value.items():
+            target = _REQUEST_KEYS.get(str(key).lower())
+            if target is None:
+                extra[str(key)] = val
+            else:
+                kwargs[target] = val
+        return cls(extra=extra, **kwargs)
+
+    @property
+    def is_whole_model(self) -> bool:
+        return not self.objects
+
+    def with_objects(self, objects: Iterable[RefreshObject]) -> RefreshRequest:
+        return replace(self, objects=tuple(objects))
+
+    def without_partitions(self) -> RefreshRequest:
+        """Collapse partition-level objects to their tables, keeping first-seen order."""
+        seen: dict[str, RefreshObject] = {}
+        for obj in self.objects:
+            seen.setdefault(obj.table, obj.table_only())
+        return self.with_objects(seen.values())
+
+    def per_object(self) -> list[RefreshRequest]:
+        """One request per object; a whole-model request yields itself."""
+        if self.is_whole_model:
+            return [self]
+        return [self.with_objects((obj,)) for obj in self.objects]
+
+    def describe(self) -> str:
+        """Short summary used in log lines."""
+        if self.is_whole_model:
+            return 'whole model'
+        return ', '.join(obj.label for obj in self.objects)
+
+    def to_payload(self) -> dict[str, Any]:
+        payload: dict[str, Any] = dict(self.extra)
+        payload['Type'] = self.refresh_type
+        payload['CommitMode'] = self.commit_mode
+        payload['MaxParallelism'] = self.max_parallelism
+        payload['RetryCount'] = self.retry_count
+        if self.objects:
+            payload['Objects'] = [obj.to_payload() for obj in self.objects]
+        return payload
+
+
+def _flatten_messages(raw: Any, prefix: str = '') -> list[str]:
+    """Turn the ``messages`` array from a status document into plain strings."""
+    out: list[str] = []
+    if not isinstance(raw, list):
+        return out
+    for item in raw:
+        if isinstance(item, Mapping):
+            text = item.get('message') or item.get('Message') or str(dict(item))
+        else:
+            text = str(item)
+        out.append(f'{prefix}{text}')
+    return out
+
+
+@dataclass(frozen=True)
+class RefreshStatus:
+    """Snapshot of ``GET .../refreshes/{refreshId}``."""
+
+    refresh_id: str
+    status: str
+    refresh_type: str | None = None
+    start_time: str | None = None
+    end_time: str | None = None
+    objects: tuple[Mapping[str, Any], ...] = ()
+    messages: tuple[str, ...] = ()
+    raw: Mapping[str, Any] = field(default_factory=dict, hash=False)
+
+    @classmethod
+    def from_response(
+        cls, refresh_id: str, payload: Mapping[str, Any]
+    ) -> RefreshStatus:
+        status = payload.get('status')
+        if not isinstance(status, str) or not status:
+            raise ValueError(f'status document has no "status" field: {payload!r}')
+        objects_raw = payload.get('objects')
+        objects: tuple[Mapping[str, Any], ...] = ()
+        if isinstance(objects_raw, list):
+            objects = tuple(o for o in objects_raw if isinstance(o, Mapping))
+        messages = _flatten_messages(payload.get('messages'))
+        for obj in objects:
+            label = _object_label(obj)
+            messages.extend(_flatten_messages(obj.get('messages'), f'[{label}] '))
+        return cls(
+            refresh_id=refresh_id,
+            status=status,
+            refresh_type=payload.get('type') or payload.get('currentRefreshType'),
+            start_time=payload.get('startTime'),
+            end_time=payload.get('endTime'),
+            objects=objects,
+            messages=tuple(messages),
+            raw=dict(payload),
+        )
+
+    @property
+    def is_active(self) -> bool:
+        return self.status.lower() in ACTIVE_STATUSES
+
+    @property
+    def is_terminal(self) -> bool:
+        return not self.is_active
+
+    @property
+    def succeeded(self) -> bool:
+        return self.status.lower() == SUCCEEDED_STATUS
+
+    def failed_objects(self) -> list[str]:
+        """Labels of objects whose own status is not succeeded."""
+        out: list[str] = []
+        for obj in self.objects:
+            obj_status = str(obj.get('status', '')).lower()
+            if obj_status and obj_status != SUCCEEDED_STATUS:
+                out.append(f'{_object_label(obj)}: {obj.get("status")}')
+        return out
+
+
+def _object_label(obj: Mapping[str, Any]) -> str:
+    table = obj.get('table') or '?'
+    partition = obj.get('partition')
+    return f'{table} / {partition}' if partition else str(table)

al_data_utils/azure_as_utils/poller.py

@@ -0,0 +1,121 @@
+"""Waits for a started refresh to reach a terminal state."""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+
+from .client import AasRestClient
+from .errors import AasHttpError, AasRefreshTimeout, AasTransportError
+from .models import RefreshStatus
+
+log = logging.getLogger(__name__)
+
+DEFAULT_POLL_INTERVAL_SECONDS = 15.0
+DEFAULT_TIMEOUT_SECONDS = 6 * 60 * 60.0
+
+
+class RefreshPoller:
+    """Polls the status endpoint with a growing interval until the refresh ends.
+
+    The interval starts at ``poll_interval_seconds`` and is multiplied by
+    ``backoff_factor`` after each check, capped at ``max_poll_interval_seconds``
+    (eight times the initial interval unless given). Up to
+    ``max_consecutive_errors`` transient failures in a row (HTTP 429 / 5xx or
+    connection problems) are tolerated before the error is raised.
+    """
+
+    def __init__(
+        self,
+        client: AasRestClient,
+        *,
+        poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
+        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+        backoff_factor: float = 1.5,
+        max_poll_interval_seconds: float | None = None,
+        max_consecutive_errors: int = 3,
+        sleep: Callable[[float], None] = time.sleep,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if poll_interval_seconds <= 0:
+            raise ValueError('poll_interval_seconds must be positive')
+        if timeout_seconds <= 0:
+            raise ValueError('timeout_seconds must be positive')
+        if backoff_factor < 1:
+            raise ValueError('backoff_factor must be >= 1')
+        if max_consecutive_errors < 0:
+            raise ValueError('max_consecutive_errors cannot be negative')
+        self._client = client
+        self._interval = float(poll_interval_seconds)
+        self._timeout = float(timeout_seconds)
+        self._backoff = float(backoff_factor)
+        self._max_interval = (
+            float(max_poll_interval_seconds)
+            if max_poll_interval_seconds is not None
+            else self._interval * 8
+        )
+        self._max_interval = max(self._max_interval, self._interval)
+        self._max_errors = max_consecutive_errors
+        self._sleep = sleep
+        self._clock = clock
+
+    @property
+    def timeout_seconds(self) -> float:
+        return self._timeout
+
+    def wait(self, model_name: str, refresh_id: str) -> RefreshStatus:
+        """Block until the refresh is no longer active and return its last status.
+
+        Raises :class:`AasRefreshTimeout` when the deadline passes first. The
+        returned status may be a failure; the caller decides what to do.
+        """
+        started = self._clock()
+        deadline = started + self._timeout
+        delay = self._interval
+        errors_in_a_row = 0
+        last_status = 'unknown'
+
+        while True:
+            try:
+                status = self._client.get_refresh(model_name, refresh_id)
+            except (AasHttpError, AasTransportError) as exc:
+                transient = isinstance(exc, AasTransportError) or exc.is_transient
+                errors_in_a_row += 1
+                if not transient or errors_in_a_row > self._max_errors:
+                    raise
+                log.warning(
+                    'model %s refresh %s: status check failed (%d/%d), retrying: %s',
+                    model_name,
+                    refresh_id,
+                    errors_in_a_row,
+                    self._max_errors,
+                    exc,
+                )
+            else:
+                errors_in_a_row = 0
+                last_status = status.status
+                if status.is_terminal:
+                    return status
+                log.info(
+                    'model %s refresh %s: %s (elapsed %s)',
+                    model_name,
+                    refresh_id,
+                    status.status,
+                    _format_elapsed(self._clock() - started),
+                )
+
+            remaining = deadline - self._clock()
+            if remaining <= 0:
+                raise AasRefreshTimeout(
+                    model_name, refresh_id, self._timeout, last_status
+                )
+            self._sleep(min(delay, remaining))
+            delay = min(delay * self._backoff, self._max_interval)
+
+
+def _format_elapsed(seconds: float) -> str:
+    total = max(int(seconds), 0)
+    hours, rest = divmod(total, 3600)
+    minutes, secs = divmod(rest, 60)
+    return f'{hours:d}:{minutes:02d}:{secs:02d}'

al_data_utils/backup_utils/__init__.py

@@ -0,0 +1,9 @@
+"""MSSQL-related utilities under backup_utils.
+
+This package is intentionally modular:
+- schema_snapshot: DataGrip-like deterministic schema snapshots (per-table files)
+- dacpac: SQLPackage/DACPAC helpers (planned / migrated from legacy mssql.py)
+- scripter: alternative scripting approach (planned)
+"""
+
+# Re-exports can be added later if you want a shorter import path.

al_data_utils/backup_utils/mssql/dacpac/__init__.py

@@ -0,0 +1 @@
+# placeholder for DACPAC/sqlpackage wrappers (migrate legacy mssql.py here)

al_data_utils/backup_utils/mssql/dacpac/sqlpackage.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Sequence
+
+from al_data_utils import logging_utils
+
+
+class SqlPackageError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class SqlPackageResult:
+    cmd: list[str]
+    returncode: int
+    stdout: str
+    stderr: str
+
+
+def _run(
+    cmd: Sequence[str], cwd: Optional[str | os.PathLike[str]] = None
+) -> SqlPackageResult:
+    cmd_list = [str(x) for x in cmd]
+    logging_utils.info(
+        f'[sqlpackage] run: {" ".join(cmd_list)} (cwd={cwd or os.getcwd()})'
+    )
+
+    p = subprocess.run(
+        cmd_list,
+        cwd=str(cwd) if cwd else None,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    out = (p.stdout or '').strip()
+    err = (p.stderr or '').strip()
+
+    if out:
+        logging_utils.info(f'[sqlpackage] stdout: {out}')
+    if err:
+        logging_utils.info(f'[sqlpackage] stderr: {err}')
+
+    if p.returncode != 0:
+        msg = err or out or 'sqlpackage failed'
+        logging_utils.error(f'[sqlpackage] fail: {msg}')
+        raise SqlPackageError(msg)
+
+    return SqlPackageResult(
+        cmd=cmd_list, returncode=p.returncode, stdout=out, stderr=err
+    )
+
+
+def _sqlpackage_path(sqlpackage_path: Optional[str] = None) -> str:
+    return sqlpackage_path or os.environ.get('SQLPACKAGE', 'sqlpackage')
+
+
+def _normalize_server(server: str) -> str:
+    s = server.strip()
+    if ':' in s and s.count(':') == 1 and ',' not in s:
+        host, port = s.split(':')
+        if port.isdigit():
+            return f'{host},{port}'
+    return s
+
+
+def extract_dacpac(
+    *,
+    source_connection_string: Optional[str] = None,
+    source_server: Optional[str] = None,
+    source_database: Optional[str] = None,
+    target_dacpac: str | os.PathLike[str],
+    sqlpackage_path: Optional[str] = None,
+    access_token: Optional[str] = None,
+    verify_extraction: bool = False,
+) -> SqlPackageResult:
+    sp = _sqlpackage_path(sqlpackage_path)
+    target = str(Path(target_dacpac))
+
+    cmd = [sp, '/Action:Extract', f'/TargetFile:{target}']
+
+    if source_connection_string:
+        cmd.append(f'/SourceConnectionString:{source_connection_string}')
+    else:
+        if not source_server or not source_database:
+            raise ValueError(
+                'Provide either source_connection_string OR (source_server and source_database).'
+            )
+        cmd.append(f'/SourceServerName:{_normalize_server(source_server)}')
+        cmd.append(f'/SourceDatabaseName:{source_database}')
+
+    if access_token:
+        cmd.append(f'/AccessToken:{access_token}')
+
+    if verify_extraction:
+        cmd.append('/p:VerifyExtraction=true')
+
+    return _run(cmd)
+
+
+def script_schema_from_dacpac(
+    *,
+    source_dacpac: str | os.PathLike[str],
+    empty_target_dacpac: str | os.PathLike[str],
+    output_sql: str | os.PathLike[str],
+    target_database_name: str,
+    sqlpackage_path: Optional[str] = None,
+    block_on_possible_data_loss: bool = True,
+    drop_objects_not_in_source: bool = False,
+) -> SqlPackageResult:
+    sp = _sqlpackage_path(sqlpackage_path)
+    src = str(Path(source_dacpac))
+    tgt = str(Path(empty_target_dacpac))
+    out = str(Path(output_sql))
+
+    Path(out).parent.mkdir(parents=True, exist_ok=True)
+
+    cmd = [
+        sp,
+        '/Action:Script',
+        f'/SourceFile:{src}',
+        f'/TargetFile:{tgt}',
+        f'/TargetDatabaseName:{target_database_name}',
+        f'/OutputPath:{out}',
+        f'/p:BlockOnPossibleDataLoss={"true" if block_on_possible_data_loss else "false"}',
+        f'/p:DropObjectsNotInSource={"true" if drop_objects_not_in_source else "false"}',
+    ]
+    return _run(cmd)
+
+
+def snapshot_schema(
+    *,
+    out_dir: str | os.PathLike[str],
+    empty_target_dacpac: str | os.PathLike[str],
+    target_database_name: str,
+    # source:
+    source_connection_string: Optional[str] = None,
+    source_server: Optional[str] = None,
+    source_database: Optional[str] = None,
+    access_token: Optional[str] = None,
+    sqlpackage_path: Optional[str] = None,
+    # behavior:
+    store_dacpac: bool = False,
+    cleanup_temp: bool = True,
+) -> dict[str, str]:
+    """
+    Default behavior:
+      - writes ONLY schema.sql to out_dir
+      - uses a temp dacpac internally
+
+    If store_dacpac=True:
+      - also writes db.dacpac to out_dir
+
+    Returns:
+      {'schema_sql': ..., 'dacpac': ... (only if stored)}
+    """
+    out_dir_p = Path(out_dir)
+    out_dir_p.mkdir(parents=True, exist_ok=True)
+
+    sql_path = out_dir_p / 'schema.sql'
+
+    # Choose where dacpac lives
+    if store_dacpac:
+        dacpac_path = out_dir_p / 'db.dacpac'
+        temp_dacpac_path = dacpac_path
+    else:
+        # temp file (still deterministic output because only schema.sql is committed)
+        tmp_dir = out_dir_p / '.tmp'
+        tmp_dir.mkdir(parents=True, exist_ok=True)
+        temp_dacpac_path = tmp_dir / 'db.dacpac'
+
+    extract_dacpac(
+        source_connection_string=source_connection_string,
+        source_server=source_server,
+        source_database=source_database,
+        target_dacpac=temp_dacpac_path,
+        sqlpackage_path=sqlpackage_path,
+        access_token=access_token,
+    )
+
+    script_schema_from_dacpac(
+        source_dacpac=temp_dacpac_path,
+        empty_target_dacpac=empty_target_dacpac,
+        output_sql=sql_path,
+        target_database_name=target_database_name,
+        sqlpackage_path=sqlpackage_path,
+    )
+
+    # cleanup temp dacpac if not stored
+    if not store_dacpac and cleanup_temp:
+        try:
+            Path(temp_dacpac_path).unlink(
+                missing_ok=True
+            )  # py3.8+ supports missing_ok on Windows too
+        except Exception as e:
+            logging_utils.info(
+                f'[sqlpackage] temp cleanup skipped: {temp_dacpac_path} ({e})'
+            )
+
+    result = {'schema_sql': str(sql_path)}
+    if store_dacpac:
+        result['dacpac'] = str(out_dir_p / 'db.dacpac')
+    return result

al_data_utils/backup_utils/mssql/schema_snapshot/__init__.py

@@ -0,0 +1,8 @@
+from .config import SnapshotConfig
+from .snapshot_processor import snapshot_schema, snapshot_schema_from_params
+
+__all__ = [
+    "SnapshotConfig",
+    "snapshot_schema",
+    "snapshot_schema_from_params",
+]