pixeltable 0.3.10__py3-none-any.whl → 0.3.12__py3-none-any.whl

pixeltable/io/label_studio.py

@@ -5,16 +5,14 @@ import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Iterator, Literal, Optional, cast
-from xml.etree import ElementTree
+from xml.etree import ElementTree as ET
 
 import label_studio_sdk  # type: ignore[import-untyped]
 import PIL.Image
 from requests.exceptions import HTTPError
 
 import pixeltable as pxt
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-from pixeltable import Column, Table
+from pixeltable import Column, Table, env, exceptions as excs
 from pixeltable.config import Config
 from pixeltable.exprs import ColumnRef, DataRow, Expr
 from pixeltable.io.external_store import Project, SyncStatus
@@ -140,7 +138,8 @@ class LabelStudioProject(Project):
             page += 1
         if unknown_task_count > 0:
             _logger.warning(
-                f'Skipped {unknown_task_count} unrecognized task(s) when syncing Label Studio project "{self.project_title}".'
+                f'Skipped {unknown_task_count} unrecognized task(s) when syncing '
+                f'Label Studio project {self.project_title!r}.'
             )
 
     def __update_tasks(self, t: Table, existing_tasks: dict[tuple, dict]) -> SyncStatus:
@@ -174,11 +173,11 @@ class LabelStudioProject(Project):
             # Send media to Label Studio by HTTP post.
             assert len(t_data_cols) == 1  # This was verified when the project was set up
             return self.__update_tasks_by_post(t, existing_tasks, t_data_cols[0], t_rl_cols, rl_info)
-        elif self.media_import_method == 'file' or self.media_import_method == 'url':
+        elif self.media_import_method in ('file', 'url'):
             # Send media to Label Studio by file reference (local file or URL).
             return self.__update_tasks_by_files(t, existing_tasks, t_data_cols, t_rl_cols, rl_info)
         else:
-            assert False
+            raise AssertionError()
 
     def __update_tasks_by_post(
         self,
@@ -227,7 +226,7 @@ class LabelStudioProject(Project):
                     )
                     for i in range(len(coco_annotations))
                 ]
-                _logger.debug(f'`predictions`: %s', predictions)
+                _logger.debug('`predictions`: {%s}', predictions)
                 self.project.create_predictions(predictions)
                 tasks_created += 1
 
@@ -358,7 +357,7 @@ class LabelStudioProject(Project):
     def __localpath_to_lspath(cls, localpath: str) -> str:
         # Transform the local path into Label Studio's bespoke path format.
         relpath = Path(localpath).relative_to(Config.get().home)
-        return f'/data/local-files/?d={str(relpath)}'
+        return f'/data/local-files/?d={relpath}'
 
     def __delete_stale_tasks(
         self, existing_tasks: dict[tuple, dict], row_ids_in_pxt: set[tuple], tasks_created: int
@@ -405,7 +404,8 @@ class LabelStudioProject(Project):
         updates = [{'_rowid': rowid, local_annotations_col.name: ann} for rowid, ann in annotations.items()]
         if len(updates) > 0:
             _logger.info(
-                f'Updating table `{t._name}`, column `{local_annotations_col.name}` with {len(updates)} total annotations.'
+                f'Updating table {t._name!r}, column {local_annotations_col.name!r} '
+                f'with {len(updates)} total annotations.'
             )
             # batch_update currently doesn't propagate from views to base tables. As a workaround, we call
             # batch_update on the actual ancestor table that holds the annotations column.
@@ -451,7 +451,7 @@ class LabelStudioProject(Project):
         Parses a Label Studio XML config, extracting the names and Pixeltable types of
         all input variables.
         """
-        root: ElementTree.Element = ElementTree.fromstring(xml_config)
+        root: ET.Element = ET.fromstring(xml_config)
         if root.tag.lower() != 'view':
             raise excs.Error('Root of Label Studio config must be a `View`')
         config = _LabelStudioConfig(
@@ -461,7 +461,7 @@ class LabelStudioProject(Project):
         return config
 
     @classmethod
-    def __parse_data_keys_config(cls, root: ElementTree.Element) -> dict[str, '_DataKey']:
+    def __parse_data_keys_config(cls, root: ET.Element) -> dict[str, '_DataKey']:
         """Parses the data keys from a Label Studio XML config."""
         config: dict[str, '_DataKey'] = {}
         for element in root:
@@ -477,7 +477,7 @@ class LabelStudioProject(Project):
         return config
 
     @classmethod
-    def __parse_rectangle_labels_config(cls, root: ElementTree.Element) -> dict[str, '_RectangleLabel']:
+    def __parse_rectangle_labels_config(cls, root: ET.Element) -> dict[str, '_RectangleLabel']:
         """Parses the RectangleLabels from a Label Studio XML config."""
         config: dict[str, '_RectangleLabel'] = {}
         for element in root:
@@ -534,7 +534,7 @@ class LabelStudioProject(Project):
         _label_studio_client().delete_project(self.project_id)
         env.Env.get().console_logger.info(f'Deleted Label Studio project: {title}')
 
-    def __eq__(self, other) -> bool:
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, LabelStudioProject) and self.project_id == other.project_id
 
     def __hash__(self) -> int:
@@ -576,7 +576,7 @@ class LabelStudioProject(Project):
                 local_annotations_column = ANNOTATIONS_COLUMN
             else:
                 local_annotations_column = next(k for k, v in col_mapping.items() if v == ANNOTATIONS_COLUMN)
-            if local_annotations_column not in t._schema.keys():
+            if local_annotations_column not in t._schema:
                 t.add_columns({local_annotations_column: pxt.JsonType(nullable=True)})
 
         resolved_col_mapping = cls.validate_columns(
@@ -591,9 +591,9 @@ class LabelStudioProject(Project):
             if media_import_method != 'url':
                 raise excs.Error("`s3_configuration` is only valid when `media_import_method == 'url'`")
             s3_configuration = copy.copy(s3_configuration)
-            if not 'bucket' in s3_configuration:
+            if 'bucket' not in s3_configuration:
                 raise excs.Error('`s3_configuration` must contain a `bucket` field')
-            if not 'title' in s3_configuration:
+            if 'title' not in s3_configuration:
                 s3_configuration['title'] = 'Pixeltable-S3-Import-Storage'
             if (
                 'aws_access_key_id' not in s3_configuration
@@ -633,7 +633,8 @@ class LabelStudioProject(Project):
                         raise excs.Error(
                             '`media_import_method` is set to `file`, but your Label Studio server is not configured '
                             'for local file storage.\nPlease set the `LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED` '
-                            'environment variable to `true` in the environment where your Label Studio server is running.'
+                            'environment variable to `true` in the environment where your Label Studio server '
+                            'is running.'
                         ) from exc
                 raise  # Handle any other exception type normally
 
@@ -663,7 +664,7 @@ class _LabelStudioConfig:
     rectangle_labels: dict[str, _RectangleLabel]
 
     def validate(self) -> None:
-        data_key_names = set(key.name for key in self.data_keys.values() if key.name is not None)
+        data_key_names = {key.name for key in self.data_keys.values() if key.name is not None}
         for name, rl in self.rectangle_labels.items():
             if rl.to_name not in data_key_names:
                 raise excs.Error(
@@ -674,7 +675,7 @@ class _LabelStudioConfig:
     @property
     def export_columns(self) -> dict[str, pxt.ColumnType]:
         data_key_cols = {key_id: key_info.column_type for key_id, key_info in self.data_keys.items()}
-        rl_cols = {name: pxt.JsonType() for name in self.rectangle_labels.keys()}
+        rl_cols = {name: pxt.JsonType() for name in self.rectangle_labels}
         return {**data_key_cols, **rl_cols}
 
 

pixeltable/io/pandas.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Optional, Union
 
 import numpy as np
@@ -52,12 +53,12 @@ def import_pandas(
 
 def import_csv(
     tbl_name: str,
-    filepath_or_buffer,
+    filepath_or_buffer: Union[str, os.PathLike],
     schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs,
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from a csv file. This is a convenience method and is equivalent
@@ -81,13 +82,13 @@ def import_csv(
 
 def import_excel(
     tbl_name: str,
-    io,
-    *args,
+    io: Union[str, os.PathLike],
+    *,
     schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs,
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from an Excel (.xlsx) file. This is a convenience method and is

pixeltable/iterators/__init__.py

@@ -12,5 +12,5 @@ __removed_symbols = {'base', 'document', 'video'}
 __all__ = sorted(__default_dir - __removed_symbols)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/metadata/__init__.py

@@ -16,7 +16,7 @@ _console_logger = ConsoleLogger(logging.getLogger('pixeltable'))
 
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 31
+VERSION = 34
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:
@@ -24,9 +24,11 @@ def create_system_info(engine: sql.engine.Engine) -> None:
     system_md = SystemInfoMd(schema_version=VERSION)
     record = SystemInfo(md=dataclasses.asdict(system_md))
     with orm.Session(engine, future=True) as session:
-        session.add(record)
-        session.flush()
-        session.commit()
+        # Write system metadata only once for idempotency
+        if session.query(SystemInfo).count() == 0:
+            session.add(record)
+            session.flush()
+            session.commit()
 
 
 # conversion functions for upgrading the metadata schema from one version to the following

pixeltable/metadata/converters/convert_24.py

@@ -19,11 +19,11 @@ def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], A
         isinstance(v, dict)
         and '_classpath' in v
         and v['_classpath']
-        in {
+        in (
             'pixeltable.func.callable_function.CallableFunction',
             'pixeltable.func.aggregate_function.AggregateFunction',
             'pixeltable.func.expr_template_function.ExprTemplateFunction',
-        }
+        )
     ):
         if 'path' in v:
             assert 'signature' not in v
@@ -50,6 +50,6 @@ def __substitute_path(path: str) -> str:
     # versions, it's necessary to resolve the function symbol to get the signature. The following
     # adjustment is necessary for function names that are stored in db artifacts of version < 25, but
     # have changed in some version > 25.
-    if path in {'pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'}:
+    if path in ('pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'):
         return 'pixeltable.functions.huggingface.clip'
     return path

pixeltable/metadata/converters/convert_25.py

@@ -13,7 +13,7 @@ def _(engine: sql.engine.Engine) -> None:
 
 def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
     if k == 'path' and (
-        v in {'pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image'}
+        v in ('pixeltable.functions.huggingface.clip_text', 'pixeltable.functions.huggingface.clip_image')
     ):
         return 'path', 'pixeltable.functions.huggingface.clip'
     return None

pixeltable/metadata/converters/convert_29.py

@@ -80,7 +80,7 @@ def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], A
                         rolled_kwargs = kwargs.pop(param['name'])
 
         if rolled_args is not None:
-            assert rolled_args['_classname'] in {'InlineArray', 'InlineList'}
+            assert rolled_args['_classname'] in ('InlineArray', 'InlineList')
             new_args.extend(rolled_args['components'])
         if rolled_kwargs is not None:
             assert rolled_kwargs['_classname'] == 'InlineDict'

pixeltable/metadata/converters/convert_31.py

@@ -0,0 +1,11 @@
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+
+
+@register_converter(version=31)
+def _(engine: sql.engine.Engine) -> None:
+    # Add a column "lock_dummy: int8" to the dirs table in the store
+    # This column is the target of an UPDATE operation to synchronize directory operations
+    with engine.begin() as conn:
+        conn.execute(sql.text('ALTER TABLE dirs ADD COLUMN lock_dummy int8'))

pixeltable/metadata/converters/convert_32.py

@@ -0,0 +1,15 @@
+from uuid import UUID
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=32)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, table_md_updater=__update_table_md)
+
+
+def __update_table_md(table_md: dict, table_id: UUID) -> None:
+    table_md['is_replica'] = False

pixeltable/metadata/converters/convert_33.py

@@ -0,0 +1,17 @@
+from uuid import UUID
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=33)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, table_md_updater=__update_table_md)
+
+
+def __update_table_md(table_md: dict, table_id: UUID) -> None:
+    """Set default value of 'is_pk' field in column metadata to False"""
+    for col_md in table_md['column_md'].values():
+        col_md['is_pk'] = False if col_md['is_pk'] is None else col_md['is_pk']

pixeltable/metadata/notes.py

@@ -2,6 +2,9 @@
 # rather than as a comment, so that the existence of a description can be enforced by
 # the unit tests when new versions are added.
 VERSION_NOTES = {
+    34: 'Set default value for is_pk field in column metadata to False',
+    33: 'Add is_replica field to table metadata',
+    32: 'Add the lock_dummy BIGINT column to the dirs table',
     31: 'Add table ids to metadata structs',
     30: 'Store default values and constant arguments as literals',
     29: 'Add user and additional_md fields to metadata structs',

pixeltable/metadata/schema.py

@@ -1,7 +1,7 @@
 import dataclasses
 import typing
 import uuid
-from typing import Any, Optional, TypeVar, Union, get_type_hints
+from typing import Any, NamedTuple, Optional, TypeVar, Union, get_type_hints
 
 import sqlalchemy as sql
 from sqlalchemy import BigInteger, ForeignKey, Integer, LargeBinary, orm
@@ -84,6 +84,8 @@ class Dir(Base):
     )
     parent_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
     md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # DirMd
+    # This field is updated to synchronize database operations across multiple sessions
+    lock_dummy: orm.Mapped[int] = orm.mapped_column(BigInteger, nullable=True)
 
 
 @dataclasses.dataclass
@@ -155,6 +157,7 @@ class ViewMd:
 class TableMd:
     tbl_id: str  # uuid.UUID
     name: str
+    is_replica: bool
 
     user: Optional[str]
 
@@ -286,3 +289,25 @@ class Function(Base):
     dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
     md: orm.Mapped[dict[str, Any]] = orm.mapped_column(JSONB, nullable=False)  # FunctionMd
     binary_obj: orm.Mapped[Optional[bytes]] = orm.mapped_column(LargeBinary, nullable=True)
+
+
+class FullTableMd(NamedTuple):
+    tbl_md: TableMd
+    version_md: TableVersionMd
+    schema_version_md: TableSchemaVersionMd
+
+    def as_dict(self) -> dict[str, Any]:
+        return {
+            'table_id': self.tbl_md.tbl_id,
+            'table_md': dataclasses.asdict(self.tbl_md),
+            'table_version_md': dataclasses.asdict(self.version_md),
+            'table_schema_version_md': dataclasses.asdict(self.schema_version_md),
+        }
+
+    @classmethod
+    def from_dict(cls, data_dict: dict[str, Any]) -> 'FullTableMd':
+        return FullTableMd(
+            tbl_md=md_from_dict(TableMd, data_dict['table_md']),
+            version_md=md_from_dict(TableVersionMd, data_dict['table_version_md']),
+            schema_version_md=md_from_dict(TableSchemaVersionMd, data_dict['table_schema_version_md']),
+        )

pixeltable/plan.py

@@ -768,8 +768,7 @@ class Planner:
         # - select list subexprs that aren't aggregates
         # - join clause subexprs
         # - subexprs of Where clause conjuncts that can't be run in SQL
-        # - all grouping exprs, if any aggregate function call can't be run in SQL (in that case, they all have to be
-        #   run in Python)
+        # - all grouping exprs
         candidates = list(
             exprs.Expr.list_subexprs(
                 analyzer.select_list,
@@ -784,7 +783,7 @@ class Planner:
             candidates.extend(
                 exprs.Expr.subexprs(analyzer.filter, filter=sql_elements.contains, traverse_matches=False)
             )
-        if is_python_agg and analyzer.group_by_clause is not None:
+        if analyzer.group_by_clause is not None:
             candidates.extend(
                 exprs.Expr.list_subexprs(analyzer.group_by_clause, filter=sql_elements.contains, traverse_matches=False)
             )

pixeltable/share/packager.py

@@ -1,4 +1,3 @@
-import dataclasses
 import io
 import json
 import logging
@@ -6,7 +5,6 @@ import tarfile
 import urllib.parse
 import urllib.request
 import uuid
-from datetime import datetime
 from pathlib import Path
 from typing import Any, Iterator, Optional
 
@@ -58,28 +56,14 @@ class TablePackager:
         self.tmp_dir = Path(Env.get().create_tmp_path())
         self.media_files = {}
 
-        # Generate metadata
-        self.md = {
-            'pxt_version': pxt.__version__,
-            'pxt_md_version': metadata.VERSION,
-            'md': {
-                'tables': [
-                    {
-                        'table_id': str(t._tbl_version.id),
-                        # These are temporary; will replace with a better solution once the concurrency
-                        # changes to catalog have been merged
-                        'table_md': dataclasses.asdict(t._tbl_version.get()._create_tbl_md()),
-                        'table_version_md': dataclasses.asdict(
-                            t._tbl_version.get()._create_version_md(datetime.now().timestamp())
-                        ),
-                        'table_schema_version_md': dataclasses.asdict(
-                            t._tbl_version.get()._create_schema_version_md(0)
-                        ),
-                    }
-                    for t in (table, *table._bases)
-                ]
-            },
-        }
+        # Load metadata
+        with Env.get().begin_xact():
+            tbl_md = catalog.Catalog.get().load_replica_md(table)
+            self.md = {
+                'pxt_version': pxt.__version__,
+                'pxt_md_version': metadata.VERSION,
+                'md': {'tables': [md.as_dict() for md in tbl_md]},
+            }
         if additional_md is not None:
             self.md.update(additional_md)
 

pixeltable/share/publish.py

@@ -1,4 +1,3 @@
-import os
 import sys
 import urllib.parse
 import urllib.request
@@ -10,22 +9,22 @@ from tqdm import tqdm
 import pixeltable as pxt
 from pixeltable import exceptions as excs
 from pixeltable.env import Env
+from pixeltable.metadata.schema import FullTableMd
 from pixeltable.utils import sha256sum
 
 from .packager import TablePackager
 
 # These URLs are abstracted out for now, but will be replaced with actual (hard-coded) URLs once the
 # pixeltable.com URLs are available.
-_PUBLISH_URL = os.environ.get('PIXELTABLE_PUBLISH_URL')
-_FINALIZE_URL = os.environ.get('PIXELTABLE_FINALIZE_URL')
+
+PIXELTABLE_API_URL = 'https://internal-api.pixeltable.com'
 
 
 def publish_snapshot(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
     packager = TablePackager(src_tbl, additional_md={'table_uri': dest_tbl_uri})
-    request_json = packager.md
-    headers_json = {'X-api-key': Env.get().pxt_api_key}
-
-    response = requests.post(_PUBLISH_URL, json=request_json, headers=headers_json)
+    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()
@@ -47,14 +46,14 @@ def publish_snapshot(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
     Env.get().console_logger.info('Finalizing snapshot ...')
 
     finalize_request_json = {
+        '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
     }
-
     # TODO: Use Pydantic for validation
-    finalize_response = requests.post(_FINALIZE_URL, json=finalize_request_json, headers=headers_json)
+    finalize_response = requests.post(PIXELTABLE_API_URL, json=finalize_request_json, headers=headers_json)
     if finalize_response.status_code != 200:
         raise excs.Error(f'Error finalizing snapshot: {finalize_response.text}')
     finalize_response_json = finalize_response.json()
@@ -66,6 +65,18 @@ def publish_snapshot(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
     return confirmed_tbl_uri
 
 
+def clone_snapshot(dest_tbl_uri: str) -> list[FullTableMd]:
+    headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
+    clone_request_json = {'operation_type': 'clone_snapshot', 'table_uri': dest_tbl_uri}
+    response = requests.post(PIXELTABLE_API_URL, json=clone_request_json, headers=headers_json)
+    if response.status_code != 200:
+        raise excs.Error(f'Error cloning snapshot: {response.text}')
+    response_json = response.json()
+    if not isinstance(response_json, dict) or 'table_uri' not in response_json:
+        raise excs.Error(f'Unexpected response from server.\n{response_json}')
+    return [FullTableMd.from_dict(t) for t in response_json['md']['tables']]
+
+
 def _upload_bundle_to_s3(bundle: Path, parsed_location: urllib.parse.ParseResult) -> None:
     from pixeltable.utils.s3 import get_client
 

pixeltable/store.py

@@ -16,6 +16,7 @@ from pixeltable import catalog, exceptions as excs, exprs
 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
 
@@ -99,9 +100,9 @@ class StoreBase:
 
         # v_min/v_max indices: speeds up base table scans needed to propagate a base table insert or delete
         idx_name = f'vmin_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using=Env.get().dbms.version_index_type))
         idx_name = f'vmax_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using=Env.get().dbms.version_index_type))
 
         self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
 
@@ -232,7 +233,6 @@ class StoreBase:
         assert col.tbl.id == self.tbl_version.id
         num_excs = 0
         num_rows = 0
-
         # 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 = [sql.Column(col.name, col.type, primary_key=True) for col in self.pk_columns()]
@@ -301,10 +301,13 @@ class StoreBase:
                 )
             log_explain(_logger, update_stmt, conn)
             conn.execute(update_stmt)
-
         finally:
-            tmp_tbl.drop(bind=conn)
-            self.sa_md.remove(tmp_tbl)
+
+            def remove_tmp_tbl() -> None:
+                self.sa_md.remove(tmp_tbl)
+                tmp_tbl.drop(bind=conn)
+
+            run_cleanup(remove_tmp_tbl, raise_error=True)
         return num_excs
 
     def insert_rows(

pixeltable/type_system.py

@@ -512,7 +512,7 @@ class StringType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.STRING, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -602,7 +602,7 @@ class TimestampType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.TIMESTAMP, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -768,7 +768,7 @@ class JsonType(ColumnType):
         a_type = a.get('type')
         b_type = b.get('type')
 
-        if a_type in {'string', 'integer', 'number', 'boolean', 'object', 'array'} and a_type == b_type:
+        if a_type in ('string', 'integer', 'number', 'boolean', 'object', 'array') and a_type == b_type:
             # a and b both have the same type designation, but are not identical. This can happen if
             # (for example) they have validators or other attributes that differ. In this case, we
             # generalize to {'type': t}, where t is their shared type, with no other qualifications.
@@ -1170,6 +1170,20 @@ class DocumentType(ColumnType):
         XML = 3
         TXT = 4
 
+        @classmethod
+        def from_extension(cls, ext: str) -> Optional['DocumentType.DocumentFormat']:
+            if ext in ('.htm', '.html'):
+                return cls.HTML
+            if ext == '.md':
+                return cls.MD
+            if ext == '.pdf':
+                return cls.PDF
+            if ext == '.xml':
+                return cls.XML
+            if ext == '.txt':
+                return cls.TXT
+            return None
+
     def __init__(self, nullable: bool = False, doc_formats: Optional[str] = None):
         super().__init__(self.Type.DOCUMENT, nullable=nullable)
         self.doc_formats = doc_formats
@@ -1203,9 +1217,7 @@ class DocumentType(ColumnType):
         assert isinstance(val, str)
         from pixeltable.utils.documents import get_document_handle
 
-        dh = get_document_handle(val)
-        if dh is None:
-            raise excs.Error(f'Not a recognized document format: {val}')
+        _ = get_document_handle(val)
 
 
 T = typing.TypeVar('T')
@@ -1240,7 +1252,7 @@ class _PxtType:
     `ColumnType`.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         raise TypeError(f'Type `{type(self)}` cannot be instantiated.')
 
     @classmethod

pixeltable/utils/console_output.py

@@ -1,4 +1,5 @@
 import logging
+from typing import TextIO
 
 
 def map_level(verbosity: int) -> int:
@@ -22,10 +23,10 @@ def map_level(verbosity: int) -> int:
 
 
 class ConsoleOutputHandler(logging.StreamHandler):
-    def __init__(self, stream):
+    def __init__(self, stream: TextIO):
         super().__init__(stream)
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         if record.msg.endswith('\n'):
             self.stream.write(record.msg)
         else: