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

pixeltable/utils/azure_store.py

@@ -0,0 +1,305 @@
+import logging
+import re
+import threading
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Iterator
+
+from azure.core.exceptions import AzureError
+
+from pixeltable import env, exceptions as excs
+from pixeltable.config import Config
+from pixeltable.utils.object_stores import ObjectPath, ObjectStoreBase, StorageObjectAddress
+
+if TYPE_CHECKING:
+    from azure.storage.blob import BlobProperties, BlobServiceClient
+
+    from pixeltable.catalog import Column
+
+
+_logger = logging.getLogger('pixeltable')
+
+
+client_lock = threading.Lock()
+
+
+@env.register_client('azure_blob')
+def _() -> dict[str, 'BlobServiceClient']:
+    return {}
+
+
+class AzureBlobStore(ObjectStoreBase):
+    """Class to handle Azure Blob Storage operations."""
+
+    # TODO: This needs to be redesigned to use asyncio.
+
+    # URI of the Azure Blob Storage container
+    # Always ends with a slash
+    __base_uri: str
+
+    # Storage account name
+    __account_name: str
+
+    # Container name extracted from the URI
+    __container_name: str
+
+    # Prefix path within the container, either empty or ending with a slash
+    __prefix_name: str
+
+    # URI scheme (wasb, wasbs, abfs, abfss, https)
+    __scheme: str
+
+    soa: StorageObjectAddress
+
+    def __init__(self, soa: StorageObjectAddress):
+        self.soa = soa
+        self.__scheme = soa.scheme
+        self.__account_name = soa.account
+        self.__container_name = soa.container
+        self.__prefix_name = soa.prefix
+
+        # Reconstruct base URI in normalized format
+        self.__base_uri = self.soa.prefix_free_uri + self.__prefix_name
+        _logger.info(
+            f'Initialized AzureBlobStore with base URI: {self.__base_uri},',
+            f'account: {self.__account_name}, container: {self.__container_name}, prefix: {self.__prefix_name}',
+        )
+
+    def client(self) -> 'BlobServiceClient':
+        """Return the Azure Blob Storage client."""
+        client_dict: dict[str, 'BlobServiceClient'] = env.Env.get().get_client('azure_blob')
+        with client_lock:
+            uri = self.soa.container_free_uri
+            if uri not in client_dict:
+                storage_account_name = Config.get().get_string_value('storage_account_name', section='azure')
+                storage_account_key = Config.get().get_string_value('storage_account_key', section='azure')
+                if (storage_account_name is None) != (storage_account_key is None):
+                    raise excs.Error(
+                        "Azure 'storage_account_name' and 'storage_account_key' must be specified together."
+                    )
+                if storage_account_name is None or storage_account_name != self.__account_name:
+                    # Attempt a connection to a public resource, with no account key
+                    client_dict[uri] = self.create_client(endpoint_url=uri)
+                else:
+                    client_dict[uri] = self.create_client(
+                        endpoint_url=uri, account_name=self.__account_name, account_key=storage_account_key
+                    )
+            return client_dict[uri]
+
+    @property
+    def account_name(self) -> str:
+        """Return the storage account name."""
+        return self.__account_name
+
+    @property
+    def container_name(self) -> str:
+        """Return the container name from the base URI."""
+        return self.__container_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 and is accessible.
+
+        Returns:
+            str: The base URI if the container exists and is accessible, None otherwise.
+        """
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+            # Check if container exists by trying to get its properties
+            container_client.get_container_properties()
+            return self.__base_uri
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'validate container {error_col_name}')
+        return None
+
+    def copy_object_to_local_file(self, src_path: str, dest_path: Path) -> None:
+        """Copies a blob to a local file. Thread safe."""
+        try:
+            blob_client = self.client().get_blob_client(container=self.container_name, blob=self.prefix + src_path)
+            with open(dest_path, 'wb') as download_file:
+                download_stream = blob_client.download_blob()
+                download_file.write(download_stream.readall())
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'download file {src_path}')
+            raise
+
+    # TODO: utils package should not include back-references to `Column`
+    def copy_local_file(self, col: 'Column', src_path: Path) -> str:
+        """Copy a local file to Azure Blob Storage, and return its new URL"""
+        prefix, filename = ObjectPath.create_prefix_raw(
+            col.get_tbl().id, col.id, col.get_tbl().version, ext=src_path.suffix
+        )
+        blob_name = f'{self.prefix}{prefix}/{filename}'
+        new_file_uri = f'{self.__base_uri}{prefix}/{filename}'
+
+        try:
+            blob_client = self.client().get_blob_client(container=self.container_name, blob=blob_name)
+            with open(src_path, 'rb') as data:
+                blob_client.upload_blob(data, overwrite=True)
+            _logger.debug(f'Media Storage: copied {src_path} to {new_file_uri}')
+            return new_file_uri
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'upload file {src_path}')
+            raise
+
+    def _get_filtered_blobs(
+        self, tbl_id: uuid.UUID | None, tbl_version: int | None = None
+    ) -> Iterator['BlobProperties']:
+        """Private method to get filtered blobs for a table, optionally filtered by version.
+
+        Args:
+            tbl_id: Table UUID to filter by
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Iterator over blob objects matching the criteria
+        """
+        # Use ObjectPath to construct the prefix for this table
+        if tbl_id is None:
+            prefix = self.prefix
+            assert tbl_version is None, 'tbl_version must be None if tbl_id is None'
+        else:
+            table_prefix = ObjectPath.table_prefix(tbl_id)
+            prefix = f'{self.prefix}{table_prefix}/'
+
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+
+            blob_iterator: Iterator['BlobProperties']
+            if tbl_version is None:
+                # Return all blobs with the table prefix
+                blob_iterator = container_client.list_blobs(name_starts_with=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]+.*'
+                )
+                # Get all blobs with the prefix and filter by version pattern
+                all_blobs = container_client.list_blobs(name_starts_with=prefix)
+                blob_iterator = (blob for blob in all_blobs if version_pattern.match(blob.name.split('/')[-1]))
+
+            return blob_iterator
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_name, f'setup iterator {self.prefix}')
+            raise
+
+    def count(self, tbl_id: uuid.UUID | None, 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 blobs for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of blobs matching the criteria
+        """
+        blob_iterator = self._get_filtered_blobs(tbl_id, tbl_version)
+        return sum(1 for _ in blob_iterator)
+
+    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 blobs for
+            tbl_version: Optional table version to filter by
+
+        Returns:
+            Number of blobs deleted
+        """
+        assert tbl_id is not None
+        blob_iterator = self._get_filtered_blobs(tbl_id, tbl_version)
+        total_deleted = 0
+
+        try:
+            container_client = self.client().get_container_client(self.container_name)
+
+            for blob in blob_iterator:
+                # TODO: Figure out now to properly use batch method delete_blobs(), it doesn't seem to work properly
+                container_client.delete_blob(blob.name)
+                total_deleted += 1
+
+            # print(f"Deleted {total_deleted} blobs from container '{self.container_name}'.")
+            return total_deleted
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.container_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 ''
+        r: list[str] = []
+        try:
+            blob_iterator = self._get_filtered_blobs(tbl_id=None, tbl_version=None)
+            for blob in blob_iterator:
+                r.append(f'{p}{blob.name}')
+                if len(r) >= n_max:
+                    return r
+
+        except AzureError as e:
+            self.handle_azure_error(e, self.__container_name, f'list objects from {self.__base_uri}')
+        return r
+
+    @classmethod
+    def handle_azure_error(
+        cls, e: 'AzureError', container_name: str, operation: str = '', *, ignore_404: bool = False
+    ) -> None:
+        from azure.core.exceptions import ClientAuthenticationError, HttpResponseError, ResourceNotFoundError
+
+        if ignore_404 and isinstance(e, ResourceNotFoundError):
+            return
+
+        if isinstance(e, ResourceNotFoundError):
+            raise excs.Error(f'Container {container_name} or blob not found during {operation}: {str(e)!r}')
+        elif isinstance(e, ClientAuthenticationError):
+            raise excs.Error(f'Authentication failed for container {container_name} during {operation}: {str(e)!r}')
+        elif isinstance(e, HttpResponseError):
+            if e.status_code == 403:
+                raise excs.Error(f'Access denied to container {container_name} during {operation}: {str(e)!r}')
+            elif e.status_code == 412:
+                raise excs.Error(f'Precondition failed for container {container_name} during {operation}: {str(e)!r}')
+            else:
+                raise excs.Error(
+                    f'HTTP error during {operation} in container {container_name}: {e.status_code} - {str(e)!r}'
+                )
+        else:
+            raise excs.Error(f'Error during {operation} in container {container_name}: {str(e)!r}')
+
+    @classmethod
+    def create_client(
+        cls, endpoint_url: str, account_name: str | None = None, account_key: str | None = None
+    ) -> 'BlobServiceClient':
+        from azure.core.credentials import AzureNamedKeyCredential
+        from azure.storage.blob import BlobServiceClient  # TODO: Use azure.storage.blob.aio instead
+
+        assert (account_name is None) == (account_key is None)
+        try:
+            #  e.g. endpoint_url: str = f'https://{account_name}.blob.core.windows.net'
+            assert endpoint_url is not None, 'No Azure Storage account information provided'
+
+            # Use empty SAS token for anonymous authentication
+            credential = None
+            if account_name is not None:
+                credential = AzureNamedKeyCredential(name=account_name, key=account_key)
+            return BlobServiceClient(
+                account_url=endpoint_url,
+                credential=credential,
+                max_single_get_size=(32 * 2**20),
+                max_chunk_get_size=(4 * 2**20),
+                connection_timeout=15,
+                read_timeout=30,
+            )
+        except Exception as e:
+            raise excs.Error(f'Failed to create Azure Blob Storage client: {str(e)!r}') from e

pixeltable/utils/code.py

@@ -1,12 +1,11 @@
 import types
-from typing import Optional
 
 from pixeltable.func import Function
 
 # Utilities related to the organization of the Pixeltable codebase.
 
 
-def local_public_names(mod_name: str, exclude: Optional[list[str]] = None) -> list[str]:
+def local_public_names(mod_name: str, exclude: list[str] | None = None) -> list[str]:
     """
     Returns a list of all functions and submodules that are local to the specified module and are
     publicly accessible. Intended to facilitate implementation of module __dir__() methods for

