You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@texera.apache.org by gi...@apache.org on 2026/08/06 17:05:54 UTC
(texera) branch gh-readonly-queue/main/pr-6944-148b40c03628114a9f6cb17a117e9feeec9a878a created (now 15d82c7b9d)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6944-148b40c03628114a9f6cb17a117e9feeec9a878a
in repository https://gitbox.apache.org/repos/asf/texera.git
at 15d82c7b9d feat(storage): per-warehouse Iceberg catalog and warehouse-scoped URIs (#6944)
This branch includes the following new commits:
new 15d82c7b9d feat(storage): per-warehouse Iceberg catalog and warehouse-scoped URIs (#6944)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(texera) 01/01: feat(storage): per-warehouse Iceberg catalog and warehouse-scoped URIs (#6944)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6944-148b40c03628114a9f6cb17a117e9feeec9a878a
in repository https://gitbox.apache.org/repos/asf/texera.git
commit 15d82c7b9d2af862741f71653247c46f060473a7
Author: Meng Wang <me...@uci.edu>
AuthorDate: Thu Aug 6 09:31:07 2026 -0700
feat(storage): per-warehouse Iceberg catalog and warehouse-scoped URIs (#6944)
### What changes were proposed in this PR?
First slice of the per-user bring-your-own-S3 warehouse feature
(umbrella #6870). This is a
**foundation-only, backward-compatible** change that makes the storage
layer warehouse-aware
without changing any behavior yet:
- Replace the single shared Iceberg catalog with a **per-warehouse
cache** on both the JVM
(`IcebergCatalogInstance`) and Python (`iceberg_catalog_instance.py`)
sides, so one worker
process can hold several REST catalogs — one per warehouse it touches.
Only the REST
(Lakekeeper) catalog varies by warehouse; the hadoop/postgres catalogs
stay warehouse-agnostic.
- Encode the warehouse as a leading `/wh/<name>` segment in VFS URIs
(`VFSURIFactory`), so a URI
fully identifies where its tables live; the read path
(`DocumentFactory`) derives the warehouse
back out of the URI.
- Add an optional `warehouse` field to `WorkflowContext` to carry the
selection through an execution.
With no warehouse in play (the default), URIs and catalog behavior are
byte-for-byte identical to
today, so nothing changes for existing deployments.
### Any related issues, documentation, discussions?
Closes #6929. Part of #6870 (design discussions #5293 and #6040).
### How was this PR tested?
Adds unit tests for the new warehouse-aware behavior:
- `VFSURIFactorySpec` — `warehouseFromURI` extraction, the `/wh/<name>`
segment (present with a
warehouse, absent without), and that a warehouse-scoped URI still
round-trips through `decodeURI`.
- `DocumentFactorySpec` — creating and finding a document for a
warehouse-scoped URI against a local
Iceberg catalog (the `/wh/` prefix is stripped from the storage key).
- `test_vfs_uri_factory.py` — matching coverage for the Python
`warehouse_from_uri`.
The existing storage/Iceberg tests continue to cover the default
(no-warehouse) path unchanged.
The full build and test suite run in CI.
### Was this PR authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
---------
Co-authored-by: Claude Opus 4.8 (1M context) <no...@anthropic.com>
---
.../main/python/core/storage/document_factory.py | 31 +++--
.../storage/iceberg/iceberg_catalog_instance.py | 62 +++++----
.../core/storage/iceberg/iceberg_document.py | 3 +-
.../main/python/core/storage/vfs_uri_factory.py | 60 +++++++-
.../src/main/python/core/util/virtual_identity.py | 25 ++--
.../packaging/test_state_materialization_e2e.py | 9 +-
.../iceberg/test_iceberg_catalog_instance.py | 152 +++++++++++++++++++++
.../python/core/storage/test_vfs_uri_factory.py | 84 ++++++++++++
.../test/python/core/util/test_virtual_identity.py | 11 ++
.../amber/core/storage/DocumentFactory.scala | 26 ++--
.../core/storage/IcebergCatalogInstance.scala | 108 +++++++++------
.../texera/amber/core/storage/VFSURIFactory.scala | 101 ++++++++++++--
.../storage/result/iceberg/IcebergDocument.scala | 6 +-
.../amber/core/workflow/WorkflowContext.scala | 3 +-
.../amber/util/serde/GlobalPortIdentitySerde.scala | 11 ++
.../amber/core/storage/DocumentFactorySpec.scala | 27 +++-
.../core/storage/LocalHadoopIcebergCatalog.scala | 26 +++-
.../amber/core/storage/VFSURIFactorySpec.scala | 116 +++++++++++++++-
18 files changed, 732 insertions(+), 129 deletions(-)
diff --git a/amber/src/main/python/core/storage/document_factory.py b/amber/src/main/python/core/storage/document_factory.py
index 25e1d08211..0db1b41392 100644
--- a/amber/src/main/python/core/storage/document_factory.py
+++ b/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 @@ class DocumentFactory:
: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 @@ class DocumentFactory:
iceberg_schema,
amber_tuples_to_arrow_table,
arrow_table_to_amber_tuples,
+ warehouse,
)
else:
@@ -116,10 +124,11 @@ class DocumentFactory:
"""
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 @@ class DocumentFactory:
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 @@ class DocumentFactory:
table.schema(),
amber_tuples_to_arrow_table,
arrow_table_to_amber_tuples,
+ warehouse,
)
return document, amber_schema
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py b/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py
index 65a4f6beec..61987c2594 100644
--- a/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_catalog_instance.py
@@ -27,49 +27,59 @@ from core.storage.storage_config import StorageConfig
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
diff --git a/amber/src/main/python/core/storage/iceberg/iceberg_document.py b/amber/src/main/python/core/storage/iceberg/iceberg_document.py
index 7a5beda916..797c1ec4e4 100644
--- a/amber/src/main/python/core/storage/iceberg/iceberg_document.py
+++ b/amber/src/main/python/core/storage/iceberg/iceberg_document.py
@@ -63,6 +63,7 @@ class IcebergDocument(VirtualDocument[T]):
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 @@ class IcebergDocument(VirtualDocument[T]):
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."""
diff --git a/amber/src/main/python/core/storage/vfs_uri_factory.py b/amber/src/main/python/core/storage/vfs_uri_factory.py
index d0eb3eab7a..37a28cbbf0 100644
--- a/amber/src/main/python/core/storage/vfs_uri_factory.py
+++ b/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 @@ class VFSURIFactory:
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)}"
)
diff --git a/amber/src/main/python/core/util/virtual_identity.py b/amber/src/main/python/core/util/virtual_identity.py
index 6893e7e8f0..0a9de0ccc4 100644
--- a/amber/src/main/python/core/util/virtual_identity.py
+++ b/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 (
diff --git a/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py b/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py
index 5db9bdf903..84c626f58b 100644
--- a/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py
+++ b/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:
diff --git a/amber/src/test/python/core/storage/iceberg/test_iceberg_catalog_instance.py b/amber/src/test/python/core/storage/iceberg/test_iceberg_catalog_instance.py
new file mode 100644
index 0000000000..92aa565182
--- /dev/null
+++ b/amber/src/test/python/core/storage/iceberg/test_iceberg_catalog_instance.py
@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from unittest.mock import patch
+
+import pytest
+
+from core.storage.iceberg import iceberg_catalog_instance
+from core.storage.iceberg.iceberg_catalog_instance import IcebergCatalogInstance
+from core.storage.storage_config import StorageConfig
+
+
+@pytest.fixture(autouse=True)
+def _isolated_catalog_cache():
+ """Save/clear/restore the class-level cache so these tests neither see nor
+ leak cached catalogs (mirrors test_state_materialization_e2e.py)."""
+ original = dict(IcebergCatalogInstance._catalogs)
+ IcebergCatalogInstance._catalogs.clear()
+ yield
+ IcebergCatalogInstance._catalogs.clear()
+ IcebergCatalogInstance._catalogs.update(original)
+
+
+def _rest_config():
+ """StorageConfig patched to the `rest` catalog type. No live Lakekeeper is
+ involved anywhere in this file: `create_rest_catalog` is patched out, which
+ is exactly the boundary the cache under test sits on."""
+ return (
+ patch.object(StorageConfig, "ICEBERG_CATALOG_TYPE", "rest"),
+ patch.object(
+ StorageConfig, "ICEBERG_REST_CATALOG_WAREHOUSE_NAME", "default-wh"
+ ),
+ patch.object(
+ StorageConfig, "ICEBERG_REST_CATALOG_URI", "http://localhost:8181"
+ ),
+ )
+
+
+class TestRestCatalogCachedPerWarehouse:
+ """
+ Under the `rest` catalog type the cache is keyed by warehouse name -- the
+ core of the per-warehouse catalog cache (#6870 Phase 0): one process can
+ hold several REST catalogs, one per warehouse it touches, while repeated
+ lookups for the same warehouse share one client.
+ """
+
+ def test_caches_one_catalog_per_warehouse(self):
+ type_p, default_p, uri_p = _rest_config()
+ with (
+ type_p,
+ default_p,
+ uri_p,
+ patch.object(
+ iceberg_catalog_instance,
+ "create_rest_catalog",
+ side_effect=lambda *args: object(),
+ ) as mock_create,
+ ):
+ first = IcebergCatalogInstance.get_instance("wh-a")
+ again = IcebergCatalogInstance.get_instance("wh-a")
+ other = IcebergCatalogInstance.get_instance("wh-b")
+
+ assert first is again, "same warehouse must reuse the cached catalog"
+ assert first is not other, "distinct warehouses must get distinct catalogs"
+ assert mock_create.call_count == 2
+ # The warehouse name is what create_rest_catalog is keyed/called with.
+ assert [call.args[1] for call in mock_create.call_args_list] == ["wh-a", "wh-b"]
+
+ def test_none_falls_back_to_the_configured_default_warehouse(self):
+ type_p, default_p, uri_p = _rest_config()
+ with (
+ type_p,
+ default_p,
+ uri_p,
+ patch.object(
+ iceberg_catalog_instance,
+ "create_rest_catalog",
+ side_effect=lambda *args: object(),
+ ) as mock_create,
+ ):
+ from_none = IcebergCatalogInstance.get_instance(None)
+ from_name = IcebergCatalogInstance.get_instance("default-wh")
+
+ assert from_none is from_name, "None must share the default warehouse's entry"
+ assert mock_create.call_count == 1
+ assert mock_create.call_args.args[1] == "default-wh"
+
+ def test_replace_instance_keys_by_warehouse(self):
+ sentinel = object()
+ type_p, default_p, uri_p = _rest_config()
+ with (
+ type_p,
+ default_p,
+ uri_p,
+ patch.object(
+ iceberg_catalog_instance,
+ "create_rest_catalog",
+ side_effect=AssertionError(
+ "cache was pre-seeded; nothing should be created"
+ ),
+ ),
+ ):
+ IcebergCatalogInstance.replace_instance(sentinel, "wh-a")
+ assert IcebergCatalogInstance.get_instance("wh-a") is sentinel
+
+
+class TestWarehouseAgnosticCatalogTypes:
+ """The postgres catalog has no warehouse concept: every caller shares the
+ one entry under the constant key, mirroring the Scala `SharedCatalogKey`."""
+
+ def test_postgres_catalog_is_created_once_and_shared(self):
+ with (
+ patch.object(StorageConfig, "ICEBERG_CATALOG_TYPE", "postgres"),
+ patch.object(
+ iceberg_catalog_instance,
+ "create_postgres_catalog",
+ side_effect=lambda *args: object(),
+ ) as mock_create,
+ ):
+ first = IcebergCatalogInstance.get_instance()
+ again = IcebergCatalogInstance.get_instance("ignored-warehouse")
+
+ assert first is again, "postgres ignores the warehouse and shares one entry"
+ assert mock_create.call_count == 1
+
+ def test_replace_instance_uses_the_shared_postgres_key(self):
+ sentinel = object()
+ with patch.object(StorageConfig, "ICEBERG_CATALOG_TYPE", "postgres"):
+ IcebergCatalogInstance.replace_instance(sentinel)
+ assert IcebergCatalogInstance.get_instance() is sentinel
+ assert IcebergCatalogInstance._catalogs == {
+ IcebergCatalogInstance._POSTGRES_KEY: sentinel
+ }
+
+ def test_unsupported_catalog_type_raises(self):
+ with patch.object(StorageConfig, "ICEBERG_CATALOG_TYPE", "bogus"):
+ with pytest.raises(ValueError, match="Unsupported catalog type: bogus"):
+ IcebergCatalogInstance.get_instance()
diff --git a/amber/src/test/python/core/storage/test_vfs_uri_factory.py b/amber/src/test/python/core/storage/test_vfs_uri_factory.py
index 5422f2ad68..24ad583ae6 100644
--- a/amber/src/test/python/core/storage/test_vfs_uri_factory.py
+++ b/amber/src/test/python/core/storage/test_vfs_uri_factory.py
@@ -133,3 +133,87 @@ class TestDecodeUriErrorPaths:
def test_rejects_unknown_resource_type(self):
with pytest.raises(ValueError, match="Unknown resource type: bogus"):
VFSURIFactory.decode_uri("vfs:///wid/1/eid/1/bogus")
+
+
+class TestWarehouseInDecodedUri:
+ def test_reports_warehouse_from_wh_segment(self):
+ components = VFSURIFactory.decode_uri("vfs:///wh/user-2-foo/wid/7/eid/3/result")
+ assert components.warehouse == "user-2-foo"
+
+ def test_warehouse_is_none_when_no_wh_segment(self):
+ # Non-BYO URIs have no /wh/ segment, so the warehouse is absent.
+ assert VFSURIFactory.decode_uri("vfs:///wid/7/eid/3/result").warehouse is None
+
+ def test_only_a_leading_wh_segment_counts(self):
+ # A `wh` deeper in the path -- e.g. inside an operator id -- must not select
+ # a warehouse; it would disagree with sanitize_uri_path, which strips only a
+ # leading one, and would route the write to another user's warehouse.
+ assert (
+ VFSURIFactory.decode_uri(
+ "vfs:///wid/1/eid/2/opid/a/wh/victim/b/result"
+ ).warehouse
+ is None
+ )
+ assert (
+ VFSURIFactory.decode_uri("vfs:///wid/1/eid/2/opid/wh/result").warehouse
+ is None
+ )
+
+ def test_percent_encoded_name_is_rejected_not_decoded(self):
+ # The raw path is split and the name must be a legal warehouse name, so a
+ # percent-encoded name resolves to no warehouse rather than to a decoded
+ # one. Decoding first would let "%2F" become a separator and pick the wrong
+ # warehouse; it would also diverge from Scala, which splits its raw path.
+ assert (
+ VFSURIFactory.decode_uri(
+ "vfs:///wh/user-2%2Dfoo/wid/7/eid/3/result"
+ ).warehouse
+ is None
+ )
+
+ def test_encoded_slash_in_name_cannot_forge_segments(self):
+ # "%2F" must stay inside its own segment: decoded-then-split, this URI would
+ # read as /wh/a/wid/999/... handing back "a" as the warehouse, and the key
+ # search would find the injected `wid` instead of the real one.
+ components = VFSURIFactory.decode_uri(
+ "vfs:///wh/a%2Fwid%2F999/wid/1/eid/2/result"
+ )
+ assert components.warehouse is None
+ assert components.workflow_id.id == 1
+ assert components.execution_id.id == 2
+
+ def test_builder_rejects_an_unsafe_warehouse_name(self):
+ with pytest.raises(ValueError, match="warehouse name must match"):
+ VFSURIFactory.create_port_base_uri(
+ WorkflowIdentity(id=7), ExecutionIdentity(id=3), _gpi(), "a/b"
+ )
+
+ def test_round_trips_a_warehouse_written_by_the_python_builder(self):
+ # Python can now write the segment, not just read it.
+ uri = VFSURIFactory.create_port_base_uri(
+ WorkflowIdentity(id=7), ExecutionIdentity(id=3), _gpi(), "user-2-foo"
+ )
+ assert uri.startswith("vfs:///wh/user-2-foo/wid/7/eid/3/")
+ assert (
+ VFSURIFactory.decode_uri(VFSURIFactory.result_uri(uri)).warehouse
+ == "user-2-foo"
+ )
+
+ def test_builder_without_warehouse_is_unchanged(self):
+ uri = VFSURIFactory.create_port_base_uri(
+ WorkflowIdentity(id=7), ExecutionIdentity(id=3), _gpi()
+ )
+ assert "/wh/" not in uri
+ assert VFSURIFactory.decode_uri(VFSURIFactory.result_uri(uri)).warehouse is None
+
+ def test_decode_round_trips_a_warehouse_scoped_uri(self):
+ # The leading /wh/<name> segment must not break decode_uri, which finds
+ # wid/eid by key rather than by position.
+ components = VFSURIFactory.decode_uri(
+ "vfs:///wh/user-2-foo/wid/11/eid/22/result"
+ )
+ assert components.workflow_id.id == 11
+ assert components.execution_id.id == 22
+ assert components.global_port_id is None
+ assert components.resource_type == VFSResourceType.RESULT
+ assert components.warehouse == "user-2-foo"
diff --git a/amber/src/test/python/core/util/test_virtual_identity.py b/amber/src/test/python/core/util/test_virtual_identity.py
index 75c8e30a74..ad60391c28 100644
--- a/amber/src/test/python/core/util/test_virtual_identity.py
+++ b/amber/src/test/python/core/util/test_virtual_identity.py
@@ -159,6 +159,17 @@ class TestSerializeGlobalPortIdentity:
with pytest.raises(ValueError, match="layerName must not contain"):
serialize_global_port_identity(_gpi(layer="main_source_0_op"))
+ def test_rejects_slash_in_logical_op_id(self):
+ # A '/' would add path segments to the VFS URI this string is interpolated
+ # into, letting an id forge structure the URI never meant to have. Matches
+ # the guard in GlobalPortIdentitySerde (Scala).
+ with pytest.raises(ValueError, match="logicalOpId must not contain"):
+ serialize_global_port_identity(_gpi(op_id="a/wh/victim/b"))
+
+ def test_rejects_slash_in_layer_name(self):
+ with pytest.raises(ValueError, match="layerName must not contain"):
+ serialize_global_port_identity(_gpi(layer="main/wh/victim"))
+
def test_rejects_negative_port_id(self):
# Port ids are array indices and must be non-negative.
with pytest.raises(ValueError, match="portId must be non-negative"):
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/DocumentFactory.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/DocumentFactory.scala
index 910c1dce44..f84158e0e3 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/DocumentFactory.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/DocumentFactory.scala
@@ -38,7 +38,7 @@ object DocumentFactory {
val ICEBERG = "iceberg"
private def sanitizeURIPath(uri: URI): String =
- uri.getPath.stripPrefix("/").replace("/", "_")
+ uri.getPath.stripPrefix("/").replaceFirst("^wh/[^/]+/", "").replace("/", "_")
private def resolveNamespace(resourceType: VFSResourceType.Value): String =
resourceType match {
@@ -76,13 +76,15 @@ object DocumentFactory {
def createDocument(uri: URI, schema: Schema): VirtualDocument[_] = {
uri.getScheme match {
case VFS_FILE_URI_SCHEME =>
- val resourceType = decodeURI(uri).resourceType
+ val components = decodeURI(uri)
+ val warehouse = components.warehouse
+ val resourceType = components.resourceType
val storageKey = sanitizeURIPath(uri)
val namespace = resolveNamespace(resourceType)
val icebergSchema = IcebergUtil.toIcebergSchema(schema)
IcebergUtil.createTable(
- IcebergCatalogInstance.getInstance(),
+ IcebergCatalogInstance.getInstance(warehouse),
namespace,
storageKey,
icebergSchema,
@@ -97,7 +99,8 @@ object DocumentFactory {
storageKey,
icebergSchema,
serde,
- deserde
+ deserde,
+ warehouse
)
case unsupportedScheme =>
throw new UnsupportedOperationException(
@@ -119,11 +122,13 @@ object DocumentFactory {
def documentExists(uri: URI): Boolean = {
uri.getScheme match {
case VFS_FILE_URI_SCHEME =>
- val resourceType = decodeURI(uri).resourceType
+ val components = decodeURI(uri)
+ val warehouse = components.warehouse
+ val resourceType = components.resourceType
val storageKey = sanitizeURIPath(uri)
val namespace = resolveNamespace(resourceType)
IcebergCatalogInstance
- .getInstance()
+ .getInstance(warehouse)
.tableExists(TableIdentifier.of(namespace, storageKey))
case unsupportedScheme =>
@@ -165,13 +170,15 @@ object DocumentFactory {
uri.getScheme match {
case DATASET_FILE_URI_SCHEME => (new DatasetFileDocument(uri), None)
case VFS_FILE_URI_SCHEME =>
- val resourceType = decodeURI(uri).resourceType
+ val components = decodeURI(uri)
+ val warehouse = components.warehouse
+ val resourceType = components.resourceType
val storageKey = sanitizeURIPath(uri)
val namespace = resolveNamespace(resourceType)
val table = IcebergUtil
.loadTableMetadata(
- IcebergCatalogInstance.getInstance(),
+ IcebergCatalogInstance.getInstance(warehouse),
namespace,
storageKey
)
@@ -190,7 +197,8 @@ object DocumentFactory {
storageKey,
table.schema(),
serde,
- deserde
+ deserde,
+ warehouse
),
Some(amberSchema)
)
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/IcebergCatalogInstance.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/IcebergCatalogInstance.scala
index cd4b3c8796..313772b0fc 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/IcebergCatalogInstance.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/IcebergCatalogInstance.scala
@@ -23,57 +23,87 @@ import org.apache.texera.common.config.StorageConfig
import org.apache.texera.amber.util.IcebergUtil
import org.apache.iceberg.catalog.Catalog
+import scala.collection.mutable
+
/**
- * IcebergCatalogInstance is a singleton that manages the Iceberg catalog instance.
- * - Provides a single shared catalog for all Iceberg table-related operations in the Texera application.
- * - Lazily initializes the catalog on first access.
- * - Supports replacing the catalog instance primarily for testing or reconfiguration.
+ * IcebergCatalogInstance manages the Iceberg catalog clients used across the Texera application.
+ *
+ * Catalogs are cached per warehouse: each distinct warehouse name gets its own lazily-created
+ * catalog client, so a single JVM may hold several catalogs, one per warehouse it touches. Callers
+ * that do not specify a warehouse use the configured default, preserving single-warehouse (non-BYO)
+ * behavior.
+ *
+ * Only the REST catalog varies by warehouse; the hadoop and postgres catalogs are warehouse-agnostic
+ * and ignore the warehouse argument.
+ *
+ * Access is synchronized because the same JVM serves multiple warehouses concurrently.
*/
object IcebergCatalogInstance {
- private var instance: Option[Catalog] = None
+ private val catalogs = mutable.Map.empty[String, Catalog]
+
+ // Cache key for the warehouse-agnostic catalog types. Not a legal warehouse name,
+ // so it cannot collide with a REST warehouse.
+ private val SharedCatalogKey = "<shared>"
+
+ private def defaultWarehouse: String = StorageConfig.icebergRESTCatalogWarehouseName
/**
- * Retrieves the singleton Iceberg catalog instance.
- * - If the catalog is not initialized, it is lazily created using the configured properties.
+ * The cache key for a warehouse. Only the REST catalog is scoped to a warehouse;
+ * hadoop and postgres ignore it, so they must share one entry. Keying them by
+ * warehouse name would build a second, fully equivalent catalog per distinct name
+ * -- for postgres a second JdbcCatalog with its own connection pool, all pointing
+ * at the same database. Mirrors the Python side, which keys those under a constant.
+ */
+ private def cacheKey(warehouse: String): String =
+ StorageConfig.icebergCatalogType match {
+ case "rest" => warehouse
+ case _ => SharedCatalogKey
+ }
+
+ /**
+ * Retrieves the catalog for the given warehouse, creating and caching it on first access.
*
- * @return the Iceberg catalog instance.
+ * @param warehouse the warehouse to obtain a catalog for; `None` uses the configured
+ * default, mirroring the Python side's `Optional[str]`.
+ * @return the Iceberg catalog for that warehouse.
*/
- def getInstance(): Catalog = {
- instance match {
- case Some(catalog) => catalog
- case None =>
- val catalog = StorageConfig.icebergCatalogType match {
- case "hadoop" =>
- IcebergUtil.createHadoopCatalog(
- "texera_iceberg",
- StorageConfig.fileStorageDirectoryPath
- )
- case "rest" =>
- IcebergUtil.createRestCatalog(
- "texera_iceberg",
- StorageConfig.icebergRESTCatalogWarehouseName
- )
- case "postgres" =>
- IcebergUtil.createPostgresCatalog(
- "texera_iceberg",
- StorageConfig.fileStorageDirectoryPath
- )
- case unsupported =>
- throw new IllegalArgumentException(s"Unsupported catalog type: $unsupported")
- }
- instance = Some(catalog)
- catalog
+ def getInstance(warehouse: Option[String] = None): Catalog = {
+ val name = warehouse.getOrElse(defaultWarehouse)
+ synchronized {
+ catalogs.getOrElseUpdate(cacheKey(name), createCatalog(name))
}
}
+ private def createCatalog(warehouse: String): Catalog =
+ StorageConfig.icebergCatalogType match {
+ case "hadoop" =>
+ IcebergUtil.createHadoopCatalog(
+ "texera_iceberg",
+ StorageConfig.fileStorageDirectoryPath
+ )
+ case "rest" =>
+ IcebergUtil.createRestCatalog(
+ "texera_iceberg",
+ warehouse
+ )
+ case "postgres" =>
+ IcebergUtil.createPostgresCatalog(
+ "texera_iceberg",
+ StorageConfig.fileStorageDirectoryPath
+ )
+ case unsupported =>
+ throw new IllegalArgumentException(s"Unsupported catalog type: $unsupported")
+ }
+
/**
- * Replaces the existing Iceberg catalog instance.
- * - This method is useful for testing or dynamically updating the catalog.
+ * Replaces the cached catalog for a warehouse, primarily for testing or reconfiguration.
*
- * @param catalog the new Iceberg catalog instance to replace the current one.
+ * @param catalog the catalog to cache.
+ * @param warehouse the warehouse to cache it under; `None` uses the configured default.
*/
- def replaceInstance(catalog: Catalog): Unit = {
- instance = Some(catalog)
- }
+ def replaceInstance(catalog: Catalog, warehouse: Option[String] = None): Unit =
+ synchronized {
+ catalogs(cacheKey(warehouse.getOrElse(defaultWarehouse))) = catalog
+ }
}
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/VFSURIFactory.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/VFSURIFactory.scala
index 5124c745aa..79e36e77f4 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/VFSURIFactory.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/VFSURIFactory.scala
@@ -45,12 +45,61 @@ case class VFSUriComponents(
workflowId: WorkflowIdentity,
executionId: ExecutionIdentity,
globalPortId: Option[GlobalPortIdentity],
- resourceType: VFSResourceType.Value
+ resourceType: VFSResourceType.Value,
+ // 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: Option[String] = None
)
object VFSURIFactory {
val VFS_FILE_URI_SCHEME = "vfs"
+ // A warehouse name becomes a URI path segment, so it is restricted to characters
+ // that carry no meaning there: no '/' to add segments, and no '%' to smuggle one
+ // in percent-encoded form. Registration applies a stricter rule still; this is the
+ // invariant the URI layer itself depends on.
+ private val warehouseNamePattern = "[A-Za-z0-9][A-Za-z0-9_-]*".r
+
+ private def isValidWarehouseName(name: String): Boolean =
+ warehouseNamePattern.pattern.matcher(name).matches()
+
+ // Warehouse is carried as a leading `/wh/<name>` path segment so a storage URI
+ // fully identifies which warehouse its table lives in. Absent for non-BYO storage.
+ private def warehousePathSegment(warehouse: Option[String]): String =
+ warehouse
+ .map { name =>
+ require(
+ isValidWarehouseName(name),
+ s"warehouse name must match ${warehouseNamePattern.regex} " +
+ s"(it becomes a URI path segment): $name"
+ )
+ s"/wh/$name"
+ }
+ .getOrElse("")
+
+ /**
+ * The warehouse encoded in a VFS URI, if present. Reported as part of
+ * [[VFSUriComponents]] by [[decodeURI]], which is the only way in.
+ *
+ * Anchored to the leading segment on purpose. The warehouse is written as a
+ * leading `/wh/<name>` prefix and `DocumentFactory` strips only a leading one,
+ * so scanning the whole path would disagree with the stripper: a later segment
+ * that happens to be `wh` -- e.g. inside a user-chosen operator id in a
+ * console-messages URI -- would select a warehouse the URI was never built for.
+ *
+ * The segments come from the RAW path (see `decodeURI`), so a percent-encoded
+ * slash stays inside the segment that contains it instead of becoming a
+ * separator. The name is then required to be a legal warehouse name, so anything
+ * that could not have been written by `warehousePathSegment` resolves to no
+ * warehouse rather than to a wrong one.
+ */
+ private def warehouseFrom(segments: List[String]): Option[String] =
+ segments match {
+ case "wh" :: name :: _ if isValidWarehouseName(name) => Some(name)
+ case _ => None
+ }
+
/**
* Parses a VFS URI and extracts its components
*
@@ -63,7 +112,11 @@ object VFSURIFactory {
throw new IllegalArgumentException(s"Invalid URI scheme: ${uri.getScheme}")
}
- val segments = uri.getPath.stripPrefix("/").split("/").toList
+ // Raw path, for the same reason as warehouseFromURI: keys are located by
+ // searching the segments, so a percent-encoded slash inside a segment must not
+ // split it and shift which `wid`/`eid` the search finds. Python's decode_uri
+ // splits the raw path too, so both languages read a URI identically.
+ val segments = uri.getRawPath.stripPrefix("/").split("/").toList
def extractValue(key: String): String = {
val index = segments.indexOf(key)
@@ -86,7 +139,13 @@ object VFSURIFactory {
.find(_.toString.toLowerCase == resourceTypeStr)
.getOrElse(throw new IllegalArgumentException(s"Unknown resource type: $resourceTypeStr"))
- VFSUriComponents(workflowId, executionId, globalPortIdOption, resourceType)
+ VFSUriComponents(
+ workflowId,
+ executionId,
+ globalPortIdOption,
+ resourceType,
+ warehouseFrom(segments)
+ )
}
/**
@@ -96,10 +155,11 @@ object VFSURIFactory {
def createPortBaseURI(
workflowId: WorkflowIdentity,
executionId: ExecutionIdentity,
- globalPortId: GlobalPortIdentity
+ globalPortId: GlobalPortIdentity,
+ warehouse: Option[String] = None
): URI =
new URI(
- s"$VFS_FILE_URI_SCHEME:///wid/${workflowId.id}/eid/${executionId.id}" +
+ s"$VFS_FILE_URI_SCHEME://${warehousePathSegment(warehouse)}/wid/${workflowId.id}/eid/${executionId.id}" +
s"/globalportid/${globalPortId.serializeAsString}"
)
@@ -115,12 +175,14 @@ object VFSURIFactory {
*/
def createRuntimeStatisticsURI(
workflowId: WorkflowIdentity,
- executionId: ExecutionIdentity
+ executionId: ExecutionIdentity,
+ warehouse: Option[String] = None
): URI = {
createNonResultVFSURI(
VFSResourceType.RUNTIME_STATISTICS,
workflowId,
- executionId
+ executionId,
+ warehouse = warehouse
)
}
@@ -130,13 +192,15 @@ object VFSURIFactory {
def createConsoleMessagesURI(
workflowId: WorkflowIdentity,
executionId: ExecutionIdentity,
- operatorId: OperatorIdentity
+ operatorId: OperatorIdentity,
+ warehouse: Option[String] = None
): URI = {
createNonResultVFSURI(
VFSResourceType.CONSOLE_MESSAGES,
workflowId,
executionId,
- Some(operatorId)
+ Some(operatorId),
+ warehouse
)
}
@@ -156,7 +220,8 @@ object VFSURIFactory {
resourceType: VFSResourceType.Value,
workflowId: WorkflowIdentity,
executionId: ExecutionIdentity,
- operatorId: Option[OperatorIdentity] = None
+ operatorId: Option[OperatorIdentity] = None,
+ warehouse: Option[String] = None
): URI = {
if (resourceType == VFSResourceType.RESULT) {
@@ -177,10 +242,22 @@ object VFSURIFactory {
)
}
+ // The operator id is user-supplied (it comes straight off the workflow JSON) and
+ // is interpolated into the URI path below. A '/' in it would add path segments,
+ // letting it forge structure the URI never meant to have -- e.g. a `wh/<name>`
+ // pair that would then be read back as a warehouse.
+ operatorId.foreach { opId =>
+ require(
+ !opId.id.contains('/'),
+ s"operatorId must not contain '/' (VFS URI parsing relies on this): ${opId.id}"
+ )
+ }
+
+ val whSegment = warehousePathSegment(warehouse)
val baseUri = operatorId match {
case Some(opId) =>
- s"$VFS_FILE_URI_SCHEME:///wid/${workflowId.id}/eid/${executionId.id}/opid/${opId.id}"
- case None => s"$VFS_FILE_URI_SCHEME:///wid/${workflowId.id}/eid/${executionId.id}"
+ s"$VFS_FILE_URI_SCHEME://$whSegment/wid/${workflowId.id}/eid/${executionId.id}/opid/${opId.id}"
+ case None => s"$VFS_FILE_URI_SCHEME://$whSegment/wid/${workflowId.id}/eid/${executionId.id}"
}
new URI(s"$baseUri/${resourceType.toString.toLowerCase}")
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala
index 182da2baac..cc414825d9 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/result/iceberg/IcebergDocument.scala
@@ -57,6 +57,7 @@ object Constants {
* @param tableSchema schema of the table.
* @param serde function to serialize T into an Iceberg Record.
* @param deserde function to deserialize an Iceberg Record into T.
+ * @param warehouse the warehouse whose catalog backs this table; `None` uses the configured default.
* @tparam T type of the data items stored in the Iceberg table.
*/
private[storage] class IcebergDocument[T >: Null <: AnyRef](
@@ -64,13 +65,14 @@ private[storage] class IcebergDocument[T >: Null <: AnyRef](
val tableName: String,
val tableSchema: org.apache.iceberg.Schema,
val serde: (org.apache.iceberg.Schema, T) => Record,
- val deserde: (org.apache.iceberg.Schema, Record) => T
+ val deserde: (org.apache.iceberg.Schema, Record) => T,
+ val warehouse: Option[String] = None
) extends VirtualDocument[T]
with OnIceberg {
private val lock = new ReentrantReadWriteLock()
- @transient lazy val catalog: Catalog = IcebergCatalogInstance.getInstance()
+ @transient lazy val catalog: Catalog = IcebergCatalogInstance.getInstance(warehouse)
/**
* Returns the URI of the table location.
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/WorkflowContext.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/WorkflowContext.scala
index df6944ed88..5d12daf18a 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/WorkflowContext.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/WorkflowContext.scala
@@ -35,5 +35,6 @@ class WorkflowContext(
var workflowId: WorkflowIdentity = DEFAULT_WORKFLOW_ID,
var executionId: ExecutionIdentity = DEFAULT_EXECUTION_ID,
var workflowSettings: WorkflowSettings = DEFAULT_WORKFLOW_SETTINGS,
- var cuid: Option[Int] = None
+ var cuid: Option[Int] = None,
+ var warehouse: Option[String] = None
)
diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/serde/GlobalPortIdentitySerde.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/serde/GlobalPortIdentitySerde.scala
index c8fd8e1a36..65a836bf7a 100644
--- a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/serde/GlobalPortIdentitySerde.scala
+++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/serde/GlobalPortIdentitySerde.scala
@@ -53,6 +53,17 @@ object GlobalPortIdentitySerde {
!layerName.contains('_'),
s"layerName must not contain '_' (VFS URI parsing relies on this): $layerName"
)
+ // A '/' would add path segments to the VFS URI these ids are interpolated
+ // into, letting a user-chosen id forge structure the URI never meant to have
+ // (e.g. a `wh/<name>` pair). Rejected for the same reason as '_'.
+ require(
+ !logicalOpId.contains('/'),
+ s"logicalOpId must not contain '/' (VFS URI parsing relies on this): $logicalOpId"
+ )
+ require(
+ !layerName.contains('/'),
+ s"layerName must not contain '/' (VFS URI parsing relies on this): $layerName"
+ )
require(portId >= 0, s"portId must be non-negative: $portId")
s"(logicalOpId=$logicalOpId,layerName=$layerName,portId=$portId,isInternal=$isInternal,isInput=$isInput)"
}
diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/DocumentFactorySpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/DocumentFactorySpec.scala
index ad8b5580c5..ecb5f5c54f 100644
--- a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/DocumentFactorySpec.scala
+++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/DocumentFactorySpec.scala
@@ -61,7 +61,11 @@ class DocumentFactorySpec extends AnyFlatSpec with Matchers with BeforeAndAfterA
override def beforeAll(): Unit = {
super.beforeAll()
- LocalHadoopIcebergCatalog.ensure()
+ // "wh-test" is registered explicitly so the warehouse-scoped case resolves to
+ // this local catalog instead of reaching for a live REST catalog under the
+ // default (`rest`) config -- otherwise the case would pass or fail according to
+ // the machine's catalog type rather than the code under test.
+ LocalHadoopIcebergCatalog.ensure("wh-test")
}
// ---------------------------------------------------------------------------
@@ -235,6 +239,27 @@ class DocumentFactorySpec extends AnyFlatSpec with Matchers with BeforeAndAfterA
DocumentFactory.documentExists(stateUri) shouldBe true
}
+ it should "create + find a document for a warehouse-scoped vfs URI (the /wh/ segment is stripped from the storage key)" in {
+ val base = VFSURIFactory.createPortBaseURI(
+ WorkflowIdentity(0),
+ ExecutionIdentity(0),
+ GlobalPortIdentity(
+ PhysicalOpIdentity(
+ logicalOpId = OperatorIdentity(s"op-${UUID.randomUUID().toString.replace("-", "")}"),
+ layerName = "main"
+ ),
+ PortIdentity()
+ ),
+ warehouse = Some("wh-test")
+ )
+ val vfsUri = VFSURIFactory.resultURI(base)
+ vfsUri.getPath should startWith("/wh/wh-test/")
+
+ val doc = DocumentFactory.createDocument(vfsUri, vfsSchema)
+ doc shouldBe an[IcebergDocument[_]]
+ DocumentFactory.documentExists(vfsUri) shouldBe true
+ }
+
"documentExists" should "report false before creation and true after for a vfs URI" in {
val vfsUri = freshResultURI()
DocumentFactory.documentExists(vfsUri) shouldBe false
diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/LocalHadoopIcebergCatalog.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/LocalHadoopIcebergCatalog.scala
index 3b92b763a3..df3f0367d6 100644
--- a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/LocalHadoopIcebergCatalog.scala
+++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/LocalHadoopIcebergCatalog.scala
@@ -19,6 +19,7 @@
package org.apache.texera.amber.core.storage
+import org.apache.iceberg.catalog.Catalog
import org.apache.texera.amber.util.IcebergUtil
import java.nio.file.{Files, Path}
@@ -42,11 +43,21 @@ import java.util.Comparator
*/
object LocalHadoopIcebergCatalog {
- private var initialized = false
+ private var catalog: Option[Catalog] = None
- def ensure(): Unit =
+ /**
+ * Installs the shared local catalog, additionally registering it under each of
+ * `warehouses`.
+ *
+ * A suite that exercises a warehouse-scoped URI must name that warehouse here:
+ * under the configured `rest` catalog type the cache is keyed by warehouse name,
+ * so an unregistered name would miss and try to build a *live* REST catalog,
+ * making the suite pass or fail according to the machine's catalog config rather
+ * than the code under test.
+ */
+ def ensure(warehouses: String*): Unit =
synchronized {
- if (!initialized) {
+ val installed = catalog.getOrElse {
val warehouse = Files.createTempDirectory("wfcore-iceberg-shared")
// Best-effort recursive cleanup of the temp warehouse on JVM exit so test runs
// don't leave wfcore-iceberg-shared* directories behind on dev machines / CI.
@@ -57,10 +68,11 @@ object LocalHadoopIcebergCatalog {
.forEach((p: Path) => Files.deleteIfExists(p))
catch { case _: Throwable => () }
}
- IcebergCatalogInstance.replaceInstance(
- IcebergUtil.createHadoopCatalog("wfcore-test", warehouse)
- )
- initialized = true
+ val created = IcebergUtil.createHadoopCatalog("wfcore-test", warehouse)
+ catalog = Some(created)
+ IcebergCatalogInstance.replaceInstance(created)
+ created
}
+ warehouses.foreach(name => IcebergCatalogInstance.replaceInstance(installed, Some(name)))
}
}
diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/VFSURIFactorySpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/VFSURIFactorySpec.scala
index aa3eb0dfdb..f163a0658e 100644
--- a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/VFSURIFactorySpec.scala
+++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/storage/VFSURIFactorySpec.scala
@@ -60,7 +60,7 @@ class VFSURIFactorySpec extends AnyFlatSpec {
assert(resultURI.getPath.endsWith("/result"))
assert(stateURI.getPath.endsWith("/state"))
- val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType) =
+ val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType, _) =
VFSURIFactory.decodeURI(resultURI)
assert(wid == workflowId)
assert(eid == executionId)
@@ -75,7 +75,7 @@ class VFSURIFactorySpec extends AnyFlatSpec {
assert(path.endsWith("/runtimestatistics"))
assert(!path.contains("/opid/"))
- val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType) = VFSURIFactory.decodeURI(uri)
+ val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType, _) = VFSURIFactory.decodeURI(uri)
assert(wid == workflowId)
assert(eid == executionId)
assert(globalPortIdOpt.isEmpty)
@@ -90,7 +90,7 @@ class VFSURIFactorySpec extends AnyFlatSpec {
// The current `decodeURI` does not extract the operator id (it has no
// "opid" branch), so we only round-trip wid/eid/resourceType here.
- val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType) = VFSURIFactory.decodeURI(uri)
+ val VFSUriComponents(wid, eid, globalPortIdOpt, resourceType, _) = VFSURIFactory.decodeURI(uri)
assert(wid == workflowId)
assert(eid == executionId)
assert(globalPortIdOpt.isEmpty)
@@ -121,4 +121,114 @@ class VFSURIFactorySpec extends AnyFlatSpec {
VFSURIFactory.decodeURI(new URI("vfs:///eid/2/wid"))
}
}
+
+ "decodeURI" should "report the warehouse from a leading /wh/<name> segment" in {
+ val uri =
+ VFSURIFactory.createPortBaseURI(
+ workflowId,
+ executionId,
+ portId,
+ warehouse = Some("user-2-foo")
+ )
+ assert(uri.getPath.startsWith("/wh/user-2-foo/"))
+ // A base URI has no resource segment, so decode the derived result URI.
+ assert(
+ VFSURIFactory.decodeURI(VFSURIFactory.resultURI(uri)).warehouse.contains("user-2-foo")
+ )
+ }
+
+ it should "return None when no /wh/ segment is present (non-BYO URIs are unchanged)" in {
+ val uri = VFSURIFactory.createPortBaseURI(workflowId, executionId, portId)
+ assert(!uri.getPath.contains("/wh/"))
+ assert(VFSURIFactory.decodeURI(VFSURIFactory.resultURI(uri)).warehouse.isEmpty)
+ }
+
+ it should "only honour a LEADING wh segment, never one deeper in the path" in {
+ // A `wh` appearing later -- e.g. inside an operator id -- must not select a
+ // warehouse: it would disagree with DocumentFactory, which strips only a
+ // leading `wh/<name>/`, and would route the write to another user's warehouse.
+ assert(
+ VFSURIFactory
+ .decodeURI(new URI("vfs:///wid/1/eid/2/opid/a/wh/victim/b/consolemessages"))
+ .warehouse
+ .isEmpty
+ )
+ // An operator literally named `wh` is likewise not a warehouse.
+ assert(
+ VFSURIFactory
+ .decodeURI(new URI("vfs:///wid/1/eid/2/opid/wh/consolemessages"))
+ .warehouse
+ .isEmpty
+ )
+ }
+
+ it should "not let a percent-encoded slash in the name forge extra segments" in {
+ // Decoded before splitting, this path would read as /wh/a/wid/999/... -- handing
+ // back "a" as the warehouse and shifting which `wid` the parser sees. Splitting
+ // the raw path keeps `%2F` inside its own segment, and the name is then rejected
+ // as illegal, so the URI resolves to no warehouse rather than to a wrong one.
+ assert(
+ VFSURIFactory
+ .decodeURI(new URI("vfs:///wh/a%2Fwid%2F999/wid/1/eid/2/result"))
+ .warehouse
+ .isEmpty
+ )
+ assert(
+ VFSURIFactory
+ .decodeURI(new URI("vfs:///wh/user-2%2Dfoo/wid/7/eid/3/result"))
+ .warehouse
+ .isEmpty
+ )
+ }
+
+ it should "locate wid/eid by raw segment, so an encoded slash cannot shift them" in {
+ // Decoded before splitting, this path reads as /wh/a/wid/999/wid/1/... and the
+ // key search finds the injected `wid` first -- resolving to execution 999 and
+ // landing this execution's data under another's storage key. Python's
+ // decode_uri splits the raw path, so decoding here would also make the two
+ // languages disagree about which execution a URI belongs to.
+ val components =
+ VFSURIFactory.decodeURI(new URI("vfs:///wh/a%2Fwid%2F999/wid/1/eid/2/result"))
+ assert(components.workflowId == WorkflowIdentity(1))
+ assert(components.executionId == ExecutionIdentity(2))
+ }
+
+ "VFSURIFactory" should "reject an operatorId containing '/' rather than let it forge URI segments" in {
+ assertThrows[IllegalArgumentException] {
+ VFSURIFactory.createConsoleMessagesURI(
+ workflowId,
+ executionId,
+ OperatorIdentity("a/wh/victim/b")
+ )
+ }
+ }
+
+ it should "reject a warehouse name that is not safe as a URI path segment" in {
+ Seq("a/b", "a%2Fb", "", "-lead", "sp ace").foreach { bad =>
+ withClue(s"warehouse name '$bad' should be rejected: ") {
+ assertThrows[IllegalArgumentException] {
+ VFSURIFactory.createPortBaseURI(workflowId, executionId, portId, Some(bad))
+ }
+ }
+ }
+ }
+
+ "A warehouse-scoped URI" should
+ "still round-trip through decodeURI (wid/eid/port/resource resolved despite the /wh/ prefix)" in {
+ val base =
+ VFSURIFactory.createPortBaseURI(
+ workflowId,
+ executionId,
+ portId,
+ warehouse = Some("user-2-foo")
+ )
+ val resultURI = VFSURIFactory.resultURI(base)
+ assert(VFSURIFactory.decodeURI(resultURI).warehouse.contains("user-2-foo"))
+
+ val components = VFSURIFactory.decodeURI(resultURI)
+ assert(components.workflowId == workflowId)
+ assert(components.executionId == executionId)
+ assert(components.globalPortId.contains(portId))
+ assert(components.resourceType == VFSResourceType.RESULT)
+ }
}