datamasque-python 1.0.0__py3-none-any.whl

datamasque/client/models/runs.py

@@ -0,0 +1,165 @@
+"""Typed request and response shapes for run-related API endpoints."""
+
+import enum
+from datetime import datetime
+from typing import Any, NewType, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+from datamasque.client.models.connection import ConnectionConfig, ConnectionId, unwrap_connection_id
+from datamasque.client.models.ruleset import Ruleset, RulesetId, unwrap_ruleset_id
+from datamasque.client.models.status import MaskingRunStatus
+
+RunId = NewType("RunId", int)
+
+
+class MaskType(enum.Enum):
+    """Type of a masking run."""
+
+    database = "database"  # Also used for schema discovery.
+    file = "file"
+    file_data_discovery = "file_data_discovery"
+
+
+class MaskingRunOptions(BaseModel):
+    """
+    Optional run-time overrides for `MaskingRunRequest.options`.
+
+    All fields optional; server applies defaults when omitted.
+    `run_secret`,
+    if supplied,
+    must be 16–256 characters and is used as the per-run encryption key;
+    the server auto-generates one when omitted.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    batch_size: Optional[int] = None
+    dry_run: Optional[bool] = None
+    continue_on_failure: Optional[bool] = None
+    max_rows: Optional[int] = None
+    diagnostic_logging: Optional[bool] = None
+    run_secret: Optional[str] = Field(default=None, min_length=16, max_length=256)
+    disable_instance_secret: Optional[bool] = None
+
+
+class MaskingRunRequest(BaseModel):
+    """
+    Request body for `POST /api/runs/`.
+
+    `connection`, `destination_connection`, and `ruleset` accept either the server-assigned ID
+    or the corresponding object returned by an earlier client call (e.g. a `ConnectionConfig`
+    or `Ruleset`); the object's `id` is extracted at construction time.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    connection: Union[ConnectionId, ConnectionConfig]
+    ruleset: Union[RulesetId, Ruleset]
+    mask_type: MaskType = MaskType.database
+    destination_connection: Optional[Union[ConnectionId, ConnectionConfig]] = None
+    options: MaskingRunOptions = Field(default_factory=MaskingRunOptions)
+    name: Optional[str] = None
+
+    @field_validator("connection", "destination_connection", mode="before")
+    @classmethod
+    def _unwrap_connection(cls, value: Any) -> Any:
+        return unwrap_connection_id(value)
+
+    @field_validator("ruleset", mode="before")
+    @classmethod
+    def _unwrap_ruleset(cls, value: Any) -> Any:
+        return unwrap_ruleset_id(value)
+
+
+class RunConnectionRef(BaseModel):
+    """A reference to a connection used in a run — just the ID and display name."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: Optional[ConnectionId] = None
+    name: str
+
+
+def _collapse_flat_connection_fields(data: Any) -> Any:
+    """
+    Collapse flat `*_connection` + `*_connection_name` pairs into nested `RunConnectionRef`s.
+
+    The admin server sends connections as two parallel fields
+    (`source_connection` holding the ID and `source_connection_name` holding the display name);
+    the client surfaces them as a single nested object.
+    Leaves the input alone if the fields are already in nested form
+    (i.e. the caller constructed the model directly).
+    """
+
+    if not isinstance(data, dict):
+        return data
+
+    data = dict(data)
+
+    if "source_connection_name" in data and not isinstance(data.get("source_connection"), dict):
+        data["source_connection"] = {
+            "id": data.pop("source_connection", None),
+            "name": data.pop("source_connection_name"),
+        }
+
+    dest_name = data.get("destination_connection_name")
+    if dest_name and not isinstance(data.get("destination_connection"), dict):
+        data["destination_connection"] = {
+            "id": data.pop("destination_connection", None),
+            "name": data.pop("destination_connection_name"),
+        }
+    elif "destination_connection_name" in data:
+        # Empty string or None — let the Optional default apply.
+        data.pop("destination_connection_name", None)
+        data.pop("destination_connection", None)
+
+    return data
+
+
+class RunInfo(BaseModel):
+    """Full record for a masking run."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: int
+    status: MaskingRunStatus
+    mask_type: MaskType
+    source_connection: RunConnectionRef
+    ruleset_name: str
+    name: Optional[str] = None
+    destination_connection: Optional[RunConnectionRef] = None
+    ruleset: Optional[RulesetId] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    options: Optional[dict[str, Any]] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def _collapse_connection_fields(cls, data: Any) -> Any:
+        return _collapse_flat_connection_fields(data)
+
+
+class UnfinishedRun(BaseModel):
+    """Represents a masking run that is queued, running, validating, or cancelling."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: int
+    source_connection: RunConnectionRef
+    ruleset_name: str
+    status: MaskingRunStatus
+    destination_connection: Optional[RunConnectionRef] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def _collapse_connection_fields(cls, data: Any) -> Any:
+        return _collapse_flat_connection_fields(data)
+
+    def __str__(self) -> str:
+        if self.destination_connection is not None:
+            connection_part = f'"{self.source_connection.name}", "{self.destination_connection.name}"'
+        else:
+            connection_part = f'"{self.source_connection.name}"'
+
+        return f'{connection_part}: Run ID {self.id} in status `{self.status.value}`, ruleset "{self.ruleset_name}"'

datamasque/client/models/status.py

@@ -0,0 +1,68 @@
+import enum
+
+
+class ValidationStatus(enum.Enum):
+    """Validation status of a ruleset or ruleset library."""
+
+    valid = "valid"
+    invalid = "invalid"
+    in_progress = "in_progress"
+    unknown = "unknown"
+
+
+class MaskingRunStatus(enum.Enum):
+    """List of valid masking run statuses."""
+
+    finished = "finished"
+    finished_with_warnings = "finished_with_warnings"
+    queued = "queued"
+    running = "running"
+    failed = "failed"
+    validating = "validating"
+    cancelling = "cancelling"
+    cancelled = "cancelled"
+
+    @classmethod
+    def get_final_states(cls) -> set["MaskingRunStatus"]:
+        """Returns the list of final statuses, i.e. the run is completed, successfully or otherwise."""
+
+        return {cls.finished, cls.finished_with_warnings, cls.cancelled, cls.failed}
+
+    @classmethod
+    def get_finished_states(cls) -> set["MaskingRunStatus"]:
+        """Returns the list of statuses that indicate the run completed successfully."""
+
+        return {cls.finished, cls.finished_with_warnings}
+
+    @property
+    def is_in_final_state(self) -> bool:
+        """Returns True if this status is a final status."""
+
+        return self in self.get_final_states()
+
+    @property
+    def is_finished(self) -> bool:
+        """Returns True if this status is a finished status."""
+
+        return self in self.get_finished_states()
+
+
+class AsyncRulesetGenerationTaskStatus(enum.Enum):
+    """List of statuses of async ruleset generation tasks."""
+
+    finished = "finished"
+    failed = "failed"
+    running = "running"
+    queued = "queued"
+
+    @classmethod
+    def get_final_states(cls) -> set["AsyncRulesetGenerationTaskStatus"]:
+        """Returns the list of final statuses, i.e. the ruleset generation has completed, successfully or otherwise."""
+
+        return {cls.finished, cls.failed}
+
+    @property
+    def is_in_final_state(self) -> bool:
+        """Returns True if this status is a final status."""
+
+        return self in self.get_final_states()

datamasque/client/models/user.py

@@ -0,0 +1,69 @@
+import secrets
+import string
+from enum import Enum
+from typing import NewType, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+UserId = NewType("UserId", int)
+
+GENERATED_PASSWORD_LENGTH = 16
+
+
+class UserRole(Enum):
+    """
+    List of supported user roles.
+
+    `ruleset_library_manager` can be optionally included alongside `mask_builder`.
+    It is not valid as a standalone role.
+    """
+
+    superuser = "admin"
+    mask_builder = "mask_builder"
+    ruleset_library_manager = "ruleset_library_managers"
+    mask_runner = "mask_runner"
+
+
+class User(BaseModel):
+    """Represents a DataMasque user account."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    username: str
+    email: str
+    roles: list[UserRole] = Field(alias="user_roles")
+    id: Optional[UserId] = None
+    password: Optional[str] = Field(default=None, exclude=True)
+
+    @staticmethod
+    def generate_password() -> str:
+        """
+        Generates a password suitable for DataMasque authentication.
+
+        The password consists of 16 characters
+        without the same character occurring three times in a row
+        and without any three consecutive characters forming an increasing or decreasing sequence.
+        """
+
+        def is_sequential(s: str) -> bool:
+            """Check if the last three characters are in an increasing or decreasing sequence."""
+
+            if len(s) < 3:
+                return False
+            return (ord(s[-1]) == ord(s[-2]) + 1 == ord(s[-3]) + 2) or (ord(s[-1]) == ord(s[-2]) - 1 == ord(s[-3]) - 2)
+
+        chars = string.ascii_letters + string.digits
+        result = secrets.choice(chars)
+
+        while len(result) < GENERATED_PASSWORD_LENGTH:
+            next_char = secrets.choice(chars)
+            if len(result) >= 2 and next_char == result[-1] == result[-2]:
+                continue
+            if is_sequential(result + next_char):
+                continue
+            result += next_char
+
+        return result
+
+    def __str__(self) -> str:
+        return self.username

datamasque/client/ruleset_libraries.py

@@ -0,0 +1,164 @@
+import logging
+from typing import Iterator, Optional
+
+from datamasque.client.base import BaseClient
+from datamasque.client.exceptions import DataMasqueApiError, DataMasqueException
+from datamasque.client.models.pagination import Page
+from datamasque.client.models.ruleset import Ruleset
+from datamasque.client.models.ruleset_library import RulesetLibrary, RulesetLibraryId
+
+logger = logging.getLogger(__name__)
+
+
+class RulesetLibraryClient(BaseClient):
+    """Ruleset library CRUD API methods. Mixed into `DataMasqueClient`."""
+
+    def iter_ruleset_libraries(self) -> Iterator[RulesetLibrary]:
+        """Lazily iterate all ruleset libraries via paginated endpoint."""
+
+        return self._iter_paginated("/api/ruleset-libraries/", model=RulesetLibrary)
+
+    def list_ruleset_libraries(self) -> list[RulesetLibrary]:
+        """
+        Lists all ruleset libraries.
+
+        Note: The YAML content is not included in the list response for performance.
+        Use `get_ruleset_library` to retrieve the full library with YAML content.
+        """
+
+        return list(self.iter_ruleset_libraries())
+
+    def get_ruleset_library(self, library_id: RulesetLibraryId) -> RulesetLibrary:
+        """Retrieves a single ruleset library by ID, including its YAML content."""
+
+        response = self.make_request("GET", f"/api/ruleset-libraries/{library_id}/")
+        return RulesetLibrary.model_validate(response.json())
+
+    def get_ruleset_library_by_name(self, name: str, namespace: str = "") -> Optional[RulesetLibrary]:
+        """
+        Looks for a ruleset library matching the given name and namespace (case-sensitive, exact match).
+
+        Returns it (with full YAML content) if found, otherwise None.
+        """
+
+        response = self.make_request(
+            "GET",
+            "/api/ruleset-libraries/",
+            params={"name_exact": name, "namespace_exact": namespace, "limit": 1},
+        )
+        page = Page[RulesetLibrary].model_validate(response.json())
+        if not page.results:
+            return None
+
+        library_id = page.results[0].id
+        if library_id is None:
+            raise DataMasqueApiError(
+                "Server returned a ruleset library list entry without an `id`.",
+                response=response,
+            )
+
+        return self.get_ruleset_library(library_id)
+
+    def create_ruleset_library(self, library: RulesetLibrary) -> RulesetLibrary:
+        """
+        Creates a new ruleset library on the server.
+
+        Sets the library's server-assigned fields (`id`, `is_valid`, `created`, `modified`) and returns the library.
+        """
+
+        data = library.model_dump(exclude_none=True, by_alias=True, mode="json")
+        response = self.make_request("POST", "/api/ruleset-libraries/", data=data)
+        created_library = RulesetLibrary.model_validate(response.json())
+        library.id = created_library.id
+        library.is_valid = created_library.is_valid
+        library.created = created_library.created
+        library.modified = created_library.modified
+        logger.info('Creation of ruleset library "%s" successful', library.name)
+        return library
+
+    def update_ruleset_library(self, library: RulesetLibrary) -> RulesetLibrary:
+        """
+        Performs a full update of the ruleset library.
+
+        The library must have its `id` set (i.e., it must have been previously created or retrieved from the server).
+        """
+
+        if library.id is None:
+            raise ValueError("Cannot update a library that has not been created yet (id is None)")
+
+        data = library.model_dump(exclude_none=True, by_alias=True, mode="json")
+        response = self.make_request("PUT", f"/api/ruleset-libraries/{library.id}/", data=data)
+        updated_library = RulesetLibrary.model_validate(response.json())
+        library.is_valid = updated_library.is_valid
+        library.modified = updated_library.modified
+        logger.debug('Update of ruleset library "%s" successful', library.name)
+        return library
+
+    def create_or_update_ruleset_library(self, library: RulesetLibrary) -> RulesetLibrary:
+        """
+        Creates the library if it doesn't exist, or updates it if a library with the same name already exists.
+
+        Sets the library's `id` property.
+        """
+
+        existing = self.get_ruleset_library_by_name(library.name, library.namespace)
+        if existing is not None:
+            library.id = existing.id
+            return self.update_ruleset_library(library)
+
+        return self.create_ruleset_library(library)
+
+    def delete_ruleset_library_by_id_if_exists(self, library_id: RulesetLibraryId, *, force: bool = False) -> None:
+        """
+        Deletes (archives) the ruleset library with the given ID.
+
+        No-op if the library does not exist.
+
+        If the library is imported by any rulesets,
+        the server will return 409 Conflict unless `force=True` is passed.
+        """
+
+        params = {"force": "true"} if force else None
+        self._delete_if_exists(f"/api/ruleset-libraries/{library_id}/", params=params)
+
+    def delete_ruleset_library_by_name_if_exists(
+        self, library_name: str, namespace: str = "", *, force: bool = False
+    ) -> None:
+        """
+        Deletes the ruleset library with the given name and namespace.
+
+        No-op if the library does not exist.
+        """
+
+        all_libraries = self.list_ruleset_libraries()
+        matching = [lib for lib in all_libraries if lib.name == library_name and lib.namespace == namespace]
+        for lib in matching:
+            if lib.id is None:
+                raise DataMasqueException(f'Server returned a ruleset library named "{lib.name}" without an `id`.')
+
+            self.delete_ruleset_library_by_id_if_exists(lib.id, force=force)
+
+    def iter_rulesets_using_library(self, library_id: RulesetLibraryId) -> Iterator[Ruleset]:
+        """Lazily iterate non-archived rulesets that import the given library."""
+
+        return self._iter_paginated(f"/api/ruleset-libraries/{library_id}/rulesets/", model=Ruleset)
+
+    def list_rulesets_using_library(self, library_id: RulesetLibraryId) -> list[Ruleset]:
+        """
+        Lists non-archived rulesets that import the given library.
+
+        Note: The YAML content is not included in the response for performance.
+        Each returned Ruleset will have an empty string for `yaml`.
+        """
+
+        return list(self.iter_rulesets_using_library(library_id))
+
+    def validate_ruleset_library(self, library_id: RulesetLibraryId) -> RulesetLibrary:
+        """
+        Triggers re-validation of the ruleset library by performing a no-op update.
+
+        Returns the updated library with the new validation status.
+        """
+
+        response = self.make_request("PATCH", f"/api/ruleset-libraries/{library_id}/", data={})
+        return RulesetLibrary.model_validate(response.json())