pixeltable/utils/dbms.py

@@ -29,9 +29,7 @@ class Dbms(abc.ABC):
     def default_system_db_url(self) -> str: ...
 
     @abc.abstractmethod
-    def create_vector_index(
-        self, index_name: str, index_value_sa_col: sql.schema.Column, conn: sql.Connection, metric: str
-    ) -> None: ...
+    def sa_vector_index(self, store_index_name: str, sa_value_col: sql.schema.Column, metric: str) -> sql.Index: ...
 
 
 class PostgresqlDbms(Dbms):
@@ -52,17 +50,14 @@ class PostgresqlDbms(Dbms):
         a = self.db_url.set(database='postgres').render_as_string(hide_password=False)
         return a
 
-    def create_vector_index(
-        self, index_name: str, index_value_sa_col: sql.schema.Column, conn: sql.Connection, metric: str
-    ) -> None:
-        idx = sql.Index(
-            index_name,
-            index_value_sa_col,
+    def sa_vector_index(self, store_index_name: str, sa_value_col: sql.schema.Column, metric: str) -> sql.Index:
+        return sql.Index(
+            store_index_name,
+            sa_value_col,
             postgresql_using='hnsw',
             postgresql_with={'m': 16, 'ef_construction': 64},
-            postgresql_ops={index_value_sa_col.name: metric},
+            postgresql_ops={sa_value_col.name: metric},
         )
-        idx.create(bind=conn)
 
 
 class CockroachDbms(Dbms):
@@ -82,11 +77,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 create_vector_index(
-        self, index_name: str, index_value_sa_col: sql.schema.Column, conn: sql.Connection, metric: str
-    ) -> None:
-        create_index_sql = sql.text(
-            f"""CREATE VECTOR INDEX {index_name} ON {index_value_sa_col.table.name}
-             ({index_value_sa_col.name} {metric})"""
-        )
-        conn.execute(create_index_sql)
+    def sa_vector_index(self, store_index_name: str, sa_value_col: sql.schema.Column, metric: str) -> sql.Index:
+        # TODO: can the Create Index statement be generated via sqlalchemy?
+        # if not, change this method to create_vector_index_stmt(...) -> str
+        # original code:
+        # create_index_sql = sql.text(
+        #     f"""CREATE VECTOR INDEX {store_index_name} ON {sa_value_col.table.name}
+        #      ({sa_value_col.name} {metric})"""
+        # )
+        return None

pixeltable/utils/description_helper.py

@@ -1,5 +1,4 @@
 import dataclasses
-from typing import Optional
 
 import pandas as pd
 from pandas.io.formats.style import Styler
@@ -11,7 +10,7 @@ class _Descriptor:
     # 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:
@@ -36,7 +35,7 @@ class DescriptionHelper:
         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,12 +1,12 @@
 import logging
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Callable, TypeVar
 
 R = TypeVar('R')
 
 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.

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)

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

@@ -5,7 +5,7 @@ import re
 import urllib.parse
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterator, Optional
+from typing import TYPE_CHECKING, Any, Iterator
 
 from google.api_core.exceptions import GoogleAPIError
 from google.cloud import storage  # type: ignore[attr-defined]
@@ -81,7 +81,7 @@ class GCSStore(ObjectStoreBase):
         """Return the prefix from the base URI."""
         return self.__prefix_name
 
-    def validate(self, error_col_name: str) -> Optional[str]:
+    def validate(self, error_col_name: str) -> str | None:
         """
         Checks if the URI exists.
 
@@ -99,7 +99,7 @@ class GCSStore(ObjectStoreBase):
             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: Optional[str] = None) -> str:
+    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.
         """
@@ -107,12 +107,12 @@ class GCSStore(ObjectStoreBase):
         parent = f'{self.__base_uri}{prefix}'
         return f'{parent}/{filename}'
 
-    def _prepare_uri(self, col: Column, ext: Optional[str] = None) -> str:
+    def _prepare_uri(self, col: Column, ext: str | None = None) -> str:
         """
         Construct a new, unique URI for a persisted media file.
         """
-        assert col.tbl is not None, 'Column must be associated with a table'
-        return self._prepare_uri_raw(col.tbl.id, col.id, col.tbl.version, ext=ext)
+        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"""
@@ -142,7 +142,7 @@ class GCSStore(ObjectStoreBase):
             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: Optional[int] = None) -> Iterator:
+    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:
@@ -168,7 +168,7 @@ class GCSStore(ObjectStoreBase):
 
         return blob_iterator
 
-    def count(self, tbl_id: uuid.UUID, tbl_version: Optional[int] = None) -> int:
+    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.
 
@@ -193,7 +193,7 @@ class GCSStore(ObjectStoreBase):
             self.handle_gcs_error(e, self.bucket_name, f'setup iterator {self.prefix}')
             raise
 
-    def delete(self, tbl_id: uuid.UUID, tbl_version: Optional[int] = None) -> int:
+    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.