pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/utils/gcs_store.py

@@ -0,0 +1,295 @@
+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
+
+    def create_presigned_url(self, soa: StorageObjectAddress, expiration_seconds: int) -> str:
+        """Create a presigned URL for downloading an object from GCS."""
+        if not soa.has_object:
+            raise excs.Error(f'StorageObjectAddress does not contain an object name: {soa}')
+
+        gcs_client = self.client()
+        bucket = gcs_client.bucket(soa.container)
+        blob = bucket.blob(soa.key)
+
+        presigned_url = blob.generate_signed_url(version='v4', expiration=expiration_seconds, method='GET')
+        return presigned_url
+
+    @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.py

@@ -0,0 +1,133 @@
+import re
+import time
+from http import HTTPStatus
+from random import random
+from typing import Any
+
+_RETRIABLE_ERROR_INDICATORS = (
+    'rate limit',
+    'too many requests',
+    '429',
+    'quota exceeded',
+    'throttled',
+    'rate exceeded',
+    'connection error',
+    'timed out',
+)
+_RETRY_AFTER_PATTERNS = (
+    r'retry after (\d+(?:\.\d+)?)\s*seconds?',
+    r'try again in (\d+(?:\.\d+)?)\s*seconds?',
+    r'wait (\d+(?:\.\d+)?)\s*seconds?',
+    r'retry-after:\s*(\d+(?:\.\d+)?)',
+)
+_RETRIABLE_HTTP_STATUSES: dict[str, int] = {
+    'TOO_MANY_REQUESTS': HTTPStatus.TOO_MANY_REQUESTS.value,
+    'SERVICE_UNAVAILABLE': HTTPStatus.SERVICE_UNAVAILABLE.value,
+    'REQUEST_TIMEOUT': HTTPStatus.REQUEST_TIMEOUT.value,
+    'GATEWAY_TIMEOUT': HTTPStatus.GATEWAY_TIMEOUT.value,
+}
+
+
+def is_retriable_error(exc: Exception) -> tuple[bool, float | None]:
+    """Attempts to guess if the exception indicates a retriable eror. If that is the case, returns True
+    and the retry delay in seconds."""
+
+    # Check for HTTP status TOO_MANY_REQUESTS in various exception classes.
+    # We look for attributes that contain status codes, instead of checking the type of the exception,
+    # in order to handle a wider variety of exception classes.
+    err_md = _extract_error_metadata(exc)
+    if (err_md is None or not err_md[0]) and hasattr(exc, 'response'):
+        err_md = _extract_error_metadata(exc.response)
+
+    if err_md is not None and err_md[0]:
+        retry_after = err_md[1]
+        return err_md[0], retry_after if retry_after is not None and retry_after >= 0 else None
+
+    # Check common rate limit keywords in exception message
+    error_msg = str(exc).lower()
+    if any(indicator in error_msg for indicator in _RETRIABLE_ERROR_INDICATORS):
+        retry_delay = _extract_retry_delay_from_message(error_msg)
+        return True, retry_delay if retry_delay is not None and retry_delay >= 0 else None
+
+    return False, None
+
+
+def _extract_error_metadata(obj: Any) -> tuple[bool, float | None] | None:
+    is_retriable: bool | None = None
+    retry_delay: float | None = None
+    for attr in ['status', 'code', 'status_code']:
+        if hasattr(obj, attr):
+            is_retriable = getattr(obj, attr) in _RETRIABLE_HTTP_STATUSES.values()
+            is_retriable |= str(getattr(obj, attr)).upper() in _RETRIABLE_HTTP_STATUSES
+
+    if hasattr(obj, 'headers'):
+        retry_delay = _extract_retry_delay_from_headers(obj.headers)
+        if retry_delay is not None:
+            is_retriable = True
+
+    return (is_retriable, retry_delay) if is_retriable is not None else None
+
+
+def _extract_retry_delay_from_headers(headers: Any | None) -> float | None:
+    """Extract retry delay from HTTP headers."""
+    if headers is None:
+        return None
+
+    # convert headers to dict-like object for consistent access
+    header_dict: dict
+    if hasattr(headers, 'get'):
+        header_dict = headers
+    else:
+        # headers are a list of tuples or other format
+        try:
+            header_dict = dict(headers)
+        except (TypeError, ValueError):
+            return None
+    # normalize dict keys: lowercase and remove dashes
+    header_dict = {k.lower().replace('-', ''): v for k, v in header_dict.items()}
+
+    # check Retry-After header
+    retry_after = header_dict.get('retryafter')
+    if retry_after is not None:
+        try:
+            return float(retry_after)
+        except (ValueError, TypeError):
+            pass
+
+    # check X-RateLimit-Reset (Unix timestamp)
+    reset_time = header_dict.get('xratelimitreset')
+    if reset_time is not None:
+        try:
+            reset_timestamp = float(reset_time)
+            delay = max(0, reset_timestamp - time.time())
+            return delay
+        except (ValueError, TypeError):
+            pass
+
+    # check X-RateLimit-Reset-After (seconds from now)
+    reset_after = header_dict.get('xratelimitresetafter')
+    if reset_after is not None:
+        try:
+            return float(reset_after)
+        except (ValueError, TypeError):
+            pass
+
+    return None
+
+
+def _extract_retry_delay_from_message(msg: str) -> float | None:
+    msg_lower = msg.lower()
+    for pattern in _RETRY_AFTER_PATTERNS:
+        match = re.search(pattern, msg_lower)
+        if match is not None:
+            try:
+                return float(match.group(1))
+            except (ValueError, TypeError):
+                continue
+    return None
+
+
+def exponential_backoff(attempt: int, base: float = 2.0, max_delay: float = 16.0) -> float:
+    """Generates the retry delay using exponential backoff strategy with jitter. Attempt count starts from 0."""
+    basic_delay = min(max_delay, base**attempt) / 2
+    return basic_delay + random() * basic_delay

