pixeltable 0.4.6__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/store.py

@@ -4,7 +4,7 @@ import abc
 import logging
 import sys
 import warnings
-from typing import Any, Iterable, Iterator, Optional, Union
+from typing import Any, Iterable, Iterator, Optional
 
 import more_itertools
 import psycopg
@@ -17,7 +17,6 @@ from pixeltable.env import Env
 from pixeltable.exec import ExecNode
 from pixeltable.metadata import schema
 from pixeltable.utils.exception_handler import run_cleanup
-from pixeltable.utils.media_store import MediaStore
 from pixeltable.utils.sql import log_explain, log_stmt
 
 _logger = logging.getLogger('pixeltable')
@@ -123,21 +122,6 @@ class StoreBase:
     def _storage_name(self) -> str:
         """Return the name of the data store table"""
 
-    def _move_tmp_media_file(self, file_url: Optional[str], col: catalog.Column) -> str:
-        src_path = MediaStore.resolve_tmp_url(file_url)
-        if src_path is None:
-            return file_url
-        assert col.tbl.id == self.tbl_version.id  # Ensure the column belongs to the same table as this store
-        new_file_url = MediaStore.relocate_local_media_file(src_path, col)
-        return new_file_url
-
-    def _move_tmp_media_files(
-        self, table_row: list[Any], media_cols_by_sql_idx: dict[int, catalog.Column], v_min: int
-    ) -> None:
-        """Move tmp media files that we generated to a permanent location"""
-        for n, col in media_cols_by_sql_idx.items():
-            table_row[n] = self._move_tmp_media_file(table_row[n], col)
-
     def count(self) -> int:
         """Return the number of rows visible in self.tbl_version"""
         stmt = (
@@ -235,7 +219,6 @@ class StoreBase:
         # create temp table to store output of exec_plan, with the same primary key as the store table
         tmp_name = f'temp_{self._storage_name()}'
         tmp_pk_cols = tuple(sql.Column(col.name, col.type, primary_key=True) for col in self.pk_columns())
-        tmp_val_col_sql_idx = len(tmp_pk_cols)
         tmp_val_col = sql.Column(col.sa_col.name, col.sa_col.type)
         tmp_cols = [*tmp_pk_cols, tmp_val_col]
         # add error columns if the store column records errors
@@ -262,9 +245,7 @@ class StoreBase:
                     if abort_on_exc and row.has_exc():
                         exc = row.get_first_exc()
                         raise excs.Error(f'Error while evaluating computed column {col.name!r}:\n{exc}') from exc
-                    table_row, num_row_exc = row_builder.create_table_row(row, None, row.pk)
-                    if col.col_type.is_media_type():
-                        table_row[tmp_val_col_sql_idx] = self._move_tmp_media_file(table_row[tmp_val_col_sql_idx], col)
+                    table_row, num_row_exc = row_builder.create_store_table_row(row, None, row.pk)
                     num_excs += num_row_exc
                     batch_table_rows.append(tuple(table_row))
 
@@ -317,7 +298,7 @@ class StoreBase:
         progress_bar: Optional[tqdm] = None  # create this only after we started executing
         row_builder = exec_plan.row_builder
 
-        store_col_names, media_cols_by_idx = row_builder.store_column_names()
+        store_col_names = row_builder.store_column_names()
 
         try:
             table_rows: list[tuple[Any]] = []
@@ -337,7 +318,7 @@ class StoreBase:
                     rowid = (next(rowids),) if rowids is not None else row.pk[:-1]
                     pk = (*rowid, v_min)
                     assert len(pk) == len(self._pk_cols)
-                    table_row, num_row_exc = row_builder.create_table_row(row, cols_with_excs, pk)
+                    table_row, num_row_exc = row_builder.create_store_table_row(row, cols_with_excs, pk)
                     num_excs += num_row_exc
 
                     if show_progress:
@@ -351,7 +332,6 @@ class StoreBase:
                             )
                         progress_bar.update(1)
 
-                    self._move_tmp_media_files(table_row, media_cols_by_idx, v_min)
                     batch_table_rows.append(tuple(table_row))
 
                 table_rows.extend(batch_table_rows)
@@ -427,7 +407,7 @@ class StoreBase:
         base_versions_clause = (
             sql.true() if len(base_versions) == 0 else self.base._versions_clause(base_versions, match_on_vmin)
         )
-        set_clause: dict[sql.Column, Union[int, sql.Column]] = {self.v_max_col: current_version}
+        set_clause: dict[sql.Column, int | sql.Column] = {self.v_max_col: current_version}
         for index_info in self.tbl_version.get().idxs_by_name.values():
             # copy value column to undo column
             set_clause[index_info.undo_col.sa_col] = index_info.val_col.sa_col

pixeltable/type_system.py

