datamasque-python 1.0.0__py3-none-any.whl

datamasque/client/models/discovery.py

@@ -0,0 +1,229 @@
+"""Typed request and response shapes for schema-discovery and ruleset-generation endpoints."""
+
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from datamasque.client.models.connection import ConnectionConfig, ConnectionId, unwrap_connection_id
+from datamasque.client.models.data_selection import HashColumnsTableConfig, Locator, UserSelection
+from datamasque.client.models.pagination import Page
+
+
+class InDataDiscoveryRule(BaseModel):
+    """A single rule for in-data discovery."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: Optional[str] = None
+    pattern: str
+
+
+class InDataDiscoveryConfig(BaseModel):
+    """In-data discovery configuration nested under `SchemaDiscoveryRequest.in_data_discovery`."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    enabled: Optional[bool] = None
+    row_sample_size: Optional[int] = None
+    custom_rules: Optional[list[InDataDiscoveryRule]] = None
+    non_sensitive_rules: Optional[list[InDataDiscoveryRule]] = None
+    force: Optional[bool] = None
+
+
+class SchemaDiscoveryRequest(BaseModel):
+    """
+    Request body for `POST /api/schema-discovery/`.
+
+    `connection` accepts either a `ConnectionId` or a full `ConnectionConfig` returned by an earlier client call.
+    Every other field uses the server's default value when omitted.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    connection: Union[ConnectionId, ConnectionConfig]
+    custom_keywords: list[str] = Field(default_factory=list)
+    ignored_keywords: list[str] = Field(default_factory=list)
+    schemas: list[str] = Field(default_factory=list)
+    in_data_discovery: Optional[InDataDiscoveryConfig] = None
+    disable_built_in_keywords: bool = False
+    disable_global_custom_keywords: bool = False
+    disable_global_ignored_keywords: bool = False
+
+    @field_validator("connection", mode="before")
+    @classmethod
+    def _unwrap_connection(cls, value: Any) -> Any:
+        return unwrap_connection_id(value)
+
+
+class RulesetGenerationRequest(BaseModel):
+    """
+    Request body for `POST /api/generate-ruleset/v2/`.
+
+    `connection` accepts either a `ConnectionId` or a full `ConnectionConfig` returned by an earlier client call.
+    `selected_columns` is the same nested `schema -> table -> [column, ...]` mapping
+    used by `SelectedColumns.columns`,
+    and `hash_columns` follows the `HashColumnsTableConfig` shape.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    connection: Union[ConnectionId, ConnectionConfig]
+    selected_columns: dict[str, dict[str, list[str]]]
+    hash_columns: Optional[dict[str, dict[str, HashColumnsTableConfig]]] = None
+
+    @field_validator("connection", mode="before")
+    @classmethod
+    def _unwrap_connection(cls, value: Any) -> Any:
+        return unwrap_connection_id(value)
+
+
+class FileRulesetGenerationRequest(BaseModel):
+    """
+    Request body for `POST /api/generate-file-ruleset/`.
+
+    `connection` accepts either a `ConnectionId` or a full `ConnectionConfig` returned by an earlier client call.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    connection: Union[ConnectionId, ConnectionConfig]
+    selected_data: list[UserSelection]
+
+    @field_validator("connection", mode="before")
+    @classmethod
+    def _unwrap_connection(cls, value: Any) -> Any:
+        return unwrap_connection_id(value)
+
+
+class DiscoveryMatch(BaseModel):
+    """A single match found by schema or file discovery."""
+
+    model_config = ConfigDict(extra="allow")
+
+    label: str
+    categories: list[str]
+    flagged_by: str
+    description: str
+    hit_ratio: Optional[int] = None  # None for metadata matches, percentage 0-100 for IDD matches.
+
+
+class ForeignKeyRef(BaseModel):
+    """A foreign key declared on a column, pointing to another column it references."""
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    referenced_column: str  # Dotted path: "schema.table.column".
+
+
+class ReferencingForeignKey(BaseModel):
+    """A foreign key declared on another column that points *at* this column."""
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    referencing_column: str  # Dotted path: "schema.table.column".
+
+
+class SchemaDiscoveryColumn(BaseModel):
+    """Column-level data in a schema discovery result."""
+
+    model_config = ConfigDict(extra="allow")
+
+    data_type: Optional[str] = None
+    max_length: Optional[int] = None
+    foreign_keys: list[ForeignKeyRef]
+    discovery_matches: list[DiscoveryMatch]
+    numeric_precision: Optional[int] = None
+    numeric_scale: Optional[int] = None
+    constraint_columns: list[str]
+    pk_constraint_name: Optional[str] = None
+    uk_constraint_name: Optional[str] = None
+    unique_index_names: list[str]
+    referencing_foreign_keys: list[ReferencingForeignKey]
+    constraint: str  # Primary or Unique, or empty string if column does not participate in a PK/UK
+
+
+class SchemaDiscoveryResult(BaseModel):
+    """A single row in the v2 schema discovery results."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    id: int
+    column: str
+    table: str
+    schema_name: Optional[str] = Field(default=None, alias="schema")  # "schema" is a reserved word in Pydantic
+    data: SchemaDiscoveryColumn
+
+
+class ConstraintColumns(BaseModel):
+    """A constraint's column list in table metadata."""
+
+    model_config = ConfigDict(extra="allow")
+
+    columns: list[str]
+
+
+class TableConstraints(BaseModel):
+    """Constraint metadata for a single table."""
+
+    model_config = ConfigDict(extra="allow")
+
+    primary_keys: Optional[list[ConstraintColumns]] = None
+    unique_keys: Optional[list[ConstraintColumns]] = None
+    foreign_keys: Optional[list[ConstraintColumns]] = None
+
+
+class SchemaDiscoveryPage(Page[SchemaDiscoveryResult]):
+    """
+    Admin-server envelope for `GET /api/schema-discovery/v2/{run_id}/`.
+
+    Extends the standard `Page` with `table_metadata`.
+    """
+
+    table_metadata: Optional[dict[str, dict[str, TableConstraints]]] = None
+
+
+class FileDiscoveryMatch(BaseModel):
+    """A single match in a file discovery locator."""
+
+    model_config = ConfigDict(extra="allow")
+
+    categories: Optional[list[str]] = None
+    flagged_by: Optional[str] = None
+    description: Optional[str] = None
+    label: Optional[str] = None
+    hit_ratio: Optional[int] = None
+
+
+class FileDiscoveryLocatorResult(BaseModel):
+    """A locator (column/path) within a discovered file."""
+
+    model_config = ConfigDict(extra="allow")
+
+    locator: Optional[Locator] = None
+    matches: Optional[list[FileDiscoveryMatch]] = None
+    data_types: Optional[list[str]] = None
+
+
+class FileDiscoveryFile(BaseModel):
+    """A file entry in a file discovery result."""
+
+    model_config = ConfigDict(extra="allow")
+
+    path: Optional[str] = None
+    file_type: Optional[str] = None
+    delimiter: Optional[str] = None
+    encoding: Optional[str] = None
+
+
+class FileDiscoveryResult(BaseModel):
+    """A single record from `GET /api/runs/{run_id}/file-discovery-results/`."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: Optional[int] = None
+    connection: Optional[Any] = None
+    file_type: Optional[str] = None
+    files: Optional[list[FileDiscoveryFile]] = None
+    results: Optional[list[FileDiscoveryLocatorResult]] = None

datamasque/client/models/dm_instance.py

@@ -0,0 +1,39 @@
+from typing import Callable, Optional
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from datamasque.client.exceptions import DataMasqueUserError
+
+
+class DataMasqueInstanceConfig(BaseModel):
+    """
+    Connection configuration for `DataMasqueClient`.
+
+    `base_url` is the root URL of the DataMasque admin server
+    (e.g. `https://datamasque.example.com/`).
+    Set `verify_ssl=False` to skip TLS certificate verification
+    (only use this with a self-signed certificate;
+    do not disable it otherwise).
+    Exactly one of `password` or `token_source` must be set.
+    `token_source` is a user-supplied callable that returns the bare API token string —
+    the hex value returned by `POST /api/auth/token/login/`;
+    the client prepends it with `Token ` when sending the `Authorization` header.
+    The client calls `token_source` on each authentication attempt,
+    so the callable is free to fetch and refresh tokens out-of-band (e.g. from a secrets manager).
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    base_url: str
+    username: str
+    password: Optional[str] = None
+    verify_ssl: bool = True
+    token_source: Optional[Callable[[], str]] = None
+
+    @model_validator(mode="after")
+    def _validate_auth_source(self) -> "DataMasqueInstanceConfig":
+        if (self.password is None) == (self.token_source is None):
+            raise DataMasqueUserError(
+                "Exactly one of `password` or `token_source` must be provided to `DataMasqueInstanceConfig`."
+            )
+        return self

