pixeltable 0.3.7__py3-none-any.whl → 0.3.9__py3-none-any.whl

pixeltable/catalog/catalog.py

@@ -1,32 +1,96 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
 import logging
-from typing import Optional, Type
+import time
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 from uuid import UUID
 
+import psycopg
 import sqlalchemy as sql
 
-import pixeltable.env as env
 import pixeltable.exceptions as excs
 import pixeltable.metadata.schema as schema
 from pixeltable.env import Env
+from pixeltable.iterators import ComponentIterator
 
 from .dir import Dir
+from .globals import IfExistsParam, IfNotExistsParam, MediaValidation
+from .insertable_table import InsertableTable
+from .path import Path
 from .schema_object import SchemaObject
 from .table import Table
 from .table_version import TableVersion
 from .table_version_handle import TableVersionHandle
 from .table_version_path import TableVersionPath
+from .view import View
+
+if TYPE_CHECKING:
+    from .. import DataFrame, exprs
 
-# from .. import InsertableTable
 
 _logger = logging.getLogger('pixeltable')
 
 
-def _join_path(path: str, name: str) -> str:
-    """Append name to path, if path is not empty."""
-    return name if path == '' else f'{path}.{name}'
+def _lock_str(for_update: bool) -> str:
+    return 'X' if for_update else 'S'
+
+
+# TODO: remove once the concurrent update behavior has been debugged
+# def _debug_print(for_update: bool, msg: str) -> None:
+#     return
+#     print(f'{datetime.datetime.now()}: {_lock_str(for_update)}: {msg}')
+
+
+def _unpack_row(
+    row: Optional[sql.engine.Row], entities: list[type[sql.orm.decl_api.DeclarativeBase]]
+) -> Optional[list[Any]]:
+    """Convert a Row result into a list of entity instances.
+
+    Assumes that the query contains a select() of exactly those entities.
+    """
+    if row is None:
+        return None
+
+    result: list[sql.orm.decl_api.DeclarativeBase] = []
+    column_offset = 0
+
+    for entity in entities:
+        num_cols = len(entity.__table__.columns)
+        data = {name: row[column_offset + i] for i, name in enumerate(entity.__table__.columns.keys())}
+        inst = entity(**data)
+        result.append(inst)
+        column_offset += num_cols
+
+    return result
+
+
+_MAX_RETRIES = 3
+T = TypeVar('T')
+
+
+def _retry_loop(op: Callable[..., T]) -> Callable[..., T]:
+    @functools.wraps(op)
+    def loop(*args: Any, **kwargs: Any) -> T:
+        num_remaining_retries = _MAX_RETRIES
+        while True:
+            try:
+                # in order for retry to work, we need to make sure that there aren't any prior db updates
+                # that are part of an ongoing transaction
+                assert not Env.get().in_xact()
+                with Env.get().begin_xact() as conn:
+                    return op(*args, **kwargs)
+            except sql.exc.DBAPIError as e:
+                if isinstance(e.orig, psycopg.errors.SerializationFailure) and num_remaining_retries > 0:
+                    num_remaining_retries -= 1
+                    print(f'serialization failure:\n{e}')
+                    print('retrying ************************************************************')
+                    time.sleep(1)
+                else:
+                    raise
+
+    return loop
 
 
 class Catalog:
@@ -59,25 +123,20 @@ class Catalog:
         self._tbls = {}  # don't use a defaultdict here, it doesn't cooperate with the debugger
         self._init_store()
 
-    def get_dir_path(self, dir_id: UUID) -> str:
+    def get_dir_path(self, dir_id: UUID) -> Path:
         """Return path for directory with given id"""
-        session = env.Env.get().session
+        conn = Env.get().conn
         names: list[str] = []
         while True:
-            dir = session.query(schema.Dir).filter(schema.Dir.id == dir_id).one()
+            q = sql.select(schema.Dir).where(schema.Dir.id == dir_id)
+            # _debug_print(for_update=False, msg=f'dir id={dir_id}')
+            row = conn.execute(q).one()
+            dir = schema.Dir(**row._mapping)
             if dir.md['name'] == '':
                 break
             names.insert(0, dir.md['name'])
             dir_id = dir.parent_id
-            assert dir_id is not None
-        return '.'.join(names)
-
-    def get_tbl_path(self, tbl_id: UUID) -> str:
-        """Return path for table with given id"""
-        session = env.Env.get().session
-        tbl = session.query(schema.Table).filter(schema.Table.id == tbl_id).one()
-        dir_path = self.get_dir_path(tbl.dir_id)
-        return _join_path(dir_path, tbl.md['name'])
+        return Path('.'.join(names), empty_is_valid=True)
 
     @dataclasses.dataclass
     class DirEntry:
@@ -85,35 +144,154 @@ class Catalog:
         dir_entries: dict[str, Catalog.DirEntry]
         table: Optional[schema.Table]
 
-    def get_dir_contents(self, dir_id: UUID, recursive: bool = False) -> dict[str, DirEntry]:
+    @_retry_loop
+    def get_dir_contents(self, dir_path: Path, recursive: bool = False) -> dict[str, DirEntry]:
+        dir = self._get_schema_object(dir_path, expected=Dir, raise_if_not_exists=True)
+        return self._get_dir_contents(dir._id, recursive=recursive)
+
+    def _get_dir_contents(self, dir_id: UUID, recursive: bool = False) -> dict[str, DirEntry]:
         """Returns a dict mapping the entry names to DirEntry objects"""
-        session = env.Env.get().session
+        conn = Env.get().conn
         result: dict[str, Catalog.DirEntry] = {}
 
-        dirs = session.query(schema.Dir).filter(schema.Dir.parent_id == dir_id).all()
-        for dir in dirs:
+        q = sql.select(schema.Dir).where(schema.Dir.parent_id == dir_id)
+        # _debug_print(for_update=False, msg=f'dirs parent_id={dir_id}')
+        rows = conn.execute(q).all()
+        for row in rows:
+            dir = schema.Dir(**row._mapping)
             dir_contents: dict[str, Catalog.DirEntry] = {}
             if recursive:
-                dir_contents = self.get_dir_contents(dir.id, recursive=True)
+                dir_contents = self._get_dir_contents(dir.id, recursive=True)
             result[dir.md['name']] = self.DirEntry(dir=dir, dir_entries=dir_contents, table=None)
 
-        tbls = session.query(schema.Table).filter(schema.Table.dir_id == dir_id).all()
-        for tbl in tbls:
+        q = sql.select(schema.Table).where(schema.Table.dir_id == dir_id)
+        # _debug_print(for_update=False, msg=f'tbls parent_id={dir_id}')
+        rows = conn.execute(q).all()
+        for row in rows:
+            tbl = schema.Table(**row._mapping)
             result[tbl.md['name']] = self.DirEntry(dir=None, dir_entries={}, table=tbl)
 
         return result
 
