deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/storage/model/manifest.py

@@ -0,0 +1,643 @@
+from __future__ import annotations
+
+import logging
+import itertools
+
+from enum import Enum
+from typing import Optional, List, Dict, Any, TYPE_CHECKING
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    from deltacat.storage.model.schema import FieldLocator
+
+from deltacat import logs
+
+from deltacat.types.media import (
+    ContentType,
+    ContentEncoding,
+    EXT_TO_CONTENT_TYPE,
+    EXT_TO_CONTENT_ENCODING,
+)
+
+import json
+import pyarrow as pa
+import posixpath
+
+from deltacat.utils.filesystem import get_file_info
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class EntryType(str, Enum):
+    """
+    Enum representing all possible content categories of a manifest entry file.
+
+    DATA: The entry contains fully qualified records compliant with the parent
+    table's schema to insert and/or update. Data files for upsert Deltas use
+    this entry's parameters to find matching fields to update. If no entry
+    parameters are specified, then the parent table's primary keys are used.
+    Only records from entries in Deltas with lower stream positions than this
+    entry will be targeted for update.
+
+    POSITIONAL_DELETE: The entry contains pointers to records in other entries
+    to delete. Deleted records will be filtered from query results at runtime.
+
+    EQUALITY_DELETE: The entry contains a subset of field values from the
+    table records to find and delete. The full record of any matching data
+    entries in Deltas with a lower stream position than this entry's Delta
+    will be deleted. The fields used for record discovery are controlled by
+    this entry's parameters. If no entry parameters are specified, then the
+    fields used for record discovery are linked to the parent table's merge
+    keys. The entry may contain additional fields not used for delete record
+    discovery which will be ignored. Deleted records will be filtered from
+    query results at runtime.
+    """
+
+    DATA = "data"
+    POSITIONAL_DELETE = "positional_delete"
+    EQUALITY_DELETE = "equality_delete"
+
+    @classmethod
+    def get_default(cls):
+        return EntryType.DATA
+
+    @classmethod
+    def list(cls):
+        return [c.value for c in EntryType]
+
+
+class EntryParams(dict):
+    """
+    Parameters that control manifest entry interpretation.
+
+    For EQUALITY_DELETE manifest entry types, parameters include equality
+    field identifiers.
+    """
+
+    @staticmethod
+    def of(
+        equality_field_locators: Optional[List["FieldLocator"]] = None,
+    ) -> EntryParams:
+        params = EntryParams()
+        if equality_field_locators is not None:
+            params["equality_field_locators"] = equality_field_locators
+        return params
+
+    @property
+    def equality_field_locators(self) -> Optional[List["FieldLocator"]]:
+        return self.get("equality_field_locators")
+
+
+class Manifest(dict):
+    """
+    A DeltaCAT manifest contains metadata common to multiple manifest formats
+    like Amazon Redshift and Apache Iceberg to simplify dataset import/export.
+    """
+
+    @staticmethod
+    def _build_manifest(
+        meta: Optional[ManifestMeta],
+        entries: Optional[ManifestEntryList],
+        author: Optional[ManifestAuthor] = None,
+        uuid: str = None,
+    ) -> Manifest:
+        if not uuid:
+            uuid = str(uuid4())
+        manifest = Manifest()
+        manifest["id"] = uuid
+        if meta is not None:
+            manifest["meta"] = meta
+        if entries is not None:
+            manifest["entries"] = entries
+        if author is not None:
+            manifest["author"] = author
+        return manifest
+
+    @staticmethod
+    def of(
+        entries: ManifestEntryList,
+        author: Optional[ManifestAuthor] = None,
+        uuid: str = None,
+        entry_type: Optional[EntryType] = None,
+        entry_params: Optional[EntryParams] = None,
+    ) -> Manifest:
+        if not uuid:
+            uuid = str(uuid4())
+        total_record_count = 0
+        total_content_length = 0
+        total_source_content_length = 0
+        content_type = None
+        content_encoding = None
+        credentials = None
+        content_type_params = None
+        schema_id = None
+        sort_scheme_id = None
+        if entries:
+            content_type = entries[0].meta.content_type
+            content_encoding = entries[0].meta.content_encoding
+            credentials = entries[0].meta.credentials
+            content_type_params = entries[0].meta.content_type_parameters
+
+            # Keep the latest schema ID
+            # Schema IDs are >= 0, and schema evolution always increments the last schema ID
+            entry_schema_ids = [
+                entry.meta.schema_id if entry.meta.schema_id is not None else -1
+                for entry in entries
+            ]
+            max_schema_id = max(entry_schema_ids) if entry_schema_ids else -1
+            schema_id = max_schema_id if max_schema_id >= 0 else None
+
+            # Handle sort_scheme_id: set to None if entries have multiple different sort_scheme_ids
+            entry_sort_scheme_ids = set(
+                entry.meta.sort_scheme_id
+                for entry in entries
+                if entry.meta.sort_scheme_id is not None
+            )
+            sort_scheme_id = (
+                list(entry_sort_scheme_ids)[0]
+                if len(entry_sort_scheme_ids) == 1
+                else None
+            )
+
+            for entry in entries:
+                meta = entry.meta
+                if meta.content_type != content_type:
+                    content_type = None
+                if meta.content_encoding != content_encoding:
+                    content_encoding = None
+                entry_content_type = meta.content_type
+                if content_type and entry_content_type != content_type:
+                    msg = (
+                        f"Expected all manifest entries to have content "
+                        f"type '{content_type}' but found "
+                        f"'{entry_content_type}'"
+                    )
+                    raise ValueError(msg)
+                entry_content_encoding = meta.get("content_encoding", None)
+                if content_encoding and entry_content_encoding != content_encoding:
+                    msg = (
+                        f"Expected all manifest entries to have content "
+                        f"encoding '{content_encoding}' but found "
+                        f"'{entry_content_encoding}'"
+                    )
+                    raise ValueError(msg)
+                actual_entry_type = meta.entry_type
+                if entry_type and (actual_entry_type != entry_type):
+                    msg = (
+                        f"Expected all manifest entries to have type "
+                        f"'{entry_type}' but found '{actual_entry_type}'"
+                    )
+                    raise ValueError(msg)
+                actual_entry_params = meta.entry_params
+                if entry_params and (actual_entry_params != entry_params):
+                    msg = (
+                        f"Expected all manifest entries to have params "
+                        f"'{entry_params}' but found '{actual_entry_params}'"
+                    )
+                    raise ValueError(msg)
+                actual_credentials = meta.credentials
+                if credentials and (actual_credentials != credentials):
+                    msg = (
+                        f"Expected all manifest entries to have credentials "
+                        f"'{credentials}' but found '{actual_credentials}'"
+                    )
+                    raise ValueError(msg)
+                actual_content_type_params = meta.content_type_parameters
+                if content_type_params and (
+                    actual_content_type_params != content_type_params
+                ):
+                    msg = (
+                        f"Expected all manifest entries to have content type params "
+                        f"'{content_type_params}' but found '{actual_content_type_params}'"
+                    )
+                    raise ValueError(msg)
+
+                total_record_count += meta.record_count or 0
+                total_content_length += meta.content_length or 0
+                total_source_content_length += meta.source_content_length or 0
+
+        meta = ManifestMeta.of(
+            record_count=total_record_count,
+            content_length=total_content_length,
+            content_type=content_type,
+            content_encoding=content_encoding,
+            source_content_length=total_source_content_length,
+            credentials=credentials,
+            content_type_parameters=content_type_params,
+            entry_type=entry_type,
+            entry_params=entry_params,
+            schema_id=schema_id,
+            sort_scheme_id=sort_scheme_id,
+        )
+        manifest = Manifest._build_manifest(meta, entries, author, uuid)
+        return manifest
+
+    @staticmethod
+    def from_json(json_string: str) -> Manifest:
+        parsed_dict = json.loads(json_string)
+        return Manifest.of(
+            entries=ManifestEntryList.of(
+                [
+                    ManifestEntry.from_dict(entry)
+                    for entry in parsed_dict.get("entries", [])
+                ]
+            ),
+            author=ManifestAuthor.from_dict(parsed_dict.get("author")),
+            uuid=parsed_dict.get("id"),
+        )
+
+    @staticmethod
+    def merge_manifests(
+        manifests: List[Manifest], author: Optional[ManifestAuthor] = None
+    ) -> Manifest:
+        all_entries = ManifestEntryList(
+            itertools.chain(*[m.entries for m in manifests])
+        )
+        merged_manifest = Manifest.of(all_entries, author)
+        return merged_manifest
+
+    @property
+    def meta(self) -> Optional[ManifestMeta]:
+        val: Dict[str, Any] = self.get("meta")
+        if val is not None and not isinstance(val, ManifestMeta):
+            self["meta"] = val = ManifestMeta(val)
+        return val
+
+    @property
+    def entries(self) -> Optional[ManifestEntryList]:
+        val: List[ManifestEntry] = self.get("entries")
+        if val is not None and not isinstance(val, ManifestEntryList):
+            self["entries"] = val = ManifestEntryList.of(val)
+        return val
+
+    @property
+    def id(self) -> str:
+        return self["id"]
+
+    @property
+    def author(self) -> Optional[ManifestAuthor]:
+        val: Dict[str, Any] = self.get("author")
+        if val is not None and not isinstance(val, ManifestAuthor):
+            self["author"] = val = ManifestAuthor(val)
+        return val
+
+
+class ManifestMeta(dict):
+    @staticmethod
+    def of(
+        record_count: Optional[int],
+        content_length: Optional[int],
+        content_type: Optional[str],
+        content_encoding: Optional[str],
+        source_content_length: Optional[int] = None,
+        credentials: Optional[Dict[str, str]] = None,
+        content_type_parameters: Optional[List[Dict[str, str]]] = None,
+        entry_type: Optional[EntryType] = None,
+        entry_params: Optional[EntryParams] = None,
+        schema_id: Optional[int] = None,
+        sort_scheme_id: Optional[str] = None,
+    ) -> ManifestMeta:
+        manifest_meta = ManifestMeta()
+        if record_count is not None:
+            manifest_meta["record_count"] = record_count
+        if content_length is not None:
+            manifest_meta["content_length"] = content_length
+        if source_content_length is not None:
+            manifest_meta["source_content_length"] = source_content_length
+        if content_type is not None:
+            manifest_meta["content_type"] = content_type
+        if content_type_parameters is not None:
+            manifest_meta["content_type_parameters"] = content_type_parameters
+        if content_encoding is not None:
+            manifest_meta["content_encoding"] = content_encoding
+        if credentials is not None:
+            manifest_meta["credentials"] = credentials
+        if entry_type is not None:
+            manifest_meta["entry_type"] = (
+                entry_type.value if isinstance(entry_type, EntryType) else entry_type
+            )
+        if entry_params is not None:
+            manifest_meta["entry_params"] = entry_params
+        if schema_id is not None:
+            manifest_meta["schema_id"] = schema_id
+        if sort_scheme_id is not None:
+            manifest_meta["sort_scheme_id"] = sort_scheme_id
+        return manifest_meta
+
+    @staticmethod
+    def from_dict(obj: dict) -> Optional[ManifestMeta]:
+        if obj is None:
+            return None
+
+        return ManifestMeta.of(
+            record_count=obj.get("record_count"),
+            content_length=obj.get("content_length"),
+            content_type=obj.get("content_type"),
+            content_encoding=obj.get("content_encoding"),
+            source_content_length=obj.get("source_content_length"),
+            credentials=obj.get("credentials"),
+            content_type_parameters=obj.get("content_type_parameters"),
+            entry_type=obj.get("entry_type"),
+            entry_params=obj.get("entry_params"),
+            schema_id=obj.get("schema_id"),
+            sort_scheme_id=obj.get("sort_scheme_id"),
+        )
+
+    @property
+    def record_count(self) -> Optional[int]:
+        return self.get("record_count")
+
+    @property
+    def content_length(self) -> Optional[int]:
+        return self.get("content_length")
+
+    @property
+    def content_type(self) -> Optional[str]:
+        return self.get("content_type")
+
+    @property
+    def content_encoding(self) -> Optional[str]:
+        return self.get("content_encoding")
+
+    @property
+    def source_content_length(self) -> Optional[int]:
+        return self.get("source_content_length")
+
+    @property
+    def content_type_parameters(self) -> Optional[List[Dict[str, str]]]:
+        return self.get("content_type_parameters")
+
+    @content_type_parameters.setter
+    def content_type_parameters(self, params: List[Dict[str, str]]) -> None:
+        self["content_type_parameters"] = params
+
+    @property
+    def credentials(self) -> Optional[Dict[str, str]]:
+        return self.get("credentials")
+
+    @property
+    def entry_type(self) -> Optional[EntryType]:
+        val = self.get("entry_type")
+        if val is not None:
+            return EntryType(self["entry_type"])
+        return val
+
+    @property
+    def entry_params(self) -> Optional[EntryParams]:
+        val: Dict[str, Any] = self.get("entry_params")
+        if val is not None and not isinstance(val, EntryParams):
+            self["entry_params"] = val = EntryParams(val)
+        return val
+
+    @property
+    def schema_id(self) -> Optional[int]:
+        return self.get("schema_id")
+
+    @property
+    def sort_scheme_id(self) -> Optional[str]:
+        return self.get("sort_scheme_id")
+
+
+class ManifestEntry(dict):
+    @staticmethod
+    def of(
+        url: Optional[str],
+        meta: Optional[ManifestMeta],
+        mandatory: bool = True,
+        uri: Optional[str] = None,
+        uuid: Optional[str] = None,
+    ) -> ManifestEntry:
+        manifest_entry = ManifestEntry()
+        if not (uri or url):
+            raise ValueError("No URI or URL specified for manifest entry contents.")
+        if (uri and url) and (uri != url):
+            raise ValueError(f"Manifest entry URI ({uri}) != URL ({url})")
+        if url:
+            manifest_entry["url"] = manifest_entry["uri"] = url
+        elif uri:
+            manifest_entry["url"] = manifest_entry["uri"] = uri
+        if meta is not None:
+            manifest_entry["meta"] = meta
+        if mandatory is not None:
+            manifest_entry["mandatory"] = mandatory
+        if uuid is not None:
+            manifest_entry["id"] = uuid
+        return manifest_entry
+
+    @staticmethod
+    def from_s3_obj_url(
+        url: str,
+        record_count: int,
+        source_content_length: Optional[int] = None,
+        credentials: Optional[Dict[str, str]] = None,
+        content_type_parameters: Optional[List[Dict[str, str]]] = None,
+        entry_type: Optional[EntryType] = None,
+        entry_params: Optional[EntryParams] = None,
+        **s3_client_kwargs,
+    ) -> ManifestEntry:
+        from deltacat.aws import s3u as s3_utils
+
+        s3_obj = s3_utils.get_object_at_url(url, **s3_client_kwargs)
+        logger.debug(f"Building manifest entry from {url}: {s3_obj}")
+        manifest_entry_meta = ManifestMeta.of(
+            record_count=record_count,
+            content_length=s3_obj["ContentLength"],
+            content_type=s3_obj["ContentType"],
+            content_encoding=s3_obj["ContentEncoding"],
+            source_content_length=source_content_length,
+            credentials=credentials,
+            content_type_parameters=content_type_parameters,
+            entry_type=entry_type,
+            entry_params=entry_params,
+        )
+        manifest_entry = ManifestEntry.of(url, manifest_entry_meta)
+        return manifest_entry
+
+    @staticmethod
+    def from_dict(obj: dict) -> ManifestEntry:
+        return ManifestEntry.of(
+            url=obj.get("url"),
+            uri=obj.get("uri"),
+            meta=ManifestMeta.from_dict(obj.get("meta")),
+            mandatory=obj.get("mandatory", True),
+            uuid=obj.get("id"),
+        )
+
+    @staticmethod
+    def from_path(
+        path: str,
+        filesystem: pa.fs.FileSystem,
+        record_count: int,
+        source_content_length: Optional[int] = None,
+        content_type: Optional[str] = None,
+        content_encoding: Optional[str] = None,
+        credentials: Optional[Dict[str, str]] = None,
+        content_type_parameters: Optional[List[Dict[str, str]]] = None,
+        entry_type: Optional[EntryType] = None,
+        entry_params: Optional[EntryParams] = None,
+        schema_id: Optional[int] = None,
+        sort_scheme_id: Optional[str] = None,
+    ) -> ManifestEntry:
+        """
+        Creates a manifest entry from a path using a pyarrow filesystem.
+
+        Args:
+            path: Path to the file
+            filesystem: PyArrow filesystem to use for accessing the file
+            record_count: Number of records in the file
+            source_content_length: Optional original content length in-memory
+                before writing to disk.
+            content_type: Optional content type override. If not provided, will
+                be derived from file extension.
+            content_encoding: Optional content encoding override. If not
+                provided, will be derived from file extension.
+            credentials: Optional credentials required to read this manifest entry.
+            content_type_parameters: Optional content type parameters.
+            entry_type: Optional entry type of this manifest entry. Defaults to DATA.
+            entry_params: Optional entry type parameters.
+            schema_id: Schema ID used to write this manifest entry.
+            sort_scheme_id: Sort scheme ID used to write this manifest entry.
+
+        Returns:
+            A ManifestEntry instance
+        """
+        file_info = get_file_info(path, filesystem)
+        if file_info.type != pa.fs.FileType.File:
+            raise FileNotFoundError(f"Path does not point to a file: {path}")
+
+        # Extract extensions from right to left
+        # First split will get potential encoding extension
+        base_path, ext1 = posixpath.splitext(path)
+
+        # Initialize with defaults for no extensions
+        derived_content_type = ContentType.BINARY
+        derived_content_encoding = ContentEncoding.IDENTITY
+
+        # Only proceed with extension checks if we found at least one extension
+        if ext1:
+            # Check if the first extension is a known encoding
+            derived_content_encoding = EXT_TO_CONTENT_ENCODING.get(
+                ext1,
+                ContentEncoding.IDENTITY,
+            )
+
+            # Get second extension only if first was an encoding
+            if derived_content_encoding != ContentEncoding.IDENTITY:
+                # Second split will get potential content type extension
+                _, ext2 = posixpath.splitext(base_path)
+                if ext2:
+                    derived_content_type = EXT_TO_CONTENT_TYPE.get(
+                        ext2,
+                        ContentType.BINARY,
+                    )
+            else:
+                # First extension wasn't an encoding, check if it's a
+                # content type
+                derived_content_type = EXT_TO_CONTENT_TYPE.get(
+                    ext1,
+                    ContentType.BINARY,
+                )
+
+        if (
+            derived_content_type == ContentType.BINARY
+            and derived_content_encoding != ContentEncoding.IDENTITY
+        ):
+            logger.debug(
+                f"Found encoding {derived_content_encoding.value} but no "
+                f"content type for {path}, assuming binary"
+            )
+
+        # Use provided values if available, otherwise use derived values
+        final_content_type = (
+            content_type if content_type is not None else derived_content_type.value
+        )
+        final_content_encoding = (
+            content_encoding
+            if content_encoding is not None
+            else derived_content_encoding.value
+        )
+
+        manifest_entry_meta = ManifestMeta.of(
+            record_count=record_count,
+            content_length=file_info.size,
+            content_type=final_content_type,
+            content_encoding=final_content_encoding,
+            source_content_length=source_content_length,
+            credentials=credentials,
+            content_type_parameters=content_type_parameters,
+            entry_type=entry_type,
+            entry_params=entry_params,
+            schema_id=schema_id,
+            sort_scheme_id=sort_scheme_id,
+        )
+        manifest_entry = ManifestEntry.of(path, manifest_entry_meta)
+        return manifest_entry
+
+    @property
+    def uri(self) -> Optional[str]:
+        return self.get("uri")
+
+    @property
+    def url(self) -> Optional[str]:
+        return self.get("url")
+
+    @property
+    def meta(self) -> Optional[ManifestMeta]:
+        val: Dict[str, Any] = self.get("meta")
+        if val is not None and not isinstance(val, ManifestMeta):
+            self["meta"] = val = ManifestMeta(val)
+        return val
+
+    @property
+    def mandatory(self) -> bool:
+        return self["mandatory"]
+
+    @property
+    def id(self) -> Optional[str]:
+        return self.get("id")
+
+
+class ManifestAuthor(dict):
+    @staticmethod
+    def of(name: Optional[str], version: Optional[str]) -> ManifestAuthor:
+        manifest_author = ManifestAuthor()
+        if name is not None:
+            manifest_author["name"] = name
+        if version is not None:
+            manifest_author["version"] = version
+        return manifest_author
+
+    @staticmethod
+    def from_dict(obj: dict) -> Optional[ManifestAuthor]:
+        if obj is None:
+            return None
+        return ManifestAuthor.of(obj.get("name"), obj.get("version"))
+
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def version(self) -> Optional[str]:
+        return self.get("version")
+
+
+class ManifestEntryList(List[ManifestEntry]):
+    @staticmethod
+    def of(entries: List[ManifestEntry]) -> ManifestEntryList:
+        manifest_entries = ManifestEntryList()
+        for entry in entries:
+            if entry is not None and not isinstance(entry, ManifestEntry):
+                entry = ManifestEntry(entry)
+            manifest_entries.append(entry)
+        return manifest_entries
+
+    def __getitem__(self, item):
+        val = super().__getitem__(item)
+        if val is not None and not isinstance(val, ManifestEntry):
+            self[item] = val = ManifestEntry(val)
+        return val
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]  # This triggers __getitem__ conversion