@@ -292,7 +292,7 @@ class ColumnType:
 
     @classmethod
     def from_python_type(
-        cls, t: Union[type, _GenericAlias], nullable_default: bool = False, allow_builtin_types: bool = True
+        cls, t: type | _GenericAlias, nullable_default: bool = False, allow_builtin_types: bool = True
     ) -> Optional[ColumnType]:
         """
         Convert a Python type into a Pixeltable `ColumnType` instance.
@@ -311,7 +311,7 @@ class ColumnType:
         if origin in (typing.Union, types.UnionType):
             # Check if `t` has the form Optional[T].
             if len(type_args) == 2 and type(None) in type_args:
-                # `t` is a type of the form Optional[T] (equivalently, Union[T, None] or Union[None, T]).
+                # `t` is a type of the form Optional[T] (equivalently, T | None or None | T).
                 # We treat it as the underlying type but with nullable=True.
                 underlying_py_type = type_args[0] if type_args[1] is type(None) else type_args[1]
                 underlying = cls.from_python_type(underlying_py_type, allow_builtin_types=allow_builtin_types)
@@ -361,10 +361,7 @@ class ColumnType:
 
     @classmethod
     def normalize_type(
-        cls,
-        t: Union[ColumnType, type, _AnnotatedAlias],
-        nullable_default: bool = False,
-        allow_builtin_types: bool = True,
+        cls, t: ColumnType | type | _AnnotatedAlias, nullable_default: bool = False, allow_builtin_types: bool = True
     ) -> ColumnType:
         """
         Convert any type recognizable by Pixeltable to its corresponding ColumnType.
@@ -389,7 +386,7 @@ class ColumnType:
     ]
 
     @classmethod
-    def __raise_exc_for_invalid_type(cls, t: Union[type, _AnnotatedAlias]) -> None:
+    def __raise_exc_for_invalid_type(cls, t: type | _AnnotatedAlias) -> None:
         for builtin_type, suggestion in cls.__TYPE_SUGGESTIONS:
             if t is builtin_type or (isinstance(t, type) and issubclass(t, builtin_type)):
                 name = t.__name__ if t.__module__ == 'builtins' else f'{t.__module__}.{t.__name__}'
@@ -405,7 +402,7 @@ class ColumnType:
         return cls.from_python_type(py_type) if py_type is not None else None
 
     @classmethod
-    def __json_schema_to_py_type(cls, schema: dict[str, Any]) -> Union[type, _GenericAlias, None]:
+    def __json_schema_to_py_type(cls, schema: dict[str, Any]) -> type | _GenericAlias | None:
         if 'type' in schema:
             if schema['type'] == 'null':
                 return type(None)

pixeltable/utils/__init__.py

@@ -2,7 +2,7 @@ import hashlib
 import urllib.parse
 import urllib.request
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 
 def print_perf_counter_delta(delta: float) -> str:
@@ -24,7 +24,7 @@ def print_perf_counter_delta(delta: float) -> str:
         return f'{delta:.2f} s'
 
 
-def sha256sum(path: Union[Path, str]) -> str:
+def sha256sum(path: Path | str) -> str:
     """
     Compute the SHA256 hash of a file.
     """

pixeltable/utils/arrow.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Any, Iterator, Optional, Union
+from typing import Any, Iterator, Optional
 
 import numpy as np
 import pyarrow as pa
@@ -88,11 +88,11 @@ def to_arrow_schema(pixeltable_schema: dict[str, Any]) -> pa.Schema:
     return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())  # type: ignore[misc]
 
 
-def to_pydict(batch: Union[pa.Table, pa.RecordBatch]) -> dict[str, Union[list, np.ndarray]]:
+def to_pydict(batch: pa.Table | pa.RecordBatch) -> dict[str, list | np.ndarray]:
     """Convert a RecordBatch to a dictionary of lists, unlike pa.lib.RecordBatch.to_pydict,
     this function will not convert numpy arrays to lists, and will preserve the original numpy dtype.
     """
-    out: dict[str, Union[list, np.ndarray]] = {}
+    out: dict[str, list | np.ndarray] = {}
     for k, name in enumerate(batch.schema.names):
         col = batch.column(k)
         if isinstance(col.type, pa.FixedShapeTensorType):
@@ -105,7 +105,7 @@ def to_pydict(batch: Union[pa.Table, pa.RecordBatch]) -> dict[str, Union[list, n
     return out
 
 
-def iter_tuples(batch: Union[pa.Table, pa.RecordBatch]) -> Iterator[dict[str, Any]]:
+def iter_tuples(batch: pa.Table | pa.RecordBatch) -> Iterator[dict[str, Any]]:
     """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
     pydict = to_pydict(batch)
     assert len(pydict) > 0, 'empty record batch'
@@ -145,7 +145,7 @@ def _ar_val_to_pxt_val(val: Any, pxt_type: ts.ColumnType) -> Any:
 
 
 def iter_tuples2(
-    batch: Union[pa.Table, pa.RecordBatch], col_mapping: Optional[dict[str, str]], schema: dict[str, ts.ColumnType]
+    batch: pa.Table | pa.RecordBatch, col_mapping: Optional[dict[str, str]], schema: dict[str, ts.ColumnType]
 ) -> Iterator[dict[str, Any]]:
     """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
     pydict = to_pydict(batch)

pixeltable/utils/description_helper.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import Optional, Union
+from typing import Optional
 
 import pandas as pd
 from pandas.io.formats.style import Styler
@@ -7,7 +7,7 @@ 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
@@ -33,7 +33,7 @@ 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,

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.
     """