datamasque/client/models/files.py

@@ -0,0 +1,89 @@
+from abc import abstractmethod
+from datetime import datetime
+from typing import NewType, Optional
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+FileId = NewType("FileId", str)
+
+
+class DataMasqueFile(BaseModel):
+    """Base class for the concrete file types (`SeedFile`, `OracleWalletFile`, `SslZipFile`, `SnowflakeKeyFile`)."""
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    created_date: datetime
+    modified_date: Optional[datetime] = None
+    id: Optional[FileId] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def _promote_filename(cls, data: dict) -> dict:
+        """The API sometimes returns `filename` instead of `name`."""
+        if isinstance(data, dict):
+            if "filename" in data and "name" not in data:
+                data["name"] = data["filename"]
+        return data
+
+    @classmethod
+    @abstractmethod
+    def get_url(cls) -> str:
+        """Returns the API URL path for files of this type."""
+
+        raise NotImplementedError  # pragma: no cover
+
+    @classmethod
+    @abstractmethod
+    def get_content_param_name(cls) -> str:
+        """Returns the multipart form field name used when uploading files of this type."""
+
+        raise NotImplementedError  # pragma: no cover
+
+
+class SeedFile(DataMasqueFile):
+    """Represents a seed file (CSV file)."""
+
+    @classmethod
+    def get_url(cls) -> str:
+        return "api/seeds/"
+
+    @classmethod
+    def get_content_param_name(cls) -> str:
+        return "seed_file"
+
+
+class OracleWalletFile(DataMasqueFile):
+    """Represents an Oracle wallet file (ZIP file)."""
+
+    @classmethod
+    def get_url(cls) -> str:
+        return "api/oracle-wallets/"
+
+    @classmethod
+    def get_content_param_name(cls) -> str:
+        return "zip_archive"
+
+
+class SslZipFile(DataMasqueFile):
+    """Represents a ZIP file of SSL certificates used to establish secure database connections."""
+
+    @classmethod
+    def get_url(cls) -> str:
+        return "api/connection-filesets/"
+
+    @classmethod
+    def get_content_param_name(cls) -> str:
+        return "zip_archive"
+
+
+class SnowflakeKeyFile(DataMasqueFile):
+    """Represents a private SSH key file for Snowflake connections."""
+
+    @classmethod
+    def get_url(cls) -> str:
+        return "api/files/snowflake-keys/"
+
+    @classmethod
+    def get_content_param_name(cls) -> str:
+        return "key_file"