-    def drop_dir(self, dir_id: UUID) -> None:
-        """Delete the directory with the given id"""
-        session = env.Env.get().session
-        session.query(schema.Dir).filter(schema.Dir.id == dir_id).delete()
+    @_retry_loop
+    def move(self, path: Path, new_path: Path) -> None:
+        _, dest_dir, src_obj = self._prepare_dir_op(
+            add_dir_path=new_path.parent,
+            add_name=new_path.name,
+            drop_dir_path=path.parent,
+            drop_name=path.name,
+            raise_if_exists=True,
+            raise_if_not_exists=True,
+        )
+        src_obj._move(new_path.name, dest_dir._id)
+
+    def _prepare_dir_op(
+        self,
+        add_dir_path: Optional[Path] = None,
+        add_name: Optional[str] = None,
+        drop_dir_path: Optional[Path] = None,
+        drop_name: Optional[str] = None,
+        drop_expected: Optional[type[SchemaObject]] = None,
+        raise_if_exists: bool = False,
+        raise_if_not_exists: bool = False,
+    ) -> tuple[Optional[SchemaObject], Optional[SchemaObject], Optional[SchemaObject]]:
+        """
+        Validates paths and acquires locks needed for a directory operation, ie, add/drop/rename (add + drop) of a
+        directory entry.
+
+        The target entry is either a table or directory. The directory operation can include
+        - adding an entry (<add_dir_path>.<add_name>)
+        - dropping an entry (<drop_dir_path>.<drop_name>)
 
-    def get_schema_object(
+        Returns: (existing SchemaObject of add path, Dir of add path, existing SchemaObject of drop path)
+
+        Locking protocol:
+        - X locks on the immediate parent directories of the added/dropped entries; this prevents concurrent
+          modifications of the parent
+        - lock parent before child
+        - if both add and drop (= two directories are involved), lock the directories in a pre-determined order
+          (in this case, by name) in order to prevent deadlocks between concurrent directory modifications
+        """
+        assert (add_dir_path is None) == (add_name is None)
+        assert (drop_dir_path is None) == (drop_name is None)
+        dir_paths: set[Path] = set()
+        if add_dir_path is not None:
+            dir_paths.add(add_dir_path)
+        if drop_dir_path is not None:
+            dir_paths.add(drop_dir_path)
+
+        add_dir: Optional[schema.Dir] = None
+        drop_dir: Optional[schema.Dir] = None
+        for p in sorted(list(dir_paths)):
+            dir = self._get_dir(p, for_update=True)
+            if dir is None:
+                raise excs.Error(f'Directory {str(p)!r} does not exist')
+            if p == add_dir_path:
+                add_dir = dir
+            if p == drop_dir_path:
+                drop_dir = dir
+
+        add_obj: Optional[SchemaObject] = None
+        if add_dir is not None:
+            add_obj = self._get_dir_entry(add_dir.id, add_name, for_update=True)
+            if add_obj is not None and raise_if_exists:
+                add_path = add_dir_path.append(add_name)
+                raise excs.Error(f'Path {str(add_path)!r} already exists')
+
+        drop_obj: Optional[SchemaObject] = None
+        if drop_dir is not None:
+            drop_path = drop_dir_path.append(drop_name)
+            drop_obj = self._get_dir_entry(drop_dir.id, drop_name, for_update=True)
+            if drop_obj is None and raise_if_not_exists:
+                raise excs.Error(f'Path {str(drop_path)!r} does not exist')
+            if drop_obj is not None and drop_expected is not None and not isinstance(drop_obj, drop_expected):
+                raise excs.Error(
+                    f'{str(drop_path)!r} needs to be a {drop_expected._display_name()} '
+                    f'but is a {type(drop_obj)._display_name()}'
+                )
+
+        add_dir_obj = Dir(add_dir.id, add_dir.parent_id, add_dir.md['name']) if add_dir is not None else None
+        return add_obj, add_dir_obj, drop_obj
+
+    def _get_dir_entry(self, dir_id: UUID, name: str, for_update: bool = False) -> Optional[SchemaObject]:
+        conn = Env.get().conn
+
+        # check for subdirectory
+        q = sql.select(schema.Dir).where(schema.Dir.parent_id == dir_id, schema.Dir.md['name'].astext == name)
+        if for_update:
+            q = q.with_for_update()
+        # _debug_print(for_update, f'dir name={name!r} parent={dir_id}')
+        # row = conn.execute(q).one_or_none()
+        # if row is not None:
+        #     dir_record = schema.Dir(**row._mapping)
+        #     return Dir(dir_record.id, dir_record.parent_id, name)
+        rows = conn.execute(q).all()
+        if len(rows) > 1:
+            assert False, rows
+        if len(rows) == 1:
+            dir_record = schema.Dir(**rows[0]._mapping)
+            return Dir(dir_record.id, dir_record.parent_id, name)
+
+        # check for table
+        q = sql.select(schema.Table.id).where(schema.Table.dir_id == dir_id, schema.Table.md['name'].astext == name)
+        if for_update:
+            q = q.with_for_update()
+        # _debug_print(for_update, f'table name={name!r} parent={dir_id}')
+        tbl_id = conn.execute(q).scalar_one_or_none()
+        if tbl_id is not None:
+            if not tbl_id in self._tbls:
+                self._tbls[tbl_id] = self._load_tbl(tbl_id)
+            return self._tbls[tbl_id]
+
+        return None
+
+    def _get_schema_object(
         self,
-        path: str,
-        expected: Optional[Type[SchemaObject]] = None,
+        path: Path,
+        expected: Optional[type[SchemaObject]] = None,
         raise_if_exists: bool = False,
         raise_if_not_exists: bool = False,
+        for_update: bool = False,
     ) -> Optional[SchemaObject]:
         """Return the schema object at the given path, or None if it doesn't exist.
 
@@ -123,42 +301,18 @@ class Catalog:
         - raise_if_not_exists is True and the path does not exist
         - expected is not None and the existing object has a different type
         """
-        session = env.Env.get().session
-        if path == '':
+        if path.is_root:
             # the root dir
             if expected is not None and expected is not Dir:
                 raise excs.Error(f'{path!r} needs to be a {expected._display_name()} but is a {Dir._display_name()}')
-            dir = self._get_dir(path)
+            dir = self._get_dir(path, for_update=for_update)
             return Dir(dir.id, dir.parent_id, dir.md['name'])
 
-        components = path.split('.')
-        parent_path = '.'.join(components[:-1])
-        parent_dir = self._get_dir('.'.join(components[:-1]))
+        parent_path = path.parent
+        parent_dir = self._get_dir(parent_path, for_update=False)
         if parent_dir is None:
             raise excs.Error(f'Directory {parent_path!r} does not exist')
-        name = components[-1]
-
-        # check if path points to a directory
-        obj: Optional[SchemaObject] = None
-        dir = (
-            session.query(schema.Dir)
-            .filter(schema.Dir.parent_id == parent_dir.id, schema.Dir.md['name'].astext == name)
-            .one_or_none()
-        )
-        if dir is not None:
-            obj = Dir(dir.id, dir.parent_id, dir.md['name'])
-        else:
-            # check if it's a table
-            row = (
-                session.query(schema.Table.id)
-                .filter(schema.Table.dir_id == parent_dir.id, schema.Table.md['name'].astext == name)
-                .one_or_none()
-            )
-            if row is not None:
-                tbl_id = row[0]
-                if not tbl_id in self._tbls:
-                    self._tbls[tbl_id] = self._load_tbl(tbl_id)
-                obj = self._tbls[tbl_id]
+        obj = self._get_dir_entry(parent_dir.id, path.name, for_update=for_update)
 
         if obj is None and raise_if_not_exists:
             raise excs.Error(f'Path {path!r} does not exist')
@@ -168,7 +322,7 @@ class Catalog:
             raise excs.Error(f'{path!r} needs to be a {expected._display_name()} but is a {type(obj)._display_name()}')
         return obj
 
-    def get_tbl(self, tbl_id: UUID) -> Optional[Table]:
+    def get_table_by_id(self, tbl_id: UUID) -> Optional[Table]:
         if not tbl_id in self._tbls:
             tbl = self._load_tbl(tbl_id)
             if tbl is None:
@@ -176,21 +330,225 @@ class Catalog:
             self._tbls[tbl_id] = tbl
         return self._tbls[tbl_id]
 
-    def add_tbl(self, tbl: Table) -> None:
-        """Explicitly add a Table"""
+    @_retry_loop
+    def create_table(
+        self,
+        path: Path,
+        schema: dict[str, Any],
+        df: 'DataFrame',
+        if_exists: IfExistsParam,
+        primary_key: Optional[list[str]],
+        num_retained_versions: int,
+        comment: str,
+        media_validation: MediaValidation,
+    ) -> Table:
+        existing = self._handle_path_collision(path, InsertableTable, False, if_exists)
+        if existing is not None:
+            assert isinstance(existing, Table)
+            return existing
+
+        dir = self._get_schema_object(path.parent, expected=Dir, raise_if_not_exists=True)
+        assert dir is not None
+
+        tbl = InsertableTable._create(
+            dir._id,
+            path.name,
+            schema,
+            df,
+            primary_key=primary_key,
+            num_retained_versions=num_retained_versions,
+            comment=comment,
+            media_validation=media_validation,
+        )
         self._tbls[tbl._id] = tbl
+        return tbl
+
+    @_retry_loop
+    def create_view(
+        self,
+        path: Path,
+        base: TableVersionPath,
+        select_list: Optional[list[tuple[exprs.Expr, Optional[str]]]],
+        where: Optional[exprs.Expr],
+        additional_columns: Optional[dict[str, Any]],
+        is_snapshot: bool,
+        iterator: Optional[tuple[type[ComponentIterator], dict[str, Any]]],
+        num_retained_versions: int,
+        comment: str,
+        media_validation: MediaValidation,
+        if_exists: IfExistsParam,
+    ) -> Table:
+        from pixeltable.utils.filecache import FileCache
+
+        existing = self._handle_path_collision(path, View, is_snapshot, if_exists)
+        if existing is not None:
+            assert isinstance(existing, View)
+            return existing
+
+        dir = self._get_schema_object(path.parent, expected=Dir, raise_if_not_exists=True)
+        assert dir is not None
+        if iterator is None:
+            iterator_class, iterator_args = None, None
+        else:
+            iterator_class, iterator_args = iterator
+        view = View._create(
+            dir._id,
+            path.name,
+            base=base,
+            select_list=select_list,
+            additional_columns=additional_columns,
+            predicate=where,
+            is_snapshot=is_snapshot,
+            iterator_cls=iterator_class,
+            iterator_args=iterator_args,
+            num_retained_versions=num_retained_versions,
+            comment=comment,
+            media_validation=media_validation,
+        )
+        FileCache.get().emit_eviction_warnings()
+        self._tbls[view._id] = view
+        return view
+
+    @_retry_loop
+    def get_table(self, path: Path) -> Table:
+        obj = Catalog.get()._get_schema_object(path, expected=Table, raise_if_not_exists=True)
+        assert isinstance(obj, Table)
+        obj._tbl_version.get().ensure_md_loaded()
+        return obj
+
+    @_retry_loop
+    def drop_table(self, path: Path, if_not_exists: IfNotExistsParam, force: bool) -> None:
+        _, _, src_obj = self._prepare_dir_op(
+            drop_dir_path=path.parent,
+            drop_name=path.name,
+            drop_expected=Table,
+            raise_if_not_exists=if_not_exists == IfNotExistsParam.ERROR and not force,
+        )
+        if src_obj is None:
+            _logger.info(f'Skipped table {str(path)!r} (does not exist).')
+            return
+        assert isinstance(src_obj, Table)
+        self._drop_tbl(src_obj, force=force, is_replace=False)
+
+    def _drop_tbl(self, tbl: Table, force: bool, is_replace: bool) -> None:
+        """
+        Drop the table (and recursively its views, if force == True).
 
-    def get_views(self, tbl_id: UUID) -> list[UUID]:
+        Locking protocol:
+        - X-lock base before X-locking any view
+        - deadlock-free wrt to TableVersion.insert() (insert propagation also proceeds top-down)
+        - X-locks parent dir prior to calling TableVersion.drop(): prevent concurrent creation of another SchemaObject
+          in the same directory with the same name (which could lead to duplicate names if we get rolled back)
+        """
+        view_ids = self.get_view_ids(tbl._id, for_update=True)
+        if len(view_ids) > 0:
+            if not force:
+                is_snapshot = tbl._tbl_version_path.is_snapshot()
+                obj_type_str = 'Snapshot' if is_snapshot else tbl._display_name().capitalize()
+                msg: str
+                if is_replace:
+                    msg = (
+                        f'{obj_type_str} {tbl._path()} already exists and has dependents. '
+                        "Use `if_exists='replace_force'` to replace it."
+                    )
+                else:
+                    msg = f'{obj_type_str} {tbl._path()} has dependents.'
+                raise excs.Error(msg)
+
+            for view_id in view_ids:
+                view = self.get_table_by_id(view_id)
+                self._drop_tbl(view, force=force, is_replace=is_replace)
+
+        _ = self.get_dir(tbl._dir_id, for_update=True)  # X-lock the parent directory
+        tbl._drop()
+        assert tbl._id in self._tbls
+        del self._tbls[tbl._id]
+        _logger.info(f'Dropped table `{tbl._path()}`.')
+
+    @_retry_loop
+    def create_dir(self, path: Path, if_exists: IfExistsParam, parents: bool) -> Dir:
+        # existing = self._handle_path_collision(path, Dir, False, if_exists)
+        # if existing is not None:
+        #     assert isinstance(existing, Dir)
+        #     return existing
+        #
+        # parent = self._get_schema_object(path.parent)
+        # assert parent is not None
+        # dir = Dir._create(parent._id, path.name)
+        # Env.get().console_logger.info(f'Created directory {path!r}.')
+        # return dir
+
+        if parents:
+            # start walking down from the root
+            last_parent: Optional[SchemaObject] = None
+            for ancestor in path.ancestors():
+                ancestor_obj = self._get_schema_object(ancestor, expected=Dir)
+                assert ancestor_obj is not None or last_parent is not None
+                last_parent = Dir._create(last_parent._id, ancestor.name) if ancestor_obj is None else ancestor_obj
+            parent = last_parent
+        else:
+            parent = self._get_schema_object(path.parent)
+        existing = self._handle_path_collision(path, Dir, False, if_exists)
+        if existing is not None:
+            assert isinstance(existing, Dir)
+            return existing
+        assert parent is not None
+        dir = Dir._create(parent._id, path.name)
+        Env.get().console_logger.info(f'Created directory {str(path)!r}.')
+        return dir
+
+    @_retry_loop
+    def drop_dir(self, path: Path, if_not_exists: IfNotExistsParam, force: bool) -> None:
+        _, _, schema_obj = self._prepare_dir_op(
+            drop_dir_path=path.parent,
+            drop_name=path.name,
+            drop_expected=Dir,
+            raise_if_not_exists=if_not_exists == IfNotExistsParam.ERROR and not force,
+        )
+        if schema_obj is None:
+            _logger.info(f'Directory {str(path)!r} does not exist; skipped drop_dir().')
+            return
+        self._drop_dir(schema_obj._id, path, force=force)
+
+    def _drop_dir(self, dir_id: UUID, dir_path: Path, force: bool = False) -> None:
+        conn = Env.get().conn
+        if not force:
+            # check for existing entries
+            q = sql.select(sql.func.count()).select_from(schema.Dir).where(schema.Dir.parent_id == dir_id)
+            num_subdirs = conn.execute(q).scalar()
+            q = sql.select(sql.func.count()).select_from(schema.Table).where(schema.Table.dir_id == dir_id)
+            num_tbls = conn.execute(q).scalar()
+            if num_subdirs + num_tbls > 0:
+                raise excs.Error(f'Directory {dir_path!r} is not empty.')
+
+        # drop existing subdirs
+        dir_q = sql.select(schema.Dir).where(schema.Dir.parent_id == dir_id).with_for_update()
+        for row in conn.execute(dir_q).all():
+            self._drop_dir(row.id, dir_path.append(row.md['name']), force=True)
+
+        # drop existing tables
+        tbl_q = sql.select(schema.Table).where(schema.Table.dir_id == dir_id).with_for_update()
+        for row in conn.execute(tbl_q).all():
+            tbl = self.get_table_by_id(row.id)
+            # this table would have been dropped already if it's a view of a base we dropped earlier
+            if tbl is not None:
+                self._drop_tbl(tbl, force=True, is_replace=False)
+
+        # self.drop_dir(dir_id)
+        # _debug_print(for_update=True, msg=f'drop dir id={dir_id}')
+        conn.execute(sql.delete(schema.Dir).where(schema.Dir.id == dir_id))
+        _logger.info(f'Removed directory {str(dir_path)!r}.')
+
+    def get_view_ids(self, tbl_id: UUID, for_update: bool = False) -> list[UUID]:
         """Return the ids of views that directly reference the given table"""
-        session = env.Env.get().session
-        q = session.query(schema.Table.id).filter(sql.text(f"md->'view_md'->'base_versions'->0->>0 = {tbl_id.hex!r}"))
-        result = [r[0] for r in q.all()]
+        conn = Env.get().conn
+        q = sql.select(schema.Table.id).where(sql.text(f"md->'view_md'->'base_versions'->0->>0 = {tbl_id.hex!r}"))
+        if for_update:
+            q = q.with_for_update()
+        # _debug_print(for_update=False, msg=f'views of tbl id={tbl_id}')
+        result = [r[0] for r in conn.execute(q).all()]
         return result
 
-    def remove_tbl(self, tbl_id: UUID) -> None:
-        assert tbl_id in self._tbls
-        del self._tbls[tbl_id]
-
     def get_tbl_version(self, tbl_id: UUID, effective_version: Optional[int]) -> Optional[TableVersion]:
         if (tbl_id, effective_version) not in self._tbl_versions:
             self._tbl_versions[(tbl_id, effective_version)] = self._load_tbl_version(tbl_id, effective_version)
@@ -208,41 +566,54 @@ class Catalog:
         assert (tbl_version.id, tbl_version.effective_version) in self._tbl_versions
         del self._tbl_versions[(tbl_version.id, tbl_version.effective_version)]
 
-    def get_dir(self, dir_id: UUID) -> Optional[Dir]:
+    def get_dir(self, dir_id: UUID, for_update: bool = False) -> Optional[Dir]:
         """Return the Dir with the given id, or None if it doesn't exist"""
-        session = env.Env.get().session
-        dir_record = session.query(schema.Dir).filter(schema.Dir.id == dir_id).one_or_none()
-        if dir_record is None:
+        conn = Env.get().conn
+        q = sql.select(schema.Dir).where(schema.Dir.id == dir_id)
+        if for_update:
+            q = q.with_for_update()
+        # _debug_print(for_update=False, msg=f'dir id={dir_id!r}')
+        row = conn.execute(q).one_or_none()
+        if row is None:
             return None
+        dir_record = schema.Dir(**row._mapping)
         return Dir(dir_record.id, dir_record.parent_id, dir_record.md['name'])
 
-    def _get_dir(self, path: str) -> Optional[schema.Dir]:
-        session = env.Env.get().session
-        assert session is not None
-        if path == '':
-            return session.query(schema.Dir).filter(schema.Dir.parent_id.is_(None)).one()
+    def _get_dir(self, path: Path, for_update: bool = False) -> Optional[schema.Dir]:
+        """
+        Locking protocol:
+        - S locks on all ancestors
+        - X lock on dir if for_update == True, otherwise also an S lock
+        """
+        conn = Env.get().conn
+        if path.is_root:
+            q = sql.select(schema.Dir).where(schema.Dir.parent_id.is_(None))
+            if for_update:
+                q = q.with_for_update()
+            # _debug_print(for_update, 'root dir')
+            row = conn.execute(q).one()
+            return schema.Dir(**row._mapping)
         else:
-            components = path.split('.')
-            parent_path = '.'.join(components[:-1])
-            parent_dir = self._get_dir(parent_path)
+            parent_dir = self._get_dir(path.parent, for_update=False)
             if parent_dir is None:
                 return None
-            name = components[-1]
-            dir = (
-                session.query(schema.Dir)
-                .filter(schema.Dir.parent_id == parent_dir.id, schema.Dir.md['name'].astext == name)
-                .one_or_none()
+            q = sql.select(schema.Dir).where(
+                schema.Dir.parent_id == parent_dir.id, schema.Dir.md['name'].astext == path.name
             )
-            return dir
+            if for_update:
+                q = q.with_for_update()
+            # _debug_print(for_update, f'dir {str(path)}')
+            row = conn.execute(q).one_or_none()
+            return schema.Dir(**row._mapping) if row is not None else None
 
     def _load_tbl(self, tbl_id: UUID) -> Optional[Table]:
         _logger.info(f'Loading table {tbl_id}')
         from .insertable_table import InsertableTable
         from .view import View
 
-        session = env.Env.get().session
-        tbl_record, schema_version_record = (
-            session.query(schema.Table, schema.TableSchemaVersion)
+        conn = Env.get().conn
+        q = (
+            sql.select(schema.Table, schema.TableSchemaVersion)
             .join(schema.TableSchemaVersion)
             .where(schema.Table.id == schema.TableSchemaVersion.tbl_id)
             # Table.md['current_schema_version'] == TableSchemaVersion.schema_version
@@ -253,10 +624,12 @@ class Catalog:
                 )
             )
             .where(schema.Table.id == tbl_id)
-            .one_or_none()
         )
-        if tbl_record is None:
+        # _debug_print(for_update=False, msg=f'load table id={tbl_id!r}')
+        row = conn.execute(q).one_or_none()
+        if row is None:
             return None
+        tbl_record, schema_version_record = _unpack_row(row, [schema.Table, schema.TableSchemaVersion])
 
         tbl_md = schema.md_from_dict(schema.TableMd, tbl_record.md)
         view_md = tbl_md.view_md