pixeltable/utils/http_server.py

@@ -2,14 +2,15 @@ import http
 import http.server
 import logging
 import pathlib
-import urllib
+import urllib.request
+from typing import Any
 
 _logger = logging.getLogger('pixeltable.http.server')
 
 
 def get_file_uri(http_address: str, file_path: str) -> str:
     """Get the URI for a file path, with the given prefix.
-        Used in the client to generate a URI
+    Used in the client to generate a URI
     """
     abs_path = pathlib.Path(file_path)
     assert abs_path.is_absolute()
@@ -19,25 +20,25 @@ def get_file_uri(http_address: str, file_path: str) -> str:
 
 class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
     """Serves all absolute paths, not just the current directory"""
+
     def translate_path(self, path: str) -> str:
         """
-            Translate a /-separated PATH to the local filename syntax.
-            overrides http.server.SimpleHTTPRequestHandler.translate_path
+        Translate a /-separated PATH to the local filename syntax.
+        overrides http.server.SimpleHTTPRequestHandler.translate_path
 
-            This is only useful for file serving.
+        This is only useful for file serving.
 
-            Code initially taken from there:
-            https://github.com/python/cpython/blob/f5406ef454662b98df107775d18ff71ae6849618/Lib/http/server.py#L834
+        Code initially taken from there:
+        https://github.com/python/cpython/blob/f5406ef454662b98df107775d18ff71ae6849618/Lib/http/server.py#L834
         """
         _logger.info(f'translate path {path=}')
         # abandon query parameters, taken from 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, *args) -> None:
+    def log_message(self, format: str, *args: Any) -> None:
         """override logging to stderr in http.server.BaseHTTPRequestHandler"""
         message = format % args
         _logger.info(message.translate(self._control_char_table))  # type: ignore[attr-defined]
@@ -46,7 +47,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
 class LoggingHTTPServer(http.server.ThreadingHTTPServer):
     """Avoids polluting stdout and stderr"""
 
-    def handle_error(self, request, client_address) -> None:
+    def handle_error(self, request, client_address) -> None:  # type: ignore[no-untyped-def]
         """override socketserver.TCPServer.handle_error which prints directly to sys.stderr"""
         import traceback
 