datamasque/client/rulesets.py

@@ -0,0 +1,57 @@
+import logging
+
+from datamasque.client.base import BaseClient
+from datamasque.client.exceptions import DataMasqueException
+from datamasque.client.models.ruleset import Ruleset, RulesetId
+from datamasque.client.models.status import ValidationStatus
+
+logger = logging.getLogger(__name__)
+
+
+class RulesetClient(BaseClient):
+    """Ruleset CRUD API methods. Mixed into `DataMasqueClient`."""
+
+    def list_rulesets(self) -> list[Ruleset]:
+        """Returns all rulesets configured on the server."""
+
+        response = self.make_request("GET", "/api/v2/rulesets/")
+        return [Ruleset.model_validate(payload) for payload in response.json()]
+
+    def create_or_update_ruleset(self, ruleset: Ruleset) -> Ruleset:
+        """
+        Creates or updates a ruleset.
+
+        Populates the given ruleset's `id` and `is_valid` fields from the server response,
+        and returns the same ruleset instance for convenience.
+        """
+
+        data = ruleset.model_dump(exclude_none=True, by_alias=True, mode="json")
+        response = self.make_request("POST", "/api/rulesets/", data=data, params={"upsert": "true"})
+        response_data = response.json()
+        ruleset.id = RulesetId(response_data["id"])
+        is_valid = response_data.get("is_valid")
+        if is_valid is not None:
+            ruleset.is_valid = ValidationStatus(is_valid)
+
+        if response.status_code == 201:
+            logger.info('Creation of ruleset "%s" successful', ruleset.name)
+        elif response.status_code == 200:
+            logger.debug('Update of ruleset "%s" successful', ruleset.name)
+
+        return ruleset
+
+    def delete_ruleset_by_id_if_exists(self, ruleset_id: RulesetId) -> None:
+        """Deletes the ruleset with the given ID. No-op if the ruleset does not exist."""
+
+        self._delete_if_exists(f"/api/rulesets/{ruleset_id}/")
+
+    def delete_ruleset_by_name_if_exists(self, ruleset_name: str) -> None:
+        """Deletes the ruleset with the given name. No-op if the ruleset does not exist."""
+
+        all_rulesets = self.list_rulesets()
+        rulesets_matching_name = [ruleset for ruleset in all_rulesets if ruleset.name == ruleset_name]
+        for ruleset in rulesets_matching_name:
+            if ruleset.id is None:
+                raise DataMasqueException(f'Server returned a ruleset named "{ruleset.name}" without an `id`.')
+
+            self.delete_ruleset_by_id_if_exists(ruleset.id)