pixeltable 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

pixeltable/share/publish.py

@@ -1,36 +1,45 @@
+import os
 import sys
 import urllib.parse
 import urllib.request
 from pathlib import Path
+from typing import Literal, Optional
 
 import requests
+from requests.adapters import HTTPAdapter
 from tqdm import tqdm
+from urllib3.util.retry import Retry
 
 import pixeltable as pxt
 from pixeltable import exceptions as excs
 from pixeltable.env import Env
 from pixeltable.utils import sha256sum
+from pixeltable.utils.media_store import TempStore
 
 from .packager import TablePackager, TableRestorer
 
 # These URLs are abstracted out for now, but will be replaced with actual (hard-coded) URLs once the
 # pixeltable.com URLs are available.
 
-PIXELTABLE_API_URL = 'https://internal-api.pixeltable.com'
+PIXELTABLE_API_URL = os.environ.get('PIXELTABLE_API_URL', 'https://internal-api.pixeltable.com')
 
 
-def push_replica(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
-    if not src_tbl._tbl_version.get().is_snapshot:
+def push_replica(
+    dest_tbl_uri: str, src_tbl: pxt.Table, bucket: str | None = None, access: Literal['public', 'private'] = 'private'
+) -> str:
+    if not src_tbl._tbl_version_path.is_snapshot():
         raise excs.Error('Only snapshots may be published.')
 
-    packager = TablePackager(src_tbl, additional_md={'table_uri': dest_tbl_uri})
+    packager = TablePackager(
+        src_tbl, additional_md={'table_uri': dest_tbl_uri, 'bucket_name': bucket, 'is_public': access == 'public'}
+    )
     request_json = packager.md | {'operation_type': 'publish_snapshot'}
     headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
     response = requests.post(PIXELTABLE_API_URL, json=request_json, headers=headers_json)
     if response.status_code != 200:
         raise excs.Error(f'Error publishing snapshot: {response.text}')
     response_json = response.json()
-    if not isinstance(response_json, dict) or response_json.get('destination') != 's3':
+    if not isinstance(response_json, dict):
         raise excs.Error(f'Error publishing snapshot: unexpected response from server.\n{response_json}')
     upload_id = response_json['upload_id']
     destination_uri = response_json['destination_uri']
@@ -42,17 +51,23 @@ def push_replica(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
     parsed_location = urllib.parse.urlparse(destination_uri)
     if parsed_location.scheme == 's3':
         _upload_bundle_to_s3(bundle, parsed_location)
+    elif parsed_location.scheme == 'https':
+        _upload_to_presigned_url(file_path=bundle, url=parsed_location.geturl())
     else:
         raise excs.Error(f'Unsupported destination: {destination_uri}')
 
     Env.get().console_logger.info('Finalizing snapshot ...')
 
     finalize_request_json = {
+        'table_uri': dest_tbl_uri,
         'operation_type': 'finalize_snapshot',
         'upload_id': upload_id,
         'datafile': bundle.name,
         'size': bundle.stat().st_size,
         'sha256': sha256sum(bundle),  # Generate our own SHA for independent verification
+        'rows': packager.md['row_count'],  # TODO rename rows to row_count once cloud side changes are complete
+        'preview_header': packager.md['preview_header'],
+        'preview_data': packager.md['preview_data'],
     }
     # TODO: Use Pydantic for validation
     finalize_response = requests.post(PIXELTABLE_API_URL, json=finalize_request_json, headers=headers_json)
@@ -107,11 +122,14 @@ def pull_replica(dest_path: str, src_tbl_uri: str) -> pxt.Table:
         raise excs.Error(f'Error cloning shapshot: unexpected response from server.\n{response_json}')
 
     primary_tbl_additional_md = response_json['md']['tables'][0]['table_md']['additional_md']
-    bundle_uri = primary_tbl_additional_md['destination_uri']
+    bundle_uri = response_json['destination_uri']
     bundle_filename = primary_tbl_additional_md['datafile']
     parsed_location = urllib.parse.urlparse(bundle_uri)
     if parsed_location.scheme == 's3':
         bundle_path = _download_bundle_from_s3(parsed_location, bundle_filename)
+    elif parsed_location.scheme == 'https':
+        bundle_path = TempStore.create_path()
+        _download_from_presigned_url(url=parsed_location.geturl(), output_path=bundle_path)
     else:
         raise excs.Error(f'Unexpected response from server: unsupported bundle uri: {bundle_uri}')
 
@@ -136,7 +154,7 @@ def _download_bundle_from_s3(parsed_location: urllib.parse.ParseResult, bundle_f
     obj = s3_client.head_object(Bucket=bucket, Key=remote_path)  # Check if the object exists
     bundle_size = obj['ContentLength']
 
-    bundle_path = Path(Env.get().create_tmp_path())
+    bundle_path = TempStore.create_path()
     progress_bar = tqdm(
         desc='Downloading',
         total=bundle_size,
@@ -149,3 +167,112 @@ def _download_bundle_from_s3(parsed_location: urllib.parse.ParseResult, bundle_f
     )
     s3_client.download_file(Bucket=bucket, Key=remote_path, Filename=str(bundle_path), Callback=progress_bar.update)
     return bundle_path
+
+
+def _create_retry_session(
+    max_retries: int = 3, backoff_factor: float = 1.0, status_forcelist: Optional[list] = None
+) -> requests.Session:
+    """Create a requests session with retry configuration"""
+    if status_forcelist is None:
+        status_forcelist = [
+            408,  # Request Timeout
+            429,  # Too Many Requests (rate limiting)
+            500,  # Internal Server Error (server-side error)
+            502,  # Bad Gateway (proxy/gateway got invalid response)
+            503,  # Service Unavailable (server overloaded or down)
+            504,  # Gateway Timeout (proxy/gateway timeout)
+        ]
+    retry_strategy = Retry(
+        total=max_retries,
+        read=max_retries,
+        connect=max_retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+        allowed_methods=['GET', 'PUT', 'POST', 'DELETE'],
+    )
+
+    session = requests.Session()
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+    session.mount('https://', adapter)
+    return session
+
+
+def _upload_to_presigned_url(file_path: Path, url: str, max_retries: int = 3) -> requests.Response:
+    """Upload file with progress bar and retries"""
+    file_size = file_path.stat().st_size
+
+    headers = {'Content-Length': str(file_size), 'Content-Type': 'application/octet-stream'}
+
+    session = _create_retry_session(max_retries=max_retries)
+    try:
+        with (
+            open(file_path, 'rb') as f,
+            tqdm.wrapattr(
+                f,
+                method='read',
+                total=file_size,
+                desc='Uploading',
+                unit='B',
+                unit_scale=True,
+                unit_divisor=1024,
+                miniters=1,  # Update every iteration (should be fine for an upload)
+                ncols=100,
+                file=sys.stdout,
+            ) as file_with_progress,
+        ):
+            response = session.put(
+                url,
+                data=file_with_progress,
+                headers=headers,
+                timeout=(60, 1800),  # 60 seconds to connect and 300 seconds for server response
+            )
+            response.raise_for_status()
+            return response
+    finally:
+        session.close()
+
+
+def _download_from_presigned_url(
+    url: str, output_path: Path, headers: Optional[dict[str, str]] = None, max_retries: int = 3
+) -> None:
+    """Download file with progress bar and retries"""
+    session = _create_retry_session(max_retries=max_retries)
+
+    try:
+        # Stream download with progress
+        response = session.get(
+            url, headers=headers, stream=True, timeout=(60, 300)
+        )  # 60 seconds to connect and 300 seconds for server response
+        response.raise_for_status()
+
+        total_size = int(response.headers.get('content-length', 0))
+        progress_bar = tqdm(
+            desc='Downloading',
+            total=total_size,
+            unit='B',
+            unit_scale=True,
+            unit_divisor=1024,
+            miniters=1,
+            ncols=100,
+            file=sys.stdout,
+        )
+        with open(output_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+                    progress_bar.update(len(chunk))
+    finally:
+        session.close()
+
+
+# TODO: This will be replaced by drop_table with cloud table uri
+def delete_replica(dest_path: str) -> None:
+    """Delete cloud replica"""
+    headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
+    delete_request_json = {'operation_type': 'delete_snapshot', 'table_uri': dest_path}
+    response = requests.post(PIXELTABLE_API_URL, json=delete_request_json, headers=headers_json)
+    if response.status_code != 200:
+        raise excs.Error(f'Error deleting replica: {response.text}')
+    response_json = response.json()
+    if not isinstance(response_json, dict) or 'table_uri' not in response_json:
+        raise excs.Error(f'Error deleting replica: unexpected response from server.\n{response_json}')

pixeltable/utils/media_store.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import glob
+import logging
 import os
 import re
 import shutil
@@ -19,10 +20,12 @@ from pixeltable import env
 if TYPE_CHECKING:
     from pixeltable.catalog import Column
 
+_logger = logging.getLogger('pixeltable')
+
 
 class MediaStore:
     """
-    Utilities to manage media files stored in Env.media_dir
+    Utilities to manage media files stored in a local filesystem directory.
 
     Media file names are a composite of: table id, column id, tbl_version, new uuid:
     the table id/column id/tbl_version are redundant but useful for identifying all files for a table
@@ -30,38 +33,70 @@ class MediaStore:
     """
 
     pattern = re.compile(r'([0-9a-fA-F]+)_(\d+)_(\d+)_([0-9a-fA-F]+)')  # tbl_id, col_id, version, uuid
+    __base_dir: Path
+
+    def __init__(self, base_dir: Path):
+        """Initialize a MediaStore with a base directory."""
+        assert isinstance(base_dir, Path), 'Base directory must be a Path instance.'
+        self.__base_dir = base_dir
 
     @classmethod
-    def _media_dir(cls) -> Path:
-        """Returns the media directory path."""
-        return env.Env.get().media_dir
+    def get(cls, base_uri: Optional[Path] = None) -> MediaStore:
+        """Get a MediaStore instance for the given base URI, or the environment's media_dir if None."""
+        if base_uri is None:
+            return MediaStore(env.Env.get().media_dir)
+        raise NotImplementedError
 
     @classmethod
-    def _tmp_dir(cls) -> Path:
-        """Returns the temporary directory path."""
-        return env.Env.get().tmp_dir
+    def _save_binary_media_file(cls, file_data: bytes, dest_path: Path, format: Optional[str]) -> Path:
+        """Save a media binary data to a file in a MediaStore. format is ignored for binary data."""
+        assert isinstance(file_data, bytes)
+        with open(dest_path, 'wb') as f:
+            f.write(file_data)
+            f.flush()  # Ensures Python buffers are written to OS
+            os.fsync(f.fileno())  # Forces OS to write to physical storage
+        return dest_path
 
     @classmethod
-    def _prepare_media_path(cls, col: Column, ext: Optional[str] = None) -> Path:
+    def _save_pil_image_file(cls, image: PIL.Image.Image, dest_path: Path, format: Optional[str]) -> Path:
+        """Save a PIL Image to a file in a MediaStore with the specified format."""
+        if dest_path.suffix != f'.{format}':
+            dest_path = dest_path.with_name(f'{dest_path.name}.{format}')
+
+        with open(dest_path, 'wb') as f:
+            image.save(f, format=format)
+            f.flush()  # Ensures Python buffers are written to OS
+            os.fsync(f.fileno())  # Forces OS to write to physical storage
+        return dest_path
+
+    def _prepare_media_path_raw(self, tbl_id: UUID, col_id: int, tbl_version: int, ext: Optional[str] = None) -> Path:
         """
         Construct a new, unique Path name for a persisted media file, and create the parent directory
         for the new Path if it does not already exist. The Path will reside in
         the environment's media_dir.
         """
         id_hex = uuid.uuid4().hex
-        parent = cls._media_dir() / col.tbl.id.hex / id_hex[:2] / id_hex[:4]
+        parent = self.__base_dir / tbl_id.hex / id_hex[:2] / id_hex[:4]
         parent.mkdir(parents=True, exist_ok=True)
-        return parent / f'{col.tbl.id.hex}_{col.id}_{col.tbl.version}_{id_hex}{ext or ""}'
+        return parent / f'{tbl_id.hex}_{col_id}_{tbl_version}_{id_hex}{ext or ""}'
 
-    @classmethod
-    def resolve_tmp_url(cls, file_url: Optional[str]) -> Optional[Path]:
-        """Return path if the given url is a tmp file.
+    def _prepare_media_path(self, col: Column, ext: Optional[str] = None) -> Path:
+        """
+        Construct a new, unique Path name for a persisted media file, and create the parent directory
+        for the new Path if it does not already exist. The Path will reside in
+        the environment's media_dir.
+        """
+        assert col.tbl is not None, 'Column must be associated with a table'
+        return self._prepare_media_path_raw(col.tbl.id, col.id, col.tbl.version, ext)
+
+    def resolve_url(self, file_url: Optional[str]) -> Optional[Path]:
+        """Return path if the given url refers to a file managed by this MediaStore, else None.
 
         Args:
-            file_url: URL of the tmp media file to check
+            file_url: URL to check
 
         Returns:
-            If the file_url is a tmp file, return a Path() to the tmp file, None, otherwise
+            If the url is a managed file, return a Path() to the file, None, otherwise
         """
         if file_url is None:
             return None
@@ -74,93 +109,76 @@ class MediaStore:
             # remote url
             return None
         src_path = urllib.parse.unquote(urllib.request.url2pathname(parsed.path))
-        pxt_tmp_dir = str(cls._tmp_dir())
-        if not src_path.startswith(pxt_tmp_dir):
+        if not src_path.startswith(str(self.__base_dir)):
             # not a tmp file
             return None
         return Path(src_path)
 
-    @classmethod
-    def relocate_local_media_file(cls, src_path: Path, col: Column) -> str:
-        """Relocate a local file to the MediaStore, and return its new URL"""
-        dest_path = cls._prepare_media_path(col, ext=src_path.suffix)
+    def relocate_local_media_file(self, src_path: Path, col: Column) -> str:
+        """Relocate a local file to a MediaStore, and return its new URL"""
+        dest_path = self._prepare_media_path(col, ext=src_path.suffix)
         src_path.rename(dest_path)
-        return urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
-
-    @classmethod
-    def save_media_object(cls, data: bytes | PIL.Image.Image, col: Column, format: Optional[str]) -> tuple[Path, str]:
-        """Save a media data to a file in the MediaStore
+        new_file_url = urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
+        _logger.debug(f'Media Storage: moved {src_path} to {new_file_url}')
+        return new_file_url
+
+    def copy_local_media_file(self, src_path: Path, col: Column) -> str:
+        """Copy a local file to a MediaStore, and return its new URL"""
+        dest_path = self._prepare_media_path(col, ext=src_path.suffix)
+        shutil.copy2(src_path, dest_path)
+        new_file_url = urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
+        _logger.debug(f'Media Storage: copied {src_path} to {new_file_url}')
+        return new_file_url
+
+    def save_media_object(self, data: bytes | PIL.Image.Image, col: Column, format: Optional[str]) -> tuple[Path, str]:
+        """Save a media data object to a file in a MediaStore
         Returns:
             dest_path: Path to the saved media file
             url: URL of the saved media file
         """
         assert col.col_type.is_media_type(), f'MediaStore: request to store non media_type Column {col.name}'
-        dest_path = cls._prepare_media_path(col)
+        dest_path = self._prepare_media_path(col)
         if isinstance(data, bytes):
-            dest_path = cls._save_binary_media_file(data, dest_path, format)
+            dest_path = self._save_binary_media_file(data, dest_path, format)
         elif isinstance(data, PIL.Image.Image):
-            dest_path = cls._save_pil_image_file(data, dest_path, format)
+            dest_path = self._save_pil_image_file(data, dest_path, format)
         else:
             raise ValueError(f'Unsupported media object type: {type(data)}')
-        url = urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
-        return dest_path, url
+        new_file_url = urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
+        return dest_path, new_file_url
 
-    @classmethod
-    def _save_binary_media_file(cls, file_data: bytes, dest_path: Path, format: Optional[str]) -> Path:
-        """Save a media binary data to a file in the MediaStore. format is ignored for binary data."""
-        assert isinstance(file_data, bytes)
-        with open(dest_path, 'wb') as f:
-            f.write(file_data)
-            f.flush()  # Ensures Python buffers are written to OS
-            os.fsync(f.fileno())  # Forces OS to write to physical storage
-        return dest_path
-
-    @classmethod
-    def _save_pil_image_file(cls, image: PIL.Image.Image, dest_path: Path, format: Optional[str]) -> Path:
-        """Save a PIL Image to a file in the MediaStore with the specified format."""
-        if dest_path.suffix != f'.{format}':
-            dest_path = dest_path.with_name(f'{dest_path.name}.{format}')
-
-        with open(dest_path, 'wb') as f:
-            image.save(f, format=format)
-            f.flush()  # Ensures Python buffers are written to OS
-            os.fsync(f.fileno())  # Forces OS to write to physical storage
-        return dest_path
-
-    @classmethod
-    def delete(cls, tbl_id: UUID, tbl_version: Optional[int] = None) -> None:
+    def delete(self, tbl_id: UUID, tbl_version: Optional[int] = None) -> None:
         """Delete all files belonging to tbl_id. If tbl_version is not None, delete
         only those files belonging to the specified tbl_version."""
         assert tbl_id is not None
         if tbl_version is None:
             # Remove the entire folder for this table id.
-            path = cls._media_dir() / tbl_id.hex
+            path = self.__base_dir / tbl_id.hex
             if path.exists():
                 shutil.rmtree(path)
         else:
             # Remove only the elements for the specified tbl_version.
-            paths = glob.glob(
-                str(cls._media_dir() / tbl_id.hex) + f'/**/{tbl_id.hex}_*_{tbl_version}_*', recursive=True
-            )
+            paths = glob.glob(str(self.__base_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*_{tbl_version}_*', recursive=True)
             for p in paths:
                 os.remove(p)
 
-    @classmethod
-    def count(cls, tbl_id: UUID) -> int:
+    def count(self, tbl_id: Optional[UUID]) -> int:
         """
         Return number of files for given tbl_id.
         """
-        paths = glob.glob(str(cls._media_dir() / tbl_id.hex) + f'/**/{tbl_id.hex}_*', recursive=True)
+        if tbl_id is None:
+            paths = glob.glob(str(self.__base_dir / '*'), recursive=True)
+        else:
+            paths = glob.glob(str(self.__base_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*', recursive=True)
         return len(paths)
 
-    @classmethod
-    def stats(cls) -> list[tuple[UUID, int, int, int]]:
-        paths = glob.glob(str(cls._media_dir()) + '/**', recursive=True)
+    def stats(self) -> list[tuple[UUID, int, int, int]]:
+        paths = glob.glob(str(self.__base_dir) + '/**', recursive=True)
         # key: (tbl_id, col_id), value: (num_files, size)
         d: dict[tuple[UUID, int], list[int]] = defaultdict(lambda: [0, 0])
         for p in paths:
             if not os.path.isdir(p):
-                matched = re.match(cls.pattern, Path(p).name)
+                matched = re.match(self.pattern, Path(p).name)
                 assert matched is not None
                 tbl_id, col_id = UUID(hex=matched[1]), int(matched[2])
                 file_info = os.stat(p)
@@ -170,3 +188,50 @@ class MediaStore:
         result = [(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()]
         result.sort(key=lambda e: e[3], reverse=True)
         return result
+
+
+class TempStore:
+    """
+    A temporary store for files of data that are not yet persisted to their destination(s).
+    A destination is typically either a MediaStore (local persisted files) or a cloud object store.
+
+    The TempStore class has no internal state. It provides functionality to manage temporary files
+    in the env.Env.get().tmp_dir directory.
+    It reuses some of the MediaStore functionality to create unique file names and save objects.
+    """
+
+    @classmethod
+    def _tmp_dir(cls) -> Path:
+        """Returns the path to the temporary directory where files are stored."""
+        from pixeltable import env
+
+        return env.Env.get().tmp_dir
+
+    @classmethod
+    def count(cls, tbl_id: Optional[UUID] = None) -> int:
+        return MediaStore(cls._tmp_dir()).count(tbl_id)
+
+    @classmethod
+    def resolve_url(cls, file_url: Optional[str]) -> Optional[Path]:
+        return MediaStore(cls._tmp_dir()).resolve_url(file_url)
+
+    @classmethod
+    def save_media_object(cls, data: bytes | PIL.Image.Image, col: Column, format: Optional[str]) -> tuple[Path, str]:
+        return MediaStore(cls._tmp_dir()).save_media_object(data, col, format)
+
+    @classmethod
+    def delete_media_file(cls, obj_path: Path) -> None:
+        """Delete a media object from the temporary store."""
+        assert obj_path is not None, 'Object path must be provided'
+        assert obj_path.exists(), f'Object path does not exist: {obj_path}'
+        assert cls.resolve_url(str(obj_path)) is not None, f'Object path is not a valid media store path: {obj_path}'
+        obj_path.unlink()
+
+    @classmethod
+    def create_path(cls, tbl_id: Optional[UUID] = None, extension: str = '') -> Path:
+        """Return a new, unique Path located in the temporary store.
+        If tbl_id is provided, the path name will be similar to a MediaStore path based on the tbl_id.
+        If tbl_id is None, a random UUID will be used to create the path."""
+        if tbl_id is not None:
+            return MediaStore(cls._tmp_dir())._prepare_media_path_raw(tbl_id, 0, 0, extension)
+        return cls._tmp_dir() / f'{uuid.uuid4()}{extension}'