feat(storage): per-warehouse Iceberg catalog and warehouse-scoped URIs

amber/src/main/python/core/storage/document_factory.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import re
 import typing
 import urllib
 from typing import Optional
@@ -67,21 +68,27 @@ def sanitize_uri_path(uri):
         :param uri: Result of urllib.parse.urlparse(). Could be quoted.
         :return: Unquoted and sanitized format of uri.
         """
-        return urllib.parse.unquote(uri.path).lstrip("/").replace("/", "_")
+        path = urllib.parse.unquote(uri.path).lstrip("/")
+        # Strip the optional leading "wh/<name>/" segment so the storage key is
+        # identical regardless of which warehouse the URI targets (matches
+        # DocumentFactory.scala).
+        path = re.sub(r"^wh/[^/]+/", "", path)
+        return path.replace("/", "_")
 
     @staticmethod
     def create_document(uri: str, schema: Schema) -> VirtualDocument:
         parsed_uri = urlparse(uri)
         if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
-            resource_type = VFSURIFactory.decode_uri(uri).resource_type
-            namespace = DocumentFactory._resolve_namespace(resource_type)
+            components = VFSURIFactory.decode_uri(uri)
+            namespace = DocumentFactory._resolve_namespace(components.resource_type)
             storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+            warehouse = components.warehouse
             # Convert Amber Schema to Iceberg Schema with LARGE_BINARY
             # field name encoding
             iceberg_schema = amber_schema_to_iceberg_schema(schema)
 
             create_table(
-                IcebergCatalogInstance.get_instance(),
+                IcebergCatalogInstance.get_instance(warehouse),
                 namespace,
                 storage_key,
                 iceberg_schema,
@@ -94,6 +101,7 @@ def create_document(uri: str, schema: Schema) -> VirtualDocument:
                 iceberg_schema,
                 amber_tuples_to_arrow_table,
                 arrow_table_to_amber_tuples,
+                warehouse,
             )
 
         else:
@@ -116,10 +124,11 @@ def document_exists(uri: str) -> bool:
         """
         parsed_uri = urlparse(uri)
         if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
-            resource_type = VFSURIFactory.decode_uri(uri).resource_type
-            namespace = DocumentFactory._resolve_namespace(resource_type)
+            components = VFSURIFactory.decode_uri(uri)
+            namespace = DocumentFactory._resolve_namespace(components.resource_type)
             storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
-            return IcebergCatalogInstance.get_instance().table_exists(
+            warehouse = components.warehouse
+            return IcebergCatalogInstance.get_instance(warehouse).table_exists(
                 f"{namespace}.{storage_key}"
             )
 
@@ -131,12 +140,13 @@ def document_exists(uri: str) -> bool:
     def open_document(uri: str) -> typing.Tuple[VirtualDocument, Optional[Schema]]:
         parsed_uri = urlparse(uri)
         if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
-            resource_type = VFSURIFactory.decode_uri(uri).resource_type
-            namespace = DocumentFactory._resolve_namespace(resource_type)
+            components = VFSURIFactory.decode_uri(uri)
+            namespace = DocumentFactory._resolve_namespace(components.resource_type)
             storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
+            warehouse = components.warehouse
 
             table = load_table_metadata(
-                IcebergCatalogInstance.get_instance(),
+                IcebergCatalogInstance.get_instance(warehouse),
                 namespace,
                 storage_key,
             )
@@ -152,6 +162,7 @@ def open_document(uri: str) -> typing.Tuple[VirtualDocument, Optional[Schema]]:
                 table.schema(),
                 amber_tuples_to_arrow_table,
                 arrow_table_to_amber_tuples,
+                warehouse,
             )
             return document, amber_schema
 

amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py

@@ -27,49 +27,59 @@
 
 class IcebergCatalogInstance:
     """
-    IcebergCatalogInstance is a singleton that manages the Iceberg catalog instance.
-    Supports postgres SQL catalog and REST catalog.
-    - Provides a single shared catalog for all Iceberg table-related operations.
-    - Lazily initializes the catalog on first access.
-    - Supports replacing the catalog instance for testing or reconfiguration.
+    Manages Iceberg catalog instances, cached per warehouse.
+    - REST catalogs are keyed by warehouse so one process can read/write tables in
+      many warehouses (Design 2, warehouse-per-execution).
+    - The postgres catalog has no warehouse concept and shares a single entry.
+    - Catalogs are lazily created on first access; entries can be replaced for
+      testing or reconfiguration.
     """
 
-    _instance: Optional[Catalog] = None
+    _catalogs: dict = {}
+    _POSTGRES_KEY = "__postgres__"
 
     @classmethod
-    def get_instance(cls):
+    def get_instance(cls, warehouse: Optional[str] = None) -> Catalog:
         """
-        Retrieves the singleton Iceberg catalog instance.
-        - If the catalog is not initialized, it is lazily created using the configured
-        properties.
-        - Supports "postgres" and "rest" catalog types.
+        Retrieves the Iceberg catalog for the given warehouse, creating and caching
+        it on first use. For REST catalogs, `warehouse` selects which warehouse's
+        catalog to use (defaults to the configured warehouse when None). For the
+        postgres catalog, `warehouse` is ignored.
+        :param warehouse: the warehouse name (REST only); None uses the default.
         :return: the Iceberg catalog instance.
         """
