diff --git a/datamasque/client/__init__.py b/datamasque/client/__init__.py index ea8e0c9..c2c3b29 100644 --- a/datamasque/client/__init__.py +++ b/datamasque/client/__init__.py @@ -38,6 +38,7 @@ DocumentDbConnectionConfig, DynamoConnectionConfig, FileConnectionConfig, + LicenseLock, MongoConnectionConfig, MountedShareConnectionConfig, MssqlLinkedServerConnectionConfig, @@ -243,6 +244,7 @@ "LengthUnit", "LengthsStatistics", "LicenseInfo", + "LicenseLock", "Locator", "MaskType", "MaskedFormEntry", diff --git a/datamasque/client/models/connection.py b/datamasque/client/models/connection.py index 4e0aa41..86861b0 100644 --- a/datamasque/client/models/connection.py +++ b/datamasque/client/models/connection.py @@ -1,7 +1,8 @@ """Connection configuration models for the DataMasque API.""" +from datetime import datetime from enum import Enum -from typing import Any, Callable, Literal, NewType, Optional +from typing import Annotated, Any, Callable, Literal, NewType, Optional from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator @@ -99,6 +100,19 @@ def _validate_kms_key(self) -> "SseConfig": return self +class LicenseLock(BaseModel): + """ + Read-only license-lock metadata for a connection. + + The server reports this only on instances whose license has a finite connection cap, + once the connection's first run has finished successfully; + otherwise it is `null`. + """ + + locked_at: datetime + editable_from: datetime + + class ConnectionConfig(BaseModel): """ Base class for all connection configurations. @@ -111,6 +125,8 @@ class ConnectionConfig(BaseModel): name: str id: Optional[ConnectionId] = None + # Server-populated and read-only, so it is excluded when serializing a config back to the API. + license_lock: Annotated[Optional[LicenseLock], Field(exclude=True)] = None class DynamoConnectionConfig(ConnectionConfig): diff --git a/tests/test_connections.py b/tests/test_connections.py index ed894f9..7eb39a5 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -1,5 +1,7 @@ """Tests for `ConnectionClient` (CRUD + Snowflake-specific behaviour).""" +from datetime import datetime, timezone + import pytest import requests_mock @@ -13,6 +15,7 @@ DatabricksConnectionConfig, DocumentDbConnectionConfig, DynamoConnectionConfig, + LicenseLock, MongoConnectionConfig, MountedShareConnectionConfig, MssqlLinkedServerConnectionConfig, @@ -682,6 +685,68 @@ def test_s3_connection_model_validate(): assert conn.iam_role_arn == "arn:aws:iam::111122223333:role/s3-role" +def test_connection_model_validate_parses_license_lock() -> None: + payload = { + "id": "88dabb63-aca5-4cc4-8f76-f78736a42f39", + "name": "s3", + "mask_type": "file", + "type": "s3_connection", + "base_directory": "data/", + "is_file_mask_source": True, + "is_file_mask_destination": False, + "bucket": "my-bucket", + "license_lock": { + "locked_at": "2026-01-01T12:00:00Z", + "editable_from": "2026-04-01T12:00:00Z", + }, + } + + conn = validate_connection(payload) + + assert isinstance(conn.license_lock, LicenseLock) + assert conn.license_lock.locked_at == datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + assert conn.license_lock.editable_from == datetime(2026, 4, 1, 12, 0, tzinfo=timezone.utc) + + +def test_connection_model_validate_null_license_lock() -> None: + payload = { + "id": "id-1", + "name": "s3", + "mask_type": "file", + "type": "s3_connection", + "base_directory": "", + "is_file_mask_source": True, + "is_file_mask_destination": False, + "bucket": "my-bucket", + "license_lock": None, + } + + conn = validate_connection(payload) + assert conn.license_lock is None + + +def test_license_lock_is_excluded_when_serializing() -> None: + # `license_lock` is server-populated and read-only, so it must never be sent back on create/update. + conn = S3ConnectionConfig.model_validate( + { + "id": "id-1", + "name": "s3", + "mask_type": "file", + "type": "s3_connection", + "base_directory": "", + "is_file_mask_source": True, + "is_file_mask_destination": False, + "bucket": "my-bucket", + "license_lock": { + "locked_at": "2026-01-01T12:00:00Z", + "editable_from": "2026-04-01T12:00:00Z", + }, + } + ) + + assert "license_lock" not in conn.model_dump(by_alias=True, mode="json") + + def test_s3_connection_model_validate_no_iam_role(): payload = { "id": "id-1",