pixeltable 0.4.19__py3-none-any.whl → 0.4.21__py3-none-any.whl

pixeltable/metadata/schema.py

@@ -2,7 +2,7 @@ import dataclasses
 import types
 import typing
 import uuid
-from typing import Any, NamedTuple, TypeVar, Union, get_type_hints
+from typing import Any, TypeVar, Union, get_type_hints
 
 import sqlalchemy as sql
 from sqlalchemy import BigInteger, ForeignKey, Integer, LargeBinary, orm
@@ -218,6 +218,12 @@ class TableMd:
             and len(self.column_md) == 0
         )
 
+    @property
+    def ancestor_ids(self) -> list[str]:
+        if self.view_md is None:
+            return []
+        return [id for id, _ in self.view_md.base_versions]
+
 
 class Table(Base):
     """
@@ -349,34 +355,3 @@ class Function(Base):
     dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
     md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # FunctionMd
     binary_obj: orm.Mapped[bytes | None] = orm.mapped_column(LargeBinary, nullable=True)
-
-
-class FullTableMd(NamedTuple):
-    tbl_md: TableMd
-    version_md: TableVersionMd
-    schema_version_md: TableSchemaVersionMd
-
-    @property
-    def is_pure_snapshot(self) -> bool:
-        return (
-            self.tbl_md.view_md is not None
-            and self.tbl_md.view_md.is_snapshot
-            and self.tbl_md.view_md.predicate is None
-            and len(self.schema_version_md.columns) == 0
-        )
-
-    def as_dict(self) -> dict[str, Any]:
-        return {
-            'table_id': self.tbl_md.tbl_id,
-            'table_md': dataclasses.asdict(self.tbl_md),
-            'table_version_md': dataclasses.asdict(self.version_md),
-            'table_schema_version_md': dataclasses.asdict(self.schema_version_md),
-        }
-
-    @classmethod
-    def from_dict(cls, data_dict: dict[str, Any]) -> 'FullTableMd':
-        return FullTableMd(
-            tbl_md=md_from_dict(TableMd, data_dict['table_md']),
-            version_md=md_from_dict(TableVersionMd, data_dict['table_version_md']),
-            schema_version_md=md_from_dict(TableSchemaVersionMd, data_dict['table_schema_version_md']),
-        )

pixeltable/share/__init__.py

@@ -1,3 +1,3 @@
 # ruff: noqa: F401
 
-from .publish import delete_replica, pull_replica, push_replica
+from .publish import delete_replica, list_table_versions, pull_replica, push_replica

pixeltable/share/packager.py

@@ -1,4 +1,5 @@
 import base64
+import dataclasses
 import io
 import json
 import logging
@@ -20,6 +21,7 @@ import sqlalchemy as sql
 
 import pixeltable as pxt
 from pixeltable import catalog, exceptions as excs, metadata, type_system as ts
+from pixeltable.catalog.table_version import TableVersionCompleteMd
 from pixeltable.env import Env
 from pixeltable.exprs.data_row import CellMd
 from pixeltable.metadata import schema
@@ -51,7 +53,7 @@ class TablePackager:
     tmp_dir: Path  # Temporary directory where the package will reside
     tables_dir: Path  # Directory where the Parquet tables will be written
     media_files: dict[Path, str]  # Mapping from local media file paths to their tarball names
-    md: dict[str, Any]
+    bundle_md: dict[str, Any]
 
     bundle_path: Path
     preview_header: dict[str, str]
@@ -62,16 +64,16 @@ class TablePackager:
         self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
-        # Load metadata
+        # Load metadata and convert to JSON immediately
         with catalog.Catalog.get().begin_xact(for_write=False):
             tbl_md = catalog.Catalog.get().load_replica_md(table)
-            self.md = {
+            self.bundle_md = {
                 'pxt_version': pxt.__version__,
                 'pxt_md_version': metadata.VERSION,
-                'md': {'tables': [md.as_dict() for md in tbl_md]},
+                'md': [dataclasses.asdict(md) for md in tbl_md],
             }
         if additional_md is not None:
-            self.md.update(additional_md)
+            self.bundle_md.update(additional_md)
 
     def package(self) -> Path:
         """
@@ -82,7 +84,7 @@ class TablePackager:
         _logger.info(f'Packaging table {self.table._path()!r} and its ancestors in: {self.tmp_dir}')
         self.tmp_dir.mkdir()
         with open(self.tmp_dir / 'metadata.json', 'w', encoding='utf8') as fp:
-            json.dump(self.md, fp)
+            json.dump(self.bundle_md, fp)
         self.tables_dir = self.tmp_dir / 'tables'
         self.tables_dir.mkdir()
         with catalog.Catalog.get().begin_xact(for_write=False):
@@ -94,10 +96,10 @@ class TablePackager:
         self.bundle_path = self.__build_tarball()
 
         _logger.info('Extracting preview data.')
-        self.md['row_count'] = self.table.count()
+        self.bundle_md['row_count'] = self.table.count()
         preview_header, preview = self.__extract_preview_data()
-        self.md['preview_header'] = preview_header
-        self.md['preview_data'] = preview
+        self.bundle_md['preview_header'] = preview_header
+        self.bundle_md['preview_data'] = preview
 
         _logger.info(f'Packaging complete: {self.bundle_path}')
         return self.bundle_path
@@ -358,19 +360,20 @@ class TableRestorer:
 
     Args:
         tbl_path: Pixeltable path (such as 'my_dir.my_table') where the materialized table will be made visible.
-        md: Optional metadata dictionary. If not provided, metadata will be read from the tarball's `metadata.json`.
+        bundle_md: Optional metadata dictionary.
+            If not provided, metadata will be read from the tarball's `metadata.json`.
             The metadata contains table_md, table_version_md, and table_schema_version_md entries for each ancestor
             of the table being restored, as written out by `TablePackager`.
     """
 
     tbl_path: str
-    md: dict[str, Any] | None
+    bundle_md: dict[str, Any] | None
     tmp_dir: Path
     media_files: dict[str, str]  # Mapping from pxtmedia:// URLs to local file:// URLs
 
-    def __init__(self, tbl_path: str, md: dict[str, Any] | None = None) -> None:
+    def __init__(self, tbl_path: str, bundle_md: dict[str, Any] | None = None) -> None:
         self.tbl_path = tbl_path
-        self.md = md
+        self.bundle_md = bundle_md
         self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
@@ -380,12 +383,12 @@ class TableRestorer:
         with tarfile.open(bundle_path, 'r:bz2') as tf:
             tf.extractall(path=self.tmp_dir)
 
-        if self.md is None:
+        if self.bundle_md is None:
             # No metadata supplied; read it from the archive
             with open(self.tmp_dir / 'metadata.json', 'r', encoding='utf8') as fp:
-                self.md = json.load(fp)
+                self.bundle_md = json.load(fp)
 
-        pxt_md_version = self.md['pxt_md_version']
+        pxt_md_version = self.bundle_md['pxt_md_version']
         assert isinstance(pxt_md_version, int)
 
         if pxt_md_version != metadata.VERSION:
@@ -393,8 +396,9 @@ class TableRestorer:
                 f'Pixeltable metadata version mismatch: {pxt_md_version} != {metadata.VERSION}.\n'
                 'Please upgrade Pixeltable to use this dataset: pip install -U pixeltable'
             )
+        # Convert tables metadata from dict to list of TableVersionCompleteMd
+        tbl_md = [schema.md_from_dict(TableVersionCompleteMd, t) for t in self.bundle_md['md']]
 
-        tbl_md = [schema.FullTableMd.from_dict(t) for t in self.md['md']['tables']]
         for md in tbl_md:
             md.tbl_md.is_replica = True
 
@@ -421,7 +425,7 @@ class TableRestorer:
 
             return cat.get_table_by_id(UUID(tbl_md[0].tbl_md.tbl_id))
 
-    def __import_table(self, bundle_path: Path, tv: catalog.TableVersion, tbl_md: schema.FullTableMd) -> None:
+    def __import_table(self, bundle_path: Path, tv: catalog.TableVersion, tbl_md: TableVersionCompleteMd) -> None:
         """
         Import the Parquet table into the Pixeltable catalog.
         """

pixeltable/share/protocol/__init__.py

@@ -0,0 +1,34 @@
+"""
+Pixeltable Core Protocol
+
+This module contains the core protocol structures for pixeltable table operations
+that can be shared between pixeltable core and cloud implementations.
+"""
+
+from .common import PxtUri, RequestBaseModel, StorageDestination
+from .operation_types import ReplicaOperationType
+from .replica import (
+    DeleteRequest,
+    DeleteResponse,
+    FinalizeRequest,
+    FinalizeResponse,
+    PublishRequest,
+    PublishResponse,
+    ReplicateRequest,
+    ReplicateResponse,
+)
+
+__all__ = [
+    'DeleteRequest',
+    'DeleteResponse',
+    'FinalizeRequest',
+    'FinalizeResponse',
+    'PublishRequest',
+    'PublishResponse',
+    'PxtUri',
+    'ReplicaOperationType',
+    'ReplicateRequest',
+    'ReplicateResponse',
+    'RequestBaseModel',
+    'StorageDestination',
+]

pixeltable/share/protocol/common.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import uuid
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any
+from urllib.parse import urlparse
+from uuid import UUID
+
+from pydantic import BaseModel, model_validator
+
+
+class StorageDestination(str, Enum):
+    """Storage destination types for table snapshots."""
+
+    S3 = 's3'
+    R2 = 'r2'
+    GCS = 'gcs'
+
+
+def is_valid_uuid(uuid_string: str) -> bool:
+    """Check if a string is a valid UUID."""
+    try:
+        uuid.UUID(uuid_string)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+class PxtUri(BaseModel):
+    """Pixeltable URI model for pxt:// URIs with validation and parsing."""
+
+    uri: str  # The full URI string
+
+    # Parsed components
+    org: str  # Organization slug from the URI
+    db: str | None  # Database slug from the URI (optional)
+    path: str | None = None  # The table or directory path (None if using UUID)
+    id: UUID | None = None  # The table UUID (None if using path)
+    version: int | None = None  # Optional version number parsed from URI (format: identifier:<version>)
+
+    def __init__(self, uri: str | dict | None = None, **kwargs: Any) -> None:
+        # Handle dict input directly (from JSON deserialization or explicit dict)
+        if isinstance(uri, dict):
+            # Dict input goes directly to Pydantic, which will call parse_uri
+            kwargs.update(uri)
+        elif uri is not None:
+            # Validate that uri is a string when passed as positional argument
+            if not isinstance(uri, str):
+                raise ValueError(f'Invalid data type for PxtUri: expected str or dict, got {type(uri)}')
+            kwargs['uri'] = uri
+        super().__init__(**kwargs)
+
+    @model_validator(mode='before')
+    @classmethod
+    def parse_uri(cls, data: Any) -> dict:
+        # Handle case where data is already a string (from JSON deserialization)
+        if isinstance(data, str):
+            uri = data
+        elif isinstance(data, dict):
+            uri = data.get('uri')
+            if uri is None:
+                raise ValueError('URI must be provided in dict with "uri" key')
+            if not isinstance(uri, str):
+                raise ValueError(f'URI in dict must be a string, got {type(uri)}')
+        else:
+            raise ValueError(f'Invalid data type for PxtUri: expected str or dict, got {type(data)}')
+
+        return {'uri': uri, **cls._parse_and_validate_uri(uri)}
+
+    def __str__(self) -> str:
+        """Return the URI string."""
+        return self.uri
+
+    @classmethod
+    def _parse_and_validate_uri(cls, uri: str) -> dict:
+        """Parse and validate a URI string, return parsed components."""
+        if not uri.startswith('pxt://'):
+            raise ValueError('URI must start with pxt://')
+
+        parsed = urlparse(uri)
+        if parsed.scheme != 'pxt':
+            raise ValueError('URI must use pxt:// scheme')
+
+        if not parsed.netloc:
+            raise ValueError('URI must have an organization')
+
+        # Parse netloc for org and optional db
+        netloc_parts = parsed.netloc.split(':')
+        org = netloc_parts[0]
+        if not org:
+            raise ValueError('URI must have an organization')
+
+        db = netloc_parts[1] if len(netloc_parts) > 1 else None
+
+        # Allow root path (/) as valid, but reject missing path
+        if parsed.path is None:
+            raise ValueError('URI must have a path')
+
+        # Get path and remove leading slash (but keep empty string for root path)
+        # path will be '/' for root directory or '/path/to/table' for regular paths
+        path_part = parsed.path.lstrip('/') if parsed.path else ''
+
+        # Handle version parsing (format: identifier:version)
+        # For root path, path_part will be empty string after lstrip
+        if path_part and ':' in path_part:
+            parts = path_part.rsplit(':', 1)  # Split from right, only once
+            if len(parts) == 2 and parts[1].isdigit():
+                identifier, version = parts[0], int(parts[1])
+            else:
+                identifier, version = path_part, None
+        else:
+            identifier, version = path_part, None
+
+        # Parse identifier into either a path string or UUID
+        path: str | None = None
+        id: UUID | None = None
+        if identifier:
+            if is_valid_uuid(identifier):
+                # It's a UUID
+                id = UUID(identifier)
+            else:
+                # It's a path - keep as string (URI format uses / as separator)
+                # Empty string means root path
+                path = identifier or ''
+        else:
+            # Empty identifier means root path
+            path = ''
+
+        return {'org': org, 'db': db, 'path': path, 'id': id, 'version': version}
+
+    @classmethod
+    def from_components(
+        cls,
+        org: str,
+        path: str | None = None,
+        id: UUID | None = None,
+        db: str | None = None,
+        version: int | None = None,
+    ) -> PxtUri:
+        """Construct a PxtUri from its components."""
+        if path is None and id is None:
+            raise ValueError('Either path or id must be provided')
+        if path is not None and id is not None:
+            raise ValueError('Cannot specify both path and id')
+
+        # Build the URI string from components
+        netloc = org if db is None else f'{org}:{db}'
+
+        # Use path or UUID as identifier
+        if id is not None:
+            identifier = str(id)
+        elif path is not None:
+            # Path is already in URI format (slash-separated)
+            identifier = path or ''
+        else:
+            identifier = ''
+
+        path_part = f'{identifier}:{version}' if version is not None else identifier
+        uri = f'pxt://{netloc}/{path_part}'
+        return cls(uri=uri)
+
+
+class RequestBaseModel(BaseModel, ABC):
+    """Abstract base model for protocol requests that must have a PxtUri."""
+
+    @abstractmethod
+    def get_pxt_uri(self) -> PxtUri:
+        """Get the PxtUri from this request. Must be implemented by subclasses."""
+        pass

pixeltable/share/protocol/operation_types.py

@@ -0,0 +1,33 @@
+"""
+Replica operation types for pixeltable table replica operations.
+
+This module defines the replica operation types that are shared between
+pixeltable core and cloud implementations.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class ReplicaOperationType(str, Enum):
+    """Replica operation types for table replica operations."""
+
+    # Table replica operations
+    PUBLISH_REPLICA = 'publish_replica'
+    FINALIZE_REPLICA = 'finalize_replica'
+    CLONE_REPLICA = 'clone_replica'
+    DELETE_REPLICA = 'delete_replica'
+
+    def is_replica_operation(self) -> bool:
+        """Check if operation is a replica operation."""
+        return self in REPLICA_OPERATIONS
+
+
+# Define the operation sets as module-level constants
+REPLICA_OPERATIONS: set[ReplicaOperationType] = {
+    ReplicaOperationType.PUBLISH_REPLICA,
+    ReplicaOperationType.FINALIZE_REPLICA,
+    ReplicaOperationType.CLONE_REPLICA,
+    ReplicaOperationType.DELETE_REPLICA,
+}

pixeltable/share/protocol/replica.py

@@ -0,0 +1,109 @@
+"""
+Core table protocol for pixeltable operations.
+
+This module contains the core table protocol structures that can be shared
+between pixeltable core and cloud implementations.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+from uuid import UUID
+
+from pydantic import AnyUrl, BaseModel
+
+from pixeltable.catalog.table_version import TableVersionCompleteMd
+
+from .common import PxtUri, RequestBaseModel, StorageDestination
+from .operation_types import ReplicaOperationType
+
+
+class PublishRequest(RequestBaseModel):
+    """Request to publish or push table replica."""
+
+    operation_type: Literal[ReplicaOperationType.PUBLISH_REPLICA] = ReplicaOperationType.PUBLISH_REPLICA
+    table_uri: PxtUri  # If PxtUri#id is not None then it's considered a push replica request
+    pxt_version: str
+    pxt_md_version: int
+    md: list[TableVersionCompleteMd]
+    is_public: bool = False
+    bucket_name: str | None = None  # Optional bucket name, falls back to org's default bucket if not provided
+
+    def get_pxt_uri(self) -> PxtUri:
+        """Get the PxtUri from this request."""
+        return self.table_uri
+
+
+class PublishResponse(BaseModel):
+    """Response from publishing a table replica."""
+
+    upload_id: UUID
+    destination: StorageDestination
+    destination_uri: AnyUrl
+    max_size: int | None = None  # Maximum size that can be used by this replica, used for R2 home buckets
+
+
+class FinalizeRequest(RequestBaseModel):
+    """Request to finalize a table replica."""
+
+    operation_type: Literal[ReplicaOperationType.FINALIZE_REPLICA] = ReplicaOperationType.FINALIZE_REPLICA
+    table_uri: PxtUri  # Use same table_uri that was given during publish replica request
+    upload_id: UUID
+    size: int
+    sha256: str
+    datafile: str
+    row_count: int
+    preview_header: dict[str, str]
+    preview_data: list[list[Any]]
+
+    def get_pxt_uri(self) -> PxtUri:
+        """Get the PxtUri from this request."""
+        return self.table_uri
+
+
+class FinalizeResponse(BaseModel):
+    """Response from finalizing a table replica."""
+
+    confirmed_table_uri: PxtUri
+    version: int | None = None  # Version that was pushed to replica
+
+
+class DeleteRequest(RequestBaseModel):
+    """Request to delete a table replica."""
+
+    operation_type: Literal[ReplicaOperationType.DELETE_REPLICA] = ReplicaOperationType.DELETE_REPLICA
+    table_uri: PxtUri
+    version: int | None = None  # Delete a version in replica
+
+    def get_pxt_uri(self) -> PxtUri:
+        """Get the PxtUri from this request."""
+        return self.table_uri
+
+
+class DeleteResponse(BaseModel):
+    """Response from deleting a table replica."""
+
+    table_uri: PxtUri
+    version: int | None = None
+
+
+class ReplicateRequest(RequestBaseModel):
+    """Request to clone a table replica."""
+
+    operation_type: Literal[ReplicaOperationType.CLONE_REPLICA] = ReplicaOperationType.CLONE_REPLICA
+    table_uri: PxtUri
+
+    def get_pxt_uri(self) -> PxtUri:
+        """Get the PxtUri from this request."""
+        return self.table_uri
+
+
+class ReplicateResponse(BaseModel):
+    """Response from cloning a table replica."""
+
+    table_uri: PxtUri
+    pxt_md_version: int
+    destination: StorageDestination
+    destination_uri: AnyUrl
+    md: list[TableVersionCompleteMd]
+    version: int | None = None