-        if cls._instance is None:
-            catalog_type = StorageConfig.ICEBERG_CATALOG_TYPE
-            if catalog_type == "postgres":
-                cls._instance = create_postgres_catalog(
+        catalog_type = StorageConfig.ICEBERG_CATALOG_TYPE
+        if catalog_type == "postgres":
+            if cls._POSTGRES_KEY not in cls._catalogs:
+                cls._catalogs[cls._POSTGRES_KEY] = create_postgres_catalog(
                     "texera_iceberg",
                     StorageConfig.ICEBERG_FILE_STORAGE_DIRECTORY_PATH,
                     StorageConfig.ICEBERG_POSTGRES_CATALOG_URI_WITHOUT_SCHEME,
                     StorageConfig.ICEBERG_POSTGRES_CATALOG_USERNAME,
                     StorageConfig.ICEBERG_POSTGRES_CATALOG_PASSWORD,
                 )
-            elif catalog_type == "rest":
-                cls._instance = create_rest_catalog(
+            return cls._catalogs[cls._POSTGRES_KEY]
+        elif catalog_type == "rest":
+            key = warehouse or StorageConfig.ICEBERG_REST_CATALOG_WAREHOUSE_NAME
+            if key not in cls._catalogs:
+                cls._catalogs[key] = create_rest_catalog(
                     "texera_iceberg",
-                    StorageConfig.ICEBERG_REST_CATALOG_WAREHOUSE_NAME,
+                    key,
                     StorageConfig.ICEBERG_REST_CATALOG_URI,
                 )
-            else:
-                raise ValueError(f"Unsupported catalog type: {catalog_type}")
-        return cls._instance
+            return cls._catalogs[key]
+        else:
+            raise ValueError(f"Unsupported catalog type: {catalog_type}")
 
     @classmethod
-    def replace_instance(cls, catalog: Catalog):
+    def replace_instance(cls, catalog: Catalog, warehouse: Optional[str] = None):
         """
-        Replaces the existing Iceberg catalog instance.
-        - This method is useful for testing or dynamically updating the catalog.
-        :param catalog: the new Iceberg catalog instance to replace the current one.
+        Replaces the cached catalog for a warehouse (testing or reconfiguration).
+        :param catalog: the new Iceberg catalog instance.
+        :param warehouse: the warehouse to replace (REST only); None uses default.
         """
-        cls._instance = catalog
+        if StorageConfig.ICEBERG_CATALOG_TYPE == "postgres":
+            key = cls._POSTGRES_KEY
+        else:
+            key = warehouse or StorageConfig.ICEBERG_REST_CATALOG_WAREHOUSE_NAME
+        cls._catalogs[key] = catalog

amber/src/main/python/core/storage/iceberg/iceberg_document.py

@@ -63,6 +63,7 @@ def __init__(
         table_schema: Schema,
         serde: Callable[[Schema, Iterable[T]], pa.Table],
         deserde: Callable[[Schema, pa.Table], Iterable[T]],
+        warehouse: Optional[str] = None,
     ):
         self.table_namespace = table_namespace
         self.table_name = table_name
@@ -71,7 +72,7 @@ def __init__(
         self.deserde = deserde
 
         self.lock = rwlock.RWLockFair()
-        self.catalog = IcebergCatalogInstance.get_instance()
+        self.catalog = IcebergCatalogInstance.get_instance(warehouse)
 
     def get_uri(self) -> ParseResult:
         """Returns the URI of the table location."""

amber/src/main/python/core/storage/vfs_uri_factory.py

@@ -17,6 +17,7 @@
 
 from enum import Enum
 from typing import NamedTuple, Optional
+import re
 from urllib.parse import urlparse
 
 from core.util.virtual_identity import (
@@ -37,6 +38,10 @@ class VFSResourceType(str, Enum):
     STATE = "state"
 
 
+# See VFSURIFactory._is_valid_warehouse_name.
+_WAREHOUSE_NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]*")
+
+
 class VFSUriComponents(NamedTuple):
     """The named components encoded in a VFS URI, as returned by
     `VFSURIFactory.decode_uri`. A NamedTuple, so positional unpacking keeps
@@ -46,6 +51,10 @@ class VFSUriComponents(NamedTuple):
     execution_id: ExecutionIdentity
     global_port_id: Optional[GlobalPortIdentity]
     resource_type: VFSResourceType
+    # The warehouse whose catalog holds this URI's table, from the optional leading
+    # "/wh/<name>" segment; None for non-BYO storage, which uses the configured
+    # default. Last so positional unpacking of the earlier fields still works.
+    warehouse: Optional[str] = None
 
 
 class VFSURIFactory:
@@ -90,15 +99,62 @@ def extract_value(key: str) -> str:
             execution_id,
             global_port_id,
             resource_type,
+            VFSURIFactory._warehouse_from(segments),
         )
 
     @staticmethod
-    def create_port_base_uri(workflow_id, execution_id, global_port_id) -> str:
+    def _is_valid_warehouse_name(name: str) -> bool:
+        """A warehouse name becomes a URI path segment, so it may not carry
+        characters that have meaning there: no "/" to add segments, and no "%" to
+        smuggle one in percent-encoded form. Mirrors VFSURIFactory (Scala).
+        """
+        return _WAREHOUSE_NAME_RE.fullmatch(name) is not None
+
+    @staticmethod
+    def _warehouse_from(segments: list) -> Optional[str]:
+        """
+        The warehouse encoded in a VFS URI, if present. Reported as part of
+        VFSUriComponents by decode_uri, which is the only way in. Mirrors
+        VFSURIFactory.warehouseFrom (Scala).
+
+        Anchored to the leading segment: a later segment that happens to be "wh"
+        -- e.g. inside a user-chosen operator id -- must not select a warehouse; it
+        would disagree with document_factory.sanitize_uri_path, which strips only a
+        leading one, and would route the write to another user's warehouse. The
+        segments come from the RAW path (see decode_uri), so a percent-encoded slash
+        stays inside its own segment, and the name must be a legal warehouse name,
+        so anything create_port_base_uri could not have written resolves to no
+        warehouse rather than to a wrong one.
+        """
+        if (
+            len(segments) >= 2
+            and segments[0] == "wh"
+            and VFSURIFactory._is_valid_warehouse_name(segments[1])
+        ):
+            return segments[1]
+        return None
+
+    @staticmethod
+    def create_port_base_uri(
+        workflow_id, execution_id, global_port_id, warehouse: Optional[str] = None
+    ) -> str:
         """Base URI for a port. Result and state URIs derive from it via
         `result_uri` / `state_uri`.
+
+        `warehouse` is written as the leading "/wh/<name>" segment, mirroring the
+        Scala side; when None the URI is byte-for-byte what it was before warehouses
+        existed.
         """
+        if warehouse is not None and not VFSURIFactory._is_valid_warehouse_name(
+            warehouse
+        ):
+            raise ValueError(
+                f"warehouse name must match {_WAREHOUSE_NAME_RE.pattern} "
+                f"(it becomes a URI path segment): {warehouse}"
+            )
+        wh_segment = f"/wh/{warehouse}" if warehouse else ""
         return (
-            f"{VFSURIFactory.VFS_FILE_URI_SCHEME}:///wid/{workflow_id.id}"
+            f"{VFSURIFactory.VFS_FILE_URI_SCHEME}://{wh_segment}/wid/{workflow_id.id}"
             f"/eid/{execution_id.id}/globalportid/"
             f"{serialize_global_port_identity(global_port_id)}"
         )

amber/src/main/python/core/util/virtual_identity.py

@@ -64,24 +64,25 @@ def serialize_global_port_identity(obj: GlobalPortIdentity) -> str:
     ``(logicalOpId=<logicalOpId>,layerName=<layerName>,
     portId=<portId.id>,isInternal=<portId.internal>,isInput=<input>)``
 
-    Raises ValueError if `logicalOpId` or `layerName` contains an underscore
-    (VFS URI parsing relies on the absence of '_'), or if `portId` is negative.
+    Raises ValueError if `logicalOpId` or `layerName` contains an underscore or a
+    slash (VFS URI parsing relies on the absence of both), or if `portId` is
+    negative. Mirrors GlobalPortIdentitySerde (Scala).
     """
     logical_op_id = obj.op_id.logical_op_id.id
     layer_name = obj.op_id.layer_name
     port_id = obj.port_id.id
     is_internal = obj.port_id.internal
     is_input_port = obj.input
-    if "_" in logical_op_id:
-        raise ValueError(
-            f"logicalOpId must not contain '_' "
-            f"(VFS URI parsing relies on this): {logical_op_id}"
-        )
-    if "_" in layer_name:
-        raise ValueError(
-            f"layerName must not contain '_' "
-            f"(VFS URI parsing relies on this): {layer_name}"
-        )
+    for field, value in (("logicalOpId", logical_op_id), ("layerName", layer_name)):
+        # '_' would collide with the separator the storage key is built from, and
+        # '/' would add path segments to the VFS URI this string is interpolated
+        # into -- letting an id forge structure the URI never meant to have.
+        for forbidden in ("_", "/"):
+            if forbidden in value:
+                raise ValueError(
+                    f"{field} must not contain '{forbidden}' "
+                    f"(VFS URI parsing relies on this): {value}"
+                )
     if port_id < 0:
         raise ValueError(f"portId must be non-negative: {port_id}")
     return (

amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py

@@ -89,8 +89,8 @@ def sqlite_iceberg_catalog():
     catalog-agnostic, so the sqlite backend exercises the same code path.
 
     Module-scoped so all tests in this file share one warehouse, and so
-    namespace creation only happens once. We save/restore the original
-    `IcebergCatalogInstance` singleton so other test modules that expect
+    namespace creation only happens once. We save/restore the
+    `IcebergCatalogInstance` cache so other test modules that expect
     a real postgres-backed catalog (e.g. test_iceberg_document.py) are
     not affected by our replacement.
     """
@@ -117,7 +117,7 @@ def sqlite_iceberg_catalog():
             s3_large_binaries_base_uri="s3://texera-large-binaries/objects/0/",
         )
 
-    original_instance = IcebergCatalogInstance._instance
+    original_catalogs = dict(IcebergCatalogInstance._catalogs)
     db_path = f"{_WAREHOUSE_DIR}/catalog.sqlite"
     catalog = SqlCatalog(
         "texera_iceberg_e2e",
@@ -134,7 +134,8 @@ def sqlite_iceberg_catalog():
     try:
         yield catalog
     finally:
-        IcebergCatalogInstance.replace_instance(original_instance)
+        IcebergCatalogInstance._catalogs.clear()
+        IcebergCatalogInstance._catalogs.update(original_catalogs)
 
 
 def _fresh_base_uri() -> str: