dq-export-connector 0.1.0.dev1__py3-none-any.whl

dq_export_connector/_models/bucket_validation_result.py

@@ -0,0 +1,63 @@
+# generated by datamodel-codegen:
+#   filename:  bucket_validation_result.json
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+
+class BucketValidationResult1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    status: Literal['OK']
+
+
+class ExportMessage(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    code: Annotated[
+        str,
+        Field(
+            description='SCREAMING_SNAKE_CASE, such as `ACCESS_DENIED`. A string rather than an\nenum, so a worker release can add codes that readers built against an\nolder ddc still accept.'
+        ),
+    ]
+    message: Annotated[
+        str,
+        Field(
+            description='Safe to show to the user: no credentials and no dataset values.'
+        ),
+    ]
+
+
+class BucketValidationResult2(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    status: Literal['OK_WITH_WARNINGS']
+    warnings: list[ExportMessage]
+
+
+class BucketValidationResult3(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    errors: list[ExportMessage]
+    status: Literal['ERROR']
+
+
+class BucketValidationResult(
+    RootModel[
+        BucketValidationResult1 | BucketValidationResult2 | BucketValidationResult3
+    ]
+):
+    root: Annotated[
+        BucketValidationResult1 | BucketValidationResult2 | BucketValidationResult3,
+        Field(
+            description='The result of a Bucket `ValidateConnection`. Connection problems are\nreported here, and the task still succeeds.',
+            title='BucketValidationResult',
+        ),
+    ]

dq_export_connector/_models/export_connector_credentials.py

@@ -0,0 +1,84 @@
+# generated by datamodel-codegen:
+#   filename:  export_connector_credentials.json
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+
+class S3(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    access_key: Annotated[str, Field(alias='accessKey')]
+    secret_key: Annotated[str, Field(alias='secretKey')]
+
+
+class BucketCredentials1(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    s3: S3
+
+
+class AzureBlob(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    sas_token: Annotated[
+        str,
+        Field(
+            alias='sasToken',
+            description='A container SAS token — the scoped, expiring grant Azure recommends\nover an account key. Only the grant: the account and container it\napplies to are pinned by the signed `BucketConnection`, so rotating\nthe token can never repoint the connector.',
+        ),
+    ]
+
+
+class BucketCredentials2(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    azure_blob: Annotated[AzureBlob, Field(alias='azureBlob')]
+
+
+class Gcs(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    service_account_json: Annotated[str, Field(alias='serviceAccountJson')]
+
+
+class BucketCredentials3(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    gcs: Gcs
+
+
+class BucketCredentials(
+    RootModel[BucketCredentials1 | BucketCredentials2 | BucketCredentials3]
+):
+    root: BucketCredentials1 | BucketCredentials2 | BucketCredentials3
+
+
+class ExportConnectorCredentials1(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    bucket: BucketCredentials
+
+
+class ExportConnectorCredentials(RootModel[ExportConnectorCredentials1]):
+    root: Annotated[
+        ExportConnectorCredentials1,
+        Field(
+            description='Keyed by Connector Kind, in step with [`ExportConnectorKind`]. The worker\nreads it at `/input/credentials.json`.',
+            title='ExportConnectorCredentials',
+        ),
+    ]

dq_export_connector/_models/export_receipt.py

@@ -0,0 +1,152 @@
+# generated by datamodel-codegen:
+#   filename:  export_receipt.json
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+
+class ExportMessage(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    code: Annotated[
+        str,
+        Field(
+            description='SCREAMING_SNAKE_CASE, such as `ACCESS_DENIED`. A string rather than an\nenum, so a worker release can add codes that readers built against an\nolder ddc still accept.'
+        ),
+    ]
+    message: Annotated[
+        str,
+        Field(
+            description='Safe to show to the user: no credentials and no dataset values.'
+        ),
+    ]
+
+
+class ExportTotals(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    bytes: Annotated[int, Field(ge=0)]
+    files: Annotated[int, Field(ge=0)]
+    rows: Annotated[
+        int | None,
+        Field(
+            description='`None` unless the worker counted the rows of every file.', ge=0
+        ),
+    ] = None
+
+
+class ExportedFile(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    bytes: Annotated[int, Field(ge=0)]
+    path: Annotated[str, Field(description='The object key in the bucket.')]
+    rows: Annotated[
+        int | None, Field(description="`None` where the worker can't count rows.", ge=0)
+    ] = None
+
+
+class InputStatus(Enum):
+    success = 'SUCCESS'
+    failed = 'FAILED'
+
+
+class BucketInputReceipt(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    errors: list[ExportMessage]
+    files: list[ExportedFile]
+    leaf_id: Annotated[str, Field(alias='leafId')]
+    status: InputStatus
+    totals: ExportTotals
+    warnings: list[ExportMessage]
+
+
+class ExportReceiptV01(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    errors: Annotated[
+        list[ExportMessage],
+        Field(
+            description='Problems tied to no single Input, such as credentials rejected before\nany upload.'
+        ),
+    ]
+    export_prefix: Annotated[
+        str | None,
+        Field(
+            alias='exportPrefix',
+            description='The export prefix with its date tokens filled; `None` without one.',
+        ),
+    ] = None
+    inputs: list[BucketInputReceipt]
+    totals: ExportTotals
+    warnings: list[ExportMessage]
+    type: Literal['BUCKET']
+
+
+class ExportReceiptV0(RootModel[ExportReceiptV01]):
+    root: ExportReceiptV01
+
+
+class ExportReceipt1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    version: Literal['V0']
+    errors: Annotated[
+        list[ExportMessage],
+        Field(
+            description='Problems tied to no single Input, such as credentials rejected before\nany upload.'
+        ),
+    ]
+    export_prefix: Annotated[
+        str | None,
+        Field(
+            alias='exportPrefix',
+            description='The export prefix with its date tokens filled; `None` without one.',
+        ),
+    ] = None
+    inputs: list[BucketInputReceipt]
+    totals: ExportTotals
+    warnings: list[ExportMessage]
+    type: Literal['BUCKET']
+
+
+class ExportReceipt(RootModel[ExportReceipt1]):
+    root: Annotated[
+        ExportReceipt1,
+        Field(
+            description='A released version never changes: a change is a new version, and readers\nkeep reading every version ever written.',
+            title='ExportReceipt',
+        ),
+    ]
+
+
+class BucketExportReceiptV0(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    errors: Annotated[
+        list[ExportMessage],
+        Field(
+            description='Problems tied to no single Input, such as credentials rejected before\nany upload.'
+        ),
+    ]
+    export_prefix: Annotated[
+        str | None,
+        Field(
+            alias='exportPrefix',
+            description='The export prefix with its date tokens filled; `None` without one.',
+        ),
+    ] = None
+    inputs: list[BucketInputReceipt]
+    totals: ExportTotals
+    warnings: list[ExportMessage]

dq_export_connector/_models/export_worker_config.py

@@ -0,0 +1,183 @@
+# generated by datamodel-codegen:
+#   filename:  export_worker_config.json
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+
+class BucketConnection1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    bucket: str
+    path_prefix: Annotated[str | None, Field(alias='pathPrefix')] = None
+    region: str
+    type: Literal['s3']
+
+
+class BucketConnection2(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    account_name: Annotated[str, Field(alias='accountName')]
+    container: str
+    path_prefix: Annotated[str | None, Field(alias='pathPrefix')] = None
+    type: Literal['azureBlob']
+
+
+class BucketConnection3(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    bucket: str
+    path_prefix: Annotated[str | None, Field(alias='pathPrefix')] = None
+    type: Literal['gcs']
+
+
+class BucketConnection(
+    RootModel[BucketConnection1 | BucketConnection2 | BucketConnection3]
+):
+    root: Annotated[
+        BucketConnection1 | BucketConnection2 | BucketConnection3,
+        Field(
+            description='Cloud-bucket addressing (no credentials). Provider-structurally paired with\nthe matching [`BucketCredentials`] variant in the confidential half.'
+        ),
+    ]
+
+
+class BucketInput(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    destination_directory: Annotated[str, Field(alias='destinationDirectory')]
+    leaf_id: Annotated[str, Field(alias='leafId')]
+
+
+class ExportConnectorKind1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    connection: BucketConnection
+    export_prefix: Annotated[
+        str | None,
+        Field(
+            alias='exportPrefix',
+            description='May hold `{yyyy}`, `{mm}`, `{dd}` and `{hh}`, which the worker fills in\nUTC when a run starts.',
+        ),
+    ] = None
+    inputs: list[BucketInput]
+    type: Literal['bucket']
+
+
+class ExportConnectorKind(RootModel[ExportConnectorKind1]):
+    root: ExportConnectorKind1
+
+
+class ExportParams1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    type: Literal['bucket']
+
+
+class ExportParams(RootModel[ExportParams1]):
+    root: ExportParams1
+
+
+class ProvisionedPolicy(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    leaf_id: Annotated[str, Field(alias='leafId')]
+    policy_id: Annotated[str, Field(alias='policyId')]
+    provisioned_at: Annotated[
+        str, Field(alias='provisionedAt', description='RFC 3339.')
+    ]
+    provisioned_by_email: Annotated[
+        str,
+        Field(
+            alias='provisionedByEmail',
+            description='Whose READ on the dataset every run re-checks.',
+        ),
+    ]
+
+
+class ValidateConnectionParams1(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    type: Literal['bucket']
+
+
+class ValidateConnectionParams(RootModel[ValidateConnectionParams1]):
+    root: ValidateConnectionParams1
+
+
+class BucketExportConfig(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    connection: BucketConnection
+    export_prefix: Annotated[
+        str | None,
+        Field(
+            alias='exportPrefix',
+            description='May hold `{yyyy}`, `{mm}`, `{dd}` and `{hh}`, which the worker fills in\nUTC when a run starts.',
+        ),
+    ] = None
+    inputs: list[BucketInput]
+
+
+class ValidateConnection(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    params: ValidateConnectionParams
+
+
+class ExportConnectorDcrComputeAction1(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    validate_connection: Annotated[
+        ValidateConnection, Field(alias='validateConnection')
+    ]
+
+
+class Export(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    params: ExportParams
+
+
+class ExportConnectorDcrComputeAction2(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+        populate_by_name=True,
+    )
+    export: Export
+
+
+class ExportConnectorDcrComputeAction(
+    RootModel[ExportConnectorDcrComputeAction1 | ExportConnectorDcrComputeAction2]
+):
+    root: Annotated[
+        ExportConnectorDcrComputeAction1 | ExportConnectorDcrComputeAction2,
+        Field(
+            description='Compute actions for an Export Connector DCR, started through `POST /jobs`.'
+        ),
+    ]
+
+
+class ExportWorkerConfig(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    action: ExportConnectorDcrComputeAction
+    connector: ExportConnectorKind
+    policies: list[ProvisionedPolicy]

dq_export_connector/bucket/datasets.py

@@ -0,0 +1,71 @@
+"""How an Input's dataset maps to files. A ZIP whose entries are all stored
+without compression is unpacked, each entry keeping its path. Anything else,
+including a ZIP with a compressed entry, is one raw file named `data`."""
+
+from __future__ import annotations
+
+import zipfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO
+
+RAW_FILE_NAME = "data"
+
+
+@dataclass(frozen=True)
+class DatasetFile:
+    """`path` is relative to the Input's destination directory."""
+
+    path: str
+    size: int
+
+
+@dataclass(frozen=True)
+class Layout:
+    """`raw_reason` says why a dataset isn't unpacked; `None` for a stored ZIP."""
+
+    files: list[DatasetFile]
+    raw_reason: str | None
+
+
+@contextmanager
+def open_dataset(dataset_path: Path) -> Iterator[tuple[Layout, _Reader]]:
+    """Reads the dataset's layout, keeping a ZIP open while its entries upload.
+    A ZIP's index sits at its end, so detecting one reads no more than that."""
+    try:
+        archive = zipfile.ZipFile(dataset_path)
+    except zipfile.BadZipFile:
+        size = dataset_path.stat().st_size
+        yield Layout([DatasetFile(RAW_FILE_NAME, size)], "not a ZIP"), _RawReader(dataset_path)
+        return
+    with archive:
+        entries = archive.infolist()
+        if any(entry.compress_type != zipfile.ZIP_STORED for entry in entries):
+            size = dataset_path.stat().st_size
+            yield Layout([DatasetFile(RAW_FILE_NAME, size)], "a ZIP with compressed entries"), _RawReader(dataset_path)
+            return
+        files = [DatasetFile(entry.filename, entry.file_size) for entry in entries if not entry.is_dir()]
+        yield Layout(files, None), _ZipReader(archive)
+
+
+class _Reader:
+    def open(self, file: DatasetFile) -> BinaryIO:
+        raise NotImplementedError
+
+
+class _RawReader(_Reader):
+    def __init__(self, dataset_path: Path) -> None:
+        self._dataset_path = dataset_path
+
+    def open(self, file: DatasetFile) -> BinaryIO:
+        return self._dataset_path.open("rb")
+
+
+class _ZipReader(_Reader):
+    def __init__(self, archive: zipfile.ZipFile) -> None:
+        self._archive = archive
+
+    def open(self, file: DatasetFile) -> BinaryIO:
+        return self._archive.open(file.path)  # type: ignore[return-value]