datamasque/client/models/ifm.py

@@ -0,0 +1,177 @@
+"""Typed request and response shapes for the IFM (in-flight masking) HTTP API."""
+
+from datetime import datetime
+from typing import Any, Callable, Optional
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from datamasque.client.exceptions import DataMasqueUserError
+
+
+class RulesetPlanOptions(BaseModel):
+    """
+    Server-defined defaults applied when a mask request omits the corresponding fields.
+
+    All keys are optional;
+    callers can supply any subset (or none) and the IFM server fills in remaining defaults.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    enabled: Optional[bool] = None
+    # NB: Encoding and charset are not currently implemented for IFM.
+    # These fields are here just to ensure we can round-trip a `RulesetPlan` object.
+    default_encoding: Optional[str] = None
+    default_charset: Optional[str] = None
+    default_log_level: Optional[str] = None
+
+
+class IfmLog(BaseModel):
+    """A single log entry produced by IFM during a mask call or a ruleset-plan validation."""
+
+    model_config = ConfigDict(extra="allow")
+
+    log_level: str
+    timestamp: str
+    message: str
+
+
+class IfmRulesetPlanRef(BaseModel):
+    """Reference to a ruleset plan embedded in a mask response."""
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    serial: int
+
+
+class RulesetPlan(BaseModel):
+    """
+    Unified model for IFM ruleset plans.
+
+    Collapses the list/detail/create/update response shapes into one model
+    with optional fields for parts that differ by endpoint.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    name: str
+    serial: int
+    created_time: datetime
+    modified_time: datetime
+    options: RulesetPlanOptions
+    ruleset_yaml: Optional[str] = None
+    logs: Optional[list[IfmLog]] = None
+    url: Optional[str] = None
+
+
+class RulesetPlanCreateRequest(BaseModel):
+    """Request body for `POST /ifm/ruleset-plans/`."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: str
+    ruleset_yaml: str
+    options: Optional[RulesetPlanOptions] = None
+
+
+class RulesetPlanUpdateRequest(BaseModel):
+    """Request body for `PUT /ifm/ruleset-plans/{name}/`."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    ruleset_yaml: str
+    options: Optional[RulesetPlanOptions] = None
+
+
+class RulesetPlanPartialUpdateRequest(BaseModel):
+    """Request body for `PATCH /ifm/ruleset-plans/{name}/` — every field is optional."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    ruleset_yaml: Optional[str] = None
+    options: Optional[RulesetPlanOptions] = None
+
+
+class IfmMaskRequest(BaseModel):
+    """
+    Request body for `POST /ruleset-plans/{name}/mask/`.
+
+    `data` is the list of records to be masked;
+    every other field overrides server defaults configured on the plan.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    data: list[Any]
+    disable_instance_secret: Optional[bool] = None
+    run_secret: Optional[str] = None
+    hash_values: Optional[Any] = None
+    log_level: Optional[str] = None
+    request_id: Optional[str] = None
+    ai_engine_url: Optional[str] = None
+
+
+class IfmMaskResult(BaseModel):
+    """
+    Response shape for `POST /ruleset-plans/{name}/mask/`.
+
+    `success` is populated by the client based on the HTTP status the server returned:
+
+    - `True` — masking completed;
+      `data` carries the masked records (possibly an empty list if the request had no input).
+    - `False` — the server rejected the request with a soft failure
+      (e.g. a masking function received an unsupported value type);
+      `data` is omitted and details surface in `logs`.
+
+    Hard failures (plan not found, auth, transport) still raise rather than producing an `IfmMaskResult`.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    success: bool
+    request_id: Optional[str] = None
+    ruleset_plan: Optional[IfmRulesetPlanRef] = None
+    logs: Optional[list[IfmLog]] = None
+    data: Optional[list[Any]] = None
+
+
+class IfmTokenInfo(BaseModel):
+    """Response body for `GET /verify-token/` — the list of scopes granted to the current JWT."""
+
+    model_config = ConfigDict(extra="allow")
+
+    scopes: list[str]
+
+
+class DataMasqueIfmInstanceConfig(BaseModel):
+    """
+    Connection configuration for `DataMasqueIfmClient`.
+
+    `admin_server_base_url` is where JWTs are obtained and refreshed;
+    `ifm_base_url` is where the IFM API itself lives
+    (typically a separate hostname or the admin server with `/ifm` prefix).
+    Exactly one of `password` or `token_source` must be set.
+    `token_source` is a user-supplied callable that returns the bare JWT access token string —
+    the value issued by the admin server's `/api/auth/jwt/login/` endpoint;
+    the client prepends it with `Bearer ` when sending the `Authorization` header.
+    The client calls `token_source` on each authentication and refresh,
+    so the callable is free to fetch and refresh tokens out-of-band (e.g. from a secrets manager).
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    admin_server_base_url: str
+    ifm_base_url: str
+    username: str
+    password: Optional[str] = None
+    verify_ssl: bool = True
+    token_source: Optional[Callable[[], str]] = None
+
+    @model_validator(mode="after")
+    def _validate_auth_source(self) -> "DataMasqueIfmInstanceConfig":
+        if (self.password is None) == (self.token_source is None):
+            raise DataMasqueUserError(
+                "Exactly one of `password` or `token_source` must be provided to `DataMasqueIfmInstanceConfig`."
+            )
+        return self

datamasque/client/models/license.py

@@ -0,0 +1,60 @@
+"""Typed response shape for the license endpoint."""
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+
+class SwitchableLicenseMetadata(BaseModel):
+    """Metadata for switchable license management (AWS Marketplace, etc.)."""
+
+    model_config = ConfigDict(extra="allow")
+
+    can_switch_license_source: Optional[bool] = None
+    license_source: Optional[str] = None
+    license_select_time: Optional[datetime] = None
+    aws_account_number: Optional[str] = None
+    last_checkout_success_time: Optional[datetime] = None
+    last_checkout_success_type: Optional[str] = None
+    last_checkout_error: Optional[str] = None
+    last_checkout_license_arn: Optional[str] = None
+    last_checkout_product_name: Optional[str] = None
+    last_checkout_contract_expiry: Optional[datetime] = None
+    last_checkout_agreement_id: Optional[str] = None
+    last_checkout_agreement_url: Optional[str] = None
+    checkout_mode: Optional[str] = None
+    selected_product_sku: Optional[str] = None
+    allow_fallback: Optional[bool] = None
+    last_checkout_success_license_count: Optional[int] = None
+    iam_role_arn: Optional[str] = None
+
+
+class LicenseInfo(BaseModel):
+    """
+    License information returned by `GET /api/license/`.
+
+    Core fields (`uuid`, `name`, `type`, `is_expired`, `uploadable`)
+    are always present in the server response.
+    Other fields vary by license type and server version.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    uuid: str
+    name: str
+    type: str
+    is_expired: bool
+    uploadable: bool
+    version: Optional[str] = None
+    raw_type: Optional[str] = None
+    expiry_date: Optional[datetime] = None
+    quota_tb: Optional[float] = None
+    maximum_node_count: Optional[int] = None
+    row_limit: Optional[int] = None
+    platform_name: Optional[str] = None
+    platform_code: Optional[str] = None
+    days_until_expiry: Optional[int] = None
+    is_contract_product: Optional[bool] = None
+    contract_license_type: Optional[str] = None
+    switchable_license_metadata: Optional[SwitchableLicenseMetadata] = None

datamasque/client/models/pagination.py

@@ -0,0 +1,29 @@
+"""Pagination envelope models matching the DataMasque admin-server and IFM list-endpoint response shapes."""
+
+from typing import Generic, Optional, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+T = TypeVar("T")
+
+
+class Page(BaseModel, Generic[T]):
+    """Admin-server paginated response envelope."""
+
+    model_config = ConfigDict(extra="allow")
+
+    count: int
+    next: Optional[str] = None
+    previous: Optional[str] = None
+    results: list[T]
+
+
+class IfmPage(BaseModel, Generic[T]):
+    """IFM paginated response envelope."""
+
+    model_config = ConfigDict(extra="allow")
+
+    items: list[T]
+    total: int
+    limit: int
+    offset: int

datamasque/client/models/ruleset.py

@@ -0,0 +1,45 @@
+import enum
+from typing import Any, NewType, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from datamasque.client.models.status import ValidationStatus
+
+RulesetId = NewType("RulesetId", str)
+
+
+def unwrap_ruleset_id(value: Any) -> Any:
+    """
+    Coerce a `Ruleset` to its `id`; pass other values through unchanged.
+
+    Used by request-model validators that accept either a `RulesetId`
+    or a full `Ruleset` for user convenience.
+    Raises `ValueError` if the ruleset has no `id`
+    (i.e. the caller hasn't yet created it on the server).
+    """
+
+    if isinstance(value, Ruleset):
+        if value.id is None:
+            raise ValueError("Ruleset has not been created yet (id is None)")
+        return value.id
+
+    return value
+
+
+class RulesetType(enum.Enum):
+    """Ruleset type (database masking or file masking)."""
+
+    file = "file"
+    database = "database"
+
+
+class Ruleset(BaseModel):
+    """Represents a ruleset."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    name: str
+    yaml: str = Field(default="", alias="config_yaml")
+    ruleset_type: RulesetType = Field(default=RulesetType.database, alias="mask_type")
+    id: Optional[RulesetId] = None
+    is_valid: Optional[ValidationStatus] = None

datamasque/client/models/ruleset_library.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from typing import NewType, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from datamasque.client.models.status import ValidationStatus
+
+RulesetLibraryId = NewType("RulesetLibraryId", str)
+
+
+class RulesetLibrary(BaseModel):
+    """Represents a ruleset library."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    name: str
+    namespace: str = ""
+    yaml: Optional[str] = Field(default=None, alias="config_yaml")
+    id: Optional[RulesetLibraryId] = None
+    is_valid: Optional[ValidationStatus] = None
+    created: Optional[datetime] = None
+    modified: Optional[datetime] = None