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-5265f207065dd06f6fd70467402c7843d9379e02 in repository https://gitbox.apache.org/repos/asf/texera.git
commit c0c3ba739ce264215be03461a10417d01465c4e1 Author: Meng Wang <[email protected]> AuthorDate: Thu Aug 6 18:00:30 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) <[email protected]> --- .../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 + + [email protected](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) + } }
