pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/utils/dbms.py

@@ -1,6 +1,6 @@
 import abc
 
-from sqlalchemy import URL
+import sqlalchemy as sql
 
 
 class Dbms(abc.ABC):
@@ -11,9 +11,9 @@ class Dbms(abc.ABC):
     name: str
     transaction_isolation_level: str
     version_index_type: str
-    db_url: URL
+    db_url: sql.URL
 
-    def __init__(self, name: str, transaction_isolation_level: str, version_index_type: str, db_url: URL) -> None:
+    def __init__(self, name: str, transaction_isolation_level: str, version_index_type: str, db_url: sql.URL) -> None:
         self.name = name
         self.transaction_isolation_level = transaction_isolation_level
         self.version_index_type = version_index_type
@@ -28,13 +28,18 @@ class Dbms(abc.ABC):
     @abc.abstractmethod
     def default_system_db_url(self) -> str: ...
 
+    @abc.abstractmethod
+    def create_vector_index_stmt(
+        self, store_index_name: str, sa_value_col: sql.Column, metric: str
+    ) -> sql.Compiled: ...
+
 
 class PostgresqlDbms(Dbms):
     """
     Implements utilities to interact with Postgres database.
     """
 
-    def __init__(self, db_url: URL):
+    def __init__(self, db_url: sql.URL):
         super().__init__('postgresql', 'SERIALIZABLE', 'brin', db_url)
 
     def drop_db_stmt(self, database: str) -> str:
@@ -47,13 +52,25 @@ class PostgresqlDbms(Dbms):
         a = self.db_url.set(database='postgres').render_as_string(hide_password=False)
         return a
 
+    def create_vector_index_stmt(self, store_index_name: str, sa_value_col: sql.Column, metric: str) -> sql.Compiled:
+        from sqlalchemy.dialects import postgresql
+
+        sa_idx = sql.Index(
+            store_index_name,
+            sa_value_col,
+            postgresql_using='hnsw',
+            postgresql_with={'m': 16, 'ef_construction': 64},
+            postgresql_ops={sa_value_col.name: metric},
+        )
+        return sql.schema.CreateIndex(sa_idx, if_not_exists=True).compile(dialect=postgresql.dialect())
+
 
 class CockroachDbms(Dbms):
     """
     Implements utilities to interact with CockroachDb database.
     """
 
-    def __init__(self, db_url: URL):
+    def __init__(self, db_url: sql.URL):
         super().__init__('cockroachdb', 'SERIALIZABLE', 'btree', db_url)
 
     def drop_db_stmt(self, database: str) -> str:
@@ -64,3 +81,12 @@ class CockroachDbms(Dbms):
 
     def default_system_db_url(self) -> str:
         return self.db_url.set(database='defaultdb').render_as_string(hide_password=False)
+
+    def sa_vector_index(self, store_index_name: str, sa_value_col: sql.schema.Column, metric: str) -> sql.Index | None:
+        return None
+
+    def create_vector_index_stmt(self, store_index_name: str, sa_value_col: sql.Column, metric: str) -> sql.Compiled:
+        return sql.text(
+            f'CREATE VECTOR INDEX IF NOT EXISTS {store_index_name} ON {sa_value_col.table.name}'
+            f'({sa_value_col.name} {metric})'
+        ).compile()

pixeltable/utils/description_helper.py

@@ -1,5 +1,4 @@
 import dataclasses
-from typing import Optional, Union
 
 import pandas as pd
 from pandas.io.formats.style import Styler
@@ -7,11 +6,11 @@ from pandas.io.formats.style import Styler
 
 @dataclasses.dataclass
 class _Descriptor:
-    body: Union[str, pd.DataFrame]
+    body: str | pd.DataFrame
     # The remaining fields only affect the behavior if `body` is a pd.DataFrame.
     show_index: bool
     show_header: bool
-    styler: Optional[Styler] = None
+    styler: Styler | None = None
 
 
 class DescriptionHelper:
@@ -33,10 +32,10 @@ class DescriptionHelper:
 
     def append(
         self,
-        descriptor: Union[str, pd.DataFrame],
+        descriptor: str | pd.DataFrame,
         show_index: bool = False,
         show_header: bool = True,
-        styler: Optional[Styler] = None,
+        styler: Styler | None = None,
     ) -> None:
         self.__descriptors.append(_Descriptor(descriptor, show_index, show_header, styler))
 

pixeltable/utils/documents.py

@@ -1,6 +1,5 @@
 import dataclasses
 import os
-from typing import Optional
 
 import bs4
 import fitz  # type: ignore[import-untyped]
@@ -13,10 +12,10 @@ from pixeltable.env import Env
 @dataclasses.dataclass
 class DocumentHandle:
     format: ts.DocumentType.DocumentFormat
-    bs_doc: Optional[bs4.BeautifulSoup] = None
-    md_ast: Optional[dict] = None
-    pdf_doc: Optional[fitz.Document] = None
-    txt_doc: Optional[str] = None
+    bs_doc: bs4.BeautifulSoup | None = None
+    md_ast: dict | None = None
+    pdf_doc: fitz.Document | None = None
+    txt_doc: str | None = None
 
 
 def get_document_handle(path: str) -> DocumentHandle:
@@ -34,7 +33,7 @@ def get_document_handle(path: str) -> DocumentHandle:
     raise excs.Error(f'Unrecognized document format: {path}')
 
 
-def get_handle_by_extension(path: str, extension: str) -> Optional[DocumentHandle]:
+def get_handle_by_extension(path: str, extension: str) -> DocumentHandle | None:
     doc_format = ts.DocumentType.DocumentFormat.from_extension(extension)
 
     try:

pixeltable/utils/exception_handler.py

@@ -1,35 +1,12 @@
 import logging
-import sys
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Callable, TypeVar
 
 R = TypeVar('R')
 
-
-def _is_in_exception() -> bool:
-    """
-    Check if code is currently executing within an exception context.
-    """
-    current_exception = sys.exc_info()[1]
-    return current_exception is not None
-
-
-def run_cleanup_on_exception(cleanup_func: Callable[..., R], *args: Any, **kwargs: Any) -> Optional[R]:
-    """
-    Runs cleanup only when running in exception context.
-
-    The function `run_cleanup_on_exception()` should be used to clean up resources when an operation fails.
-    This is typically done using a try, except, and finally block, with the resource cleanup logic placed within
-    the except block. However, this pattern may not handle KeyboardInterrupt exceptions.
-    To ensure that resources are always cleaned up at least once when an exception or KeyboardInterrupt occurs,
-    create an idempotent function for cleaning up resources and pass it to the `run_cleanup_on_exception()` function
-    from the finally block.
-    """
-    if _is_in_exception():
-        return run_cleanup(cleanup_func, *args, raise_error=False, **kwargs)
-    return None
+logger = logging.getLogger('pixeltable')
 
 
-def run_cleanup(cleanup_func: Callable[..., R], *args: Any, raise_error: bool = True, **kwargs: Any) -> Optional[R]:
+def run_cleanup(cleanup_func: Callable[..., R], *args: Any, raise_error: bool = True, **kwargs: Any) -> R | None:
     """
     Runs a cleanup function. If interrupted, retry cleanup.
     The `run_cleanup()` function ensures that the `cleanup_func()` function executes at least once.
@@ -40,20 +17,20 @@ def run_cleanup(cleanup_func: Callable[..., R], *args: Any, raise_error: bool =
         raise_error: raise an exception if an error occurs during cleanup.
     """
     try:
-        logging.debug(f'Running cleanup function: {cleanup_func.__name__!r}')
+        logger.debug(f'Running cleanup function: {cleanup_func.__name__!r}')
         return cleanup_func(*args, **kwargs)
     except KeyboardInterrupt as interrupt:
         # Save original exception and re-attempt cleanup
         original_exception = interrupt
-        logging.debug(f'Cleanup {cleanup_func.__name__!r} interrupted, retrying')
+        logger.debug(f'Cleanup {cleanup_func.__name__!r} interrupted, retrying')
         try:
             return cleanup_func(*args, **kwargs)
         except Exception as e:
             # Suppress this exception
-            logging.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e}')
+            logger.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e.__class__}: {e}')
         raise KeyboardInterrupt from original_exception
     except Exception as e:
-        logging.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e}')
+        logger.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e.__class__}: {e}')
         if raise_error:
             raise e
     return None

pixeltable/utils/filecache.py

@@ -9,7 +9,7 @@ from collections import OrderedDict, defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 from uuid import UUID
 
 import pixeltable.exceptions as excs
@@ -58,7 +58,7 @@ class FileCache:
     - implement MRU eviction for queries that exceed the capacity
     """
 
-    __instance: Optional[FileCache] = None
+    __instance: FileCache | None = None
 
     cache: OrderedDict[str, CacheEntry]
     total_size: int
@@ -126,12 +126,12 @@ class FileCache:
             return 0
         return int(self.total_size / len(self.cache))
 
-    def num_files(self, tbl_id: Optional[UUID] = None) -> int:
+    def num_files(self, tbl_id: UUID | None = None) -> int:
         if tbl_id is None:
             return len(self.cache)
         return sum(e.tbl_id == tbl_id for e in self.cache.values())
 
-    def clear(self, tbl_id: Optional[UUID] = None) -> None:
+    def clear(self, tbl_id: UUID | None = None) -> None:
         """
         For testing purposes: allow resetting capacity and stats.
         """
@@ -174,7 +174,7 @@ class FileCache:
         h.update(url.encode())
         return h.hexdigest()
 
-    def lookup(self, url: str) -> Optional[Path]:
+    def lookup(self, url: str) -> Path | None:
         self.num_requests += 1
         key = self._url_hash(url)
         entry = self.cache.get(key, None)
@@ -214,7 +214,7 @@ class FileCache:
         new_path = entry.path
         os.rename(str(path), str(new_path))
         new_path.touch(exist_ok=True)
-        _logger.debug(f'added entry for cell {url} to file cache')
+        _logger.debug(f'FileCache: cached url {url} with file name {new_path}')
         return new_path
 
     def ensure_capacity(self, size: int) -> None:

pixeltable/utils/formatter.py

@@ -4,7 +4,7 @@ import io
 import json
 import logging
 import mimetypes
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 import av
 import numpy as np
@@ -39,7 +39,7 @@ class Formatter:
         self.__num_cols = num_cols
         self.__http_address = http_address
 
-    def get_pandas_formatter(self, col_type: ts.ColumnType) -> Optional[Callable]:
+    def get_pandas_formatter(self, col_type: ts.ColumnType) -> Callable | None:
         if col_type.is_string_type():
             return self.format_string
         if col_type.is_float_type():
@@ -184,7 +184,7 @@ class Formatter:
         """
 
     @classmethod
-    def extract_first_video_frame(cls, file_path: str) -> Optional[Image.Image]:
+    def extract_first_video_frame(cls, file_path: str) -> Image.Image | None:
         with av.open(file_path) as container:
             try:
                 img = next(container.decode(video=0)).to_image()
@@ -224,9 +224,7 @@ class Formatter:
         """
 
     @classmethod
-    def make_document_thumbnail(
-        cls, file_path: str, max_width: int = 320, max_height: int = 320
-    ) -> Optional[Image.Image]:
+    def make_document_thumbnail(cls, file_path: str, max_width: int = 320, max_height: int = 320) -> Image.Image | None:
         """
         Returns a thumbnail image of a document.
         """

pixeltable/utils/gcs_store.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import logging
+import re
+import urllib.parse
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Iterator
+
+from google.api_core.exceptions import GoogleAPIError
+from google.cloud import storage  # type: ignore[attr-defined]
+from google.cloud.exceptions import Forbidden, NotFound
+from google.cloud.storage.client import Client  # type: ignore[import-untyped]
+
+from pixeltable import env, exceptions as excs
+from pixeltable.utils.object_stores import ObjectPath, ObjectStoreBase, StorageObjectAddress, StorageTarget
+
+if TYPE_CHECKING:
+    from pixeltable.catalog import Column
+
+_logger = logging.getLogger('pixeltable')
+
+
+@env.register_client('gcs_store')
+def _() -> 'Client':
+    """Create and return a GCS client, using default credentials if available,
+    otherwise creating an anonymous client for public buckets.
+    """
+    try:
+        # Create a client with default credentials
+        # Note that if the default credentials have expired, gcloud will still create a client,
+        # which will report the expiry error when it is used.
+        # To create and use an anonymous client, expired credentials must be removed.
+        # For application default credentials, delete the file in ~/.config/gcloud/, or
+        #   gcloud auth application-default revoke
+        # OR
+        # For service account keys, you must delete the downloaded key file.
+        client = storage.Client()
+        return client
+    except Exception:
+        # If no credentials are found, create an anonymous client which can be used for public buckets.
+        client = storage.Client.create_anonymous_client()
+        return client
+
+
+class GCSStore(ObjectStoreBase):
+    """Class to handle Google Cloud Storage operations."""
+
+    # URI of the GCS bucket in the format gs://bucket_name/prefix/
+    # Always ends with a slash
+    __base_uri: str
+
+    # bucket name extracted from the URI
+    __bucket_name: str
+
+    # prefix path within the bucket, either empty or ending with a slash
+    __prefix_name: str
+
+    # The parsed form of the given destination address
+    soa: StorageObjectAddress
+
+    def __init__(self, soa: StorageObjectAddress):
+        assert soa.storage_target == StorageTarget.GCS_STORE, f'Expected storage_target "gs", got {soa.storage_target}'
+        self.soa = soa
+        self.__base_uri = soa.prefix_free_uri + soa.prefix
+        self.__bucket_name = soa.container
+        self.__prefix_name = soa.prefix
+
+    @classmethod
+    def client(cls) -> 'Client':
+        """Return the GCS client."""
+        return env.Env.get().get_client('gcs_store')
+
+    @property
+    def bucket_name(self) -> str:
+        """Return the bucket name from the base URI."""
+        return self.__bucket_name
+
+    @property
+    def prefix(self) -> str:
+        """Return the prefix from the base URI."""
+        return self.__prefix_name
+
+    def validate(self, error_col_name: str) -> str | None:
+        """
+        Checks if the URI exists.
+
+        Returns:
+            str: The base URI if the GCS bucket exists and is accessible, None otherwise.
+        """
+        try:
+            client = self.client()
+            bucket = client.bucket(self.bucket_name)
+            blobs = bucket.list_blobs(max_results=1)
+            # This will raise an exception if the destination doesn't exist or cannot be listed
+            _ = list(blobs)  # Force evaluation to check access
+            return self.__base_uri
+        except (NotFound, Forbidden, GoogleAPIError) as e:
+            self.handle_gcs_error(e, self.bucket_name, f'validate bucket {error_col_name}')
+        return None
+
+    def _prepare_uri_raw(self, tbl_id: uuid.UUID, col_id: int, tbl_version: int, ext: str | None = None) -> str:
+        """
+        Construct a new, unique URI for a persisted media file.
+        """
+        prefix, filename = ObjectPath.create_prefix_raw(tbl_id, col_id, tbl_version, ext)
+        parent = f'{self.__base_uri}{prefix}'
+        return f'{parent}/{filename}'
+
+    def _prepare_uri(self, col: Column, ext: str | None = None) -> str:
+        """
+        Construct a new, unique URI for a persisted media file.
+        """
+        assert col.get_tbl() is not None, 'Column must be associated with a table'
+        return self._prepare_uri_raw(col.get_tbl().id, col.id, col.get_tbl().version, ext=ext)
+
+    def copy_local_file(self, col: Column, src_path: Path) -> str:
+        """Copy a local file, and return its new URL"""
+        new_file_uri = self._prepare_uri(col, ext=src_path.suffix)
+        parsed = urllib.parse.urlparse(new_file_uri)
+        blob_name = parsed.path.lstrip('/')
+
+        try:
+            client = self.client()
+            bucket = client.bucket(self.bucket_name)
+            blob = bucket.blob(blob_name)
+            blob.upload_from_filename(str(src_path))
+            _logger.debug(f'Media Storage: copied {src_path} to {new_file_uri}')
+            return new_file_uri
+        except GoogleAPIError as e:
+            self.handle_gcs_error(e, self.bucket_name, f'upload file {src_path}')
+            raise
+
+    def copy_object_to_local_file(self, src_path: str, dest_path: Path) -> None:
+        """Copies an object to a local file. Thread safe"""
+        try:
+            client = self.client()
+            bucket = client.bucket(self.bucket_name)
+            blob = bucket.blob(self.prefix + src_path)
+            blob.download_to_filename(str(dest_path))
+        except GoogleAPIError as e:
+            self.handle_gcs_error(e, self.bucket_name, f'download file {src_path}')
+            raise
+
+    def _get_filtered_objects(self, bucket: Any, tbl_id: uuid.UUID, tbl_version: int | None = None) -> Iterator:
+        """Private method to get filtered objects for a table, optionally filtered by version.
+
+        Args:
+            tbl_id: Table UUID to filter by
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Tuple of (iterator over GCS objects matching the criteria, bucket object)
+        """
+        table_prefix = ObjectPath.table_prefix(tbl_id)
+        prefix = f'{self.prefix}{table_prefix}/'
+
+        if tbl_version is None:
+            # Return all blobs with the table prefix
+            blob_iterator = bucket.list_blobs(prefix=prefix)
+        else:
+            # Filter by both table_id and table_version using the ObjectPath pattern
+            # Pattern: tbl_id_col_id_version_uuid
+            version_pattern = re.compile(rf'{re.escape(table_prefix)}_\d+_{re.escape(str(tbl_version))}_[0-9a-fA-F]+.*')
+            # Return filtered collection - this still uses lazy loading
+            all_blobs = bucket.list_blobs(prefix=prefix)
+            blob_iterator = (blob for blob in all_blobs if version_pattern.match(blob.name.split('/')[-1]))
+
+        return blob_iterator
+
+    def count(self, tbl_id: uuid.UUID, tbl_version: int | None = None) -> int:
+        """Count the number of files belonging to tbl_id. If tbl_version is not None,
+        count only those files belonging to the specified tbl_version.
+
+        Args:
+            tbl_id: Table UUID to count objects for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of objects matching the criteria
+        """
+        assert tbl_id is not None
+
+        try:
+            client = self.client()
+            bucket = client.bucket(self.bucket_name)
+
+            blob_iterator = self._get_filtered_objects(bucket, tbl_id, tbl_version)
+
+            return sum(1 for _ in blob_iterator)
+
+        except GoogleAPIError as e:
+            self.handle_gcs_error(e, self.bucket_name, f'setup iterator {self.prefix}')
+            raise
+
+    def delete(self, tbl_id: uuid.UUID, tbl_version: int | None = None) -> int:
+        """Delete all files belonging to tbl_id. If tbl_version is not None, delete
+        only those files belonging to the specified tbl_version.
+
+        Args:
+            tbl_id: Table UUID to delete objects for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of objects deleted
+        """
+        assert tbl_id is not None
+        total_deleted = 0
+
+        try:
+            client = self.client()
+            bucket = client.bucket(self.bucket_name)
+            blob_iterator = self._get_filtered_objects(bucket, tbl_id, tbl_version)
+
+            # Collect blob names for batch deletion
+            blobs_to_delete = []
+
+            for blob in blob_iterator:
+                blobs_to_delete.append(blob)
+
+                # Process in batches for efficiency
+                if len(blobs_to_delete) >= 100:
+                    with client.batch():
+                        for b in blobs_to_delete:
+                            b.delete()
+                    total_deleted += len(blobs_to_delete)
+                    blobs_to_delete = []
+
+            # Delete any remaining blobs in the final batch
+            if len(blobs_to_delete) > 0:
+                with client.batch():
+                    for b in blobs_to_delete:
+                        b.delete()
+                total_deleted += len(blobs_to_delete)
+
+            return total_deleted
+
+        except GoogleAPIError as e:
+            self.handle_gcs_error(e, self.bucket_name, f'deleting with {self.prefix}')
+            raise
+
+    def list_objects(self, return_uri: bool, n_max: int = 10) -> list[str]:
+        """Return a list of objects found in the specified destination bucket.
+        Each returned object includes the full set of prefixes.
+        if return_uri is True, full URI's are returned; otherwise, just the object keys.
+        """
+        p = self.soa.prefix_free_uri if return_uri else ''
+        gcs_client = self.client()
+        r: list[str] = []
+
+        try:
+            bucket = gcs_client.bucket(self.bucket_name)
+            # List blobs with the given prefix, limiting to n_max
+            blobs = bucket.list_blobs(prefix=self.prefix, max_results=n_max)
+
+            for blob in blobs:
+                r.append(f'{p}{blob.name}')
+                if len(r) >= n_max:
+                    break
+
+        except GoogleAPIError as e:
+            self.handle_gcs_error(e, self.bucket_name, f'list objects from {self.prefix}')
+        return r
+
+    @classmethod
+    def handle_gcs_error(cls, e: Exception, bucket_name: str, operation: str = '', *, ignore_404: bool = False) -> None:
+        """Handle GCS-specific errors and convert them to appropriate exceptions"""
+        if isinstance(e, NotFound):
+            if ignore_404:
+                return
+            raise excs.Error(f'Bucket or object {bucket_name} not found during {operation}: {str(e)!r}')
+        elif isinstance(e, Forbidden):
+            raise excs.Error(f'Access denied to bucket {bucket_name} during {operation}: {str(e)!r}')
+        elif isinstance(e, GoogleAPIError):
+            # Handle other Google API errors
+            error_message = str(e)
+            if 'Precondition' in error_message:
+                raise excs.Error(f'Precondition failed for bucket {bucket_name} during {operation}: {error_message}')
+            else:
+                raise excs.Error(f'Error during {operation} in bucket {bucket_name}: {error_message}')
+        else:
+            # Generic error handling
+            raise excs.Error(f'Unexpected error during {operation} in bucket {bucket_name}: {str(e)!r}')

pixeltable/utils/http_server.py

@@ -2,7 +2,7 @@ import http
 import http.server
 import logging
 import pathlib
-import urllib
+import urllib.request
 from typing import Any
 
 _logger = logging.getLogger('pixeltable.http.server')
@@ -36,8 +36,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
         path = path.split('?', 1)[0]
         path = path.split('#', 1)[0]
 
-        path = pathlib.Path(urllib.request.url2pathname(path))
-        return str(path)
+        return str(pathlib.Path(urllib.request.url2pathname(path)))
 
     def log_message(self, format: str, *args: Any) -> None:
         """override logging to stderr in http.server.BaseHTTPRequestHandler"""

pixeltable/utils/iceberg.py

@@ -1,10 +1,9 @@
 from pathlib import Path
-from typing import Union
 
 from pyiceberg.catalog.sql import SqlCatalog
 
 
-def sqlite_catalog(warehouse_path: Union[str, Path], name: str = 'pixeltable') -> SqlCatalog:
+def sqlite_catalog(warehouse_path: str | Path, name: str = 'pixeltable') -> SqlCatalog:
     """
     Instantiate a sqlite Iceberg catalog at the specified path. If no catalog exists, one will be created.
     """

pixeltable/utils/image.py

@@ -0,0 +1,17 @@
+import base64
+from io import BytesIO
+
+import PIL.Image
+
+
+def default_format(img: PIL.Image.Image) -> str:
+    # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
+    # In that case, use WebP instead.
+    return 'webp' if img.has_transparency_data else 'jpeg'
+
+
+def to_base64(image: PIL.Image.Image, format: str | None = None) -> str:
+    buffer = BytesIO()
+    image.save(buffer, format=format or image.format)
+    image_bytes = buffer.getvalue()
+    return base64.b64encode(image_bytes).decode('utf-8')