@@ -57,11 +58,11 @@ class LoggingHTTPServer(http.server.ThreadingHTTPServer):
 
 
 def make_server(address: str, port: int) -> http.server.HTTPServer:
-    """Create a file server with pixeltable specific config """
+    """Create a file server with pixeltable specific config"""
     return LoggingHTTPServer((address, port), AbsolutePathHandler)
 
 
 if __name__ == '__main__':
     httpd = make_server('127.0.0.1', 8000)
     print(f'about to server HTTP on {httpd.server_address}')
-    httpd.serve_forever()
+    httpd.serve_forever()

pixeltable/utils/iceberg.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+
+from pyiceberg.catalog.sql import 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.
+    """
+    if isinstance(warehouse_path, str):
+        warehouse_path = Path(warehouse_path)
+    warehouse_path.mkdir(exist_ok=True)
+    return SqlCatalog(name, uri=f'sqlite:///{warehouse_path}/catalog.db', warehouse=f'file://{warehouse_path}')

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')

pixeltable/utils/lancedb.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+from typing import Literal
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.catalog import Catalog
+from pixeltable.env import Env
+
+_logger = logging.getLogger('pixeltable')
+
+
+def export_lancedb(
+    table_or_query: pxt.Table | pxt.Query,
+    db_uri: Path,
+    table_name: str,
+    batch_size_bytes: int = 128 * 2**20,
+    if_exists: Literal['error', 'overwrite', 'append'] = 'error',
+) -> None:
+    """
+    Exports a Query's data to a LanceDB table.
+
+    This utilizes LanceDB's streaming interface for efficient table creation, via a sequence of in-memory pyarrow
+    `RecordBatches`, the size of which can be controlled with the `batch_size_bytes` parameter.
+
+    __Requirements:__
+
+    - `pip install lancedb`
+
+    Args:
+        table_or_query : Table or Query to export.
+        db_uri: Local Path to the LanceDB database.
+        table_name : Name of the table in the LanceDB database.
+        batch_size_bytes : Maximum size in bytes for each batch.
+        if_exists: Determines the behavior if the table already exists. Must be one of the following:
+
+            - `'error'`: raise an error
+            - `'overwrite'`: overwrite the existing table
+            - `'append'`: append to the existing table
+    """
+    Env.get().require_package('lancedb')
+
+    import lancedb  # type: ignore[import-untyped]
+
+    from pixeltable.utils.arrow import to_arrow_schema, to_record_batches
+
+    if if_exists not in ('error', 'overwrite', 'append'):
+        raise excs.Error("export_lancedb(): 'if_exists' must be one of: ['error', 'overwrite', 'append']")
+
+    query: pxt.Query
+    if isinstance(table_or_query, pxt.catalog.Table):
+        query = table_or_query.select()
+    else:
+        query = table_or_query
+
+    db_exists = False
+    if db_uri.exists():
+        if not db_uri.is_dir():
+            raise excs.Error(f"export_lancedb(): '{db_uri!s}' exists and is not a directory")
+        db_exists = True
+
+    try:
+        db = lancedb.connect(str(db_uri))
+        lance_tbl: lancedb.LanceTable | None = None
+        try:
+            lance_tbl = db.open_table(table_name)
+            if if_exists == 'error':
+                raise excs.Error(f'export_lancedb(): table {table_name!r} already exists in {db_uri!r}')
+        except ValueError:
+            # table doesn't exist
+            pass
+
+        with Catalog.get().begin_xact(for_write=False):
+            if lance_tbl is None or if_exists == 'overwrite':
+                mode = 'overwrite' if lance_tbl is not None else 'create'
+                arrow_schema = to_arrow_schema(query.schema)
+                _ = db.create_table(
+                    table_name, to_record_batches(query, batch_size_bytes), schema=arrow_schema, mode=mode
+                )
+            else:
+                lance_tbl.add(to_record_batches(query, batch_size_bytes))
+
+    except Exception as e:
+        # cleanup
+        if not db_exists:
+            shutil.rmtree(db_uri)
+        raise e