@@ -293,9 +666,9 @@ class Catalog:
 
     def _load_tbl_version(self, tbl_id: UUID, effective_version: Optional[int]) -> Optional[TableVersion]:
         _logger.info(f'Loading table version: {tbl_id}:{effective_version}')
-        session = env.Env.get().session
+        conn = Env.get().conn
         q = (
-            session.query(schema.Table, schema.TableSchemaVersion)
+            sql.select(schema.Table, schema.TableSchemaVersion)
             .select_from(schema.Table)
             .where(schema.Table.id == tbl_id)
             .join(schema.TableSchemaVersion)
@@ -337,19 +710,20 @@ class Catalog:
                 )
             )
 
-        tbl_record, schema_version_record = q.one_or_none()
+        row = conn.execute(q).one_or_none()
+        tbl_record, schema_version_record = _unpack_row(row, [schema.Table, schema.TableSchemaVersion])
         tbl_md = schema.md_from_dict(schema.TableMd, tbl_record.md)
         schema_version_md = schema.md_from_dict(schema.TableSchemaVersionMd, schema_version_record.md)
         view_md = tbl_md.view_md
 
         # load mutable view ids
-        q = session.query(schema.Table.id).filter(
+        q = sql.select(schema.Table.id).where(
             sql.text(
                 f"md->'view_md'->'base_versions'->0->>0 = {tbl_id.hex!r} "
                 "AND md->'view_md'->'base_versions'->0->1 IS NULL"
             )
         )
-        mutable_view_ids = [r[0] for r in q.all()]
+        mutable_view_ids = [r[0] for r in conn.execute(q).all()]
         mutable_views = [TableVersionHandle(id, None) for id in mutable_view_ids]
 
         if view_md is None:
@@ -384,8 +758,8 @@ class Catalog:
 
     def _init_store(self) -> None:
         """One-time initialization of the stored catalog. Idempotent."""
-        with env.Env.get().begin_xact():
-            session = env.Env.get().session
+        with Env.get().begin_xact():
+            session = Env.get().session
             if session.query(sql.func.count(schema.Dir.id)).scalar() > 0:
                 return
             # create a top-level directory, so that every schema object has a directory
@@ -393,5 +767,37 @@ class Catalog:
             dir_record = schema.Dir(parent_id=None, md=dataclasses.asdict(dir_md))
             session.add(dir_record)
             session.flush()
-            session.commit()
             _logger.info(f'Initialized catalog')
+
+    def _handle_path_collision(
+        self, path: Path, expected_obj_type: type[SchemaObject], expected_snapshot: bool, if_exists: IfExistsParam
+    ) -> Optional[SchemaObject]:
+        obj, _, _ = self._prepare_dir_op(add_dir_path=path.parent, add_name=path.name)
+
+        if if_exists == IfExistsParam.ERROR and obj is not None:
+            raise excs.Error(f'Path {path!r} is an existing {type(obj)._display_name()}')
+        else:
+            is_snapshot = isinstance(obj, View) and obj._tbl_version_path.is_snapshot()
+            if obj is not None and (not isinstance(obj, expected_obj_type) or (expected_snapshot and not is_snapshot)):
+                obj_type_str = 'snapshot' if expected_snapshot else expected_obj_type._display_name()
+                raise excs.Error(
+                    f'Path {path!r} already exists but is not a {obj_type_str}. Cannot {if_exists.name.lower()} it.'
+                )
+
+        if obj is None:
+            return None
+        if if_exists == IfExistsParam.IGNORE:
+            return obj
+
+        # drop the existing schema object
+        if isinstance(obj, Dir):
+            dir_contents = self._get_dir_contents(obj._id)
+            if len(dir_contents) > 0 and if_exists == IfExistsParam.REPLACE:
+                raise excs.Error(
+                    f'Directory {path!r} already exists and is not empty. Use `if_exists="replace_force"` to replace it.'
+                )
+            self._drop_dir(obj._id, path, force=True)
+        else:
+            assert isinstance(obj, Table)
+            self._drop_tbl(obj, force=if_exists == IfExistsParam.REPLACE_FORCE, is_replace=True)
+        return None