From a28aa88d727ca0476db490859e6573d78bf89cc3 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 17 Oct 2026 09:19:31 +0800 Subject: [PATCH 1/5] [core] Support REST catalog branch management --- bindings/python/project-description.md | 17 + .../python/pypaimon_rust/datafusion.pyi | 9 + bindings/python/src/context.rs | 108 ++++-- .../python/tests/test_catalog_branches.py | 154 +++++++++ crates/paimon/src/api/api_request.rs | 12 + crates/paimon/src/api/api_response.rs | 6 + crates/paimon/src/api/mod.rs | 15 +- crates/paimon/src/api/resource_paths.rs | 27 ++ crates/paimon/src/api/rest_api.rs | 97 +++++- crates/paimon/src/catalog/mod.rs | 37 ++ .../paimon/src/catalog/rest/rest_catalog.rs | 84 +++++ crates/paimon/src/error.rs | 10 + crates/paimon/tests/rest_branch_test.rs | 327 ++++++++++++++++++ 13 files changed, 867 insertions(+), 36 deletions(-) create mode 100644 bindings/python/tests/test_catalog_branches.py create mode 100644 crates/paimon/tests/rest_branch_test.rs diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md index 0db99785f..44633f866 100644 --- a/bindings/python/project-description.md +++ b/bindings/python/project-description.md @@ -148,4 +148,21 @@ ctx = SessionContext() ctx.register_catalog_provider("paimon", catalog) ``` +REST catalogs also support branch management through Rust core: + +```python +identifier = ("db", "table") +catalog.create_branch(identifier, "empty-branch") +catalog.create_branch(identifier, "release-branch", from_tag="release-tag") +branches = catalog.list_branches(identifier) +release = catalog.get_table(("db", "table$branch_release-branch")) +catalog.fast_forward(identifier, "release-branch") +catalog.drop_branch(identifier, "empty-branch") +``` + +The server validates branch names and publishes metadata changes. Fast-forward +updates the identifier's current branch (main by default); it does not change +the source branch. These catalog methods currently require a REST catalog. +Other catalog implementations raise `NotImplementedError`. + Time travel queries are not supported in the Python binding at this time. diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 47fa04660..cf0525721 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -755,6 +755,15 @@ class PaimonCatalog: def list_databases(self) -> List[str]: ... def list_tables(self, database_name: str) -> List[str]: ... def get_table(self, identifier: Union[str, Tuple[str, str]]) -> Table: ... + def create_branch(self, identifier: Union[str, Tuple[str, str]], branch_name: str, + from_tag: Optional[str] = None) -> None: + """Create a REST branch, empty unless from_tag is supplied.""" + ... + def drop_branch(self, identifier: Union[str, Tuple[str, str]], branch_name: str) -> None: ... + def fast_forward(self, identifier: Union[str, Tuple[str, str]], branch_name: str) -> None: + """Replace the identifier's current branch metadata through the REST server.""" + ... + def list_branches(self, identifier: Union[str, Tuple[str, str]]) -> List[str]: ... class PythonScalarUDF: def __init__( diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 8177a75b0..0954c5022 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -278,6 +278,36 @@ pub struct PaimonCatalog { provider: Arc, } +fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult { + let (database, object_name) = if let Ok(name) = identifier.extract::() { + let (database, table) = name.split_once('.').ok_or_else(|| { + PyValueError::new_err(format!( + "expected identifier in 'db.table' format, got '{name}'" + )) + })?; + (database.to_string(), table.to_string()) + } else { + identifier.extract::<(String, String)>()? + }; + if database.is_empty() || object_name.is_empty() { + return Err(PyValueError::new_err( + "database and table names must not be empty", + )); + } + let identifier = Identifier::new(&database, &object_name); + if identifier + .parsed_object_name() + .map_err(to_py_err)? + .system_table() + .is_some() + { + return Err(PyValueError::new_err( + "Catalog table operations do not support system-table identifiers", + )); + } + Ok(identifier) +} + #[pymethods] impl PaimonCatalog { /// Create a Paimon catalog that can be registered into a DataFusion session. @@ -334,29 +364,9 @@ impl PaimonCatalog { /// Get a table by "db.table" or a (database, table) tuple. /// Tuple components preserve dots; the table can include a branch suffix. fn get_table(&self, py: Python<'_>, identifier: &Bound<'_, PyAny>) -> PyResult { - let (database, object_name) = if let Ok(name) = identifier.extract::() { - let (database, table) = name.split_once('.').ok_or_else(|| { - PyValueError::new_err(format!( - "expected identifier in 'db.table' format, got '{name}'" - )) - })?; - (database.to_string(), table.to_string()) - } else { - identifier.extract::<(String, String)>()? - }; - if database.is_empty() || object_name.is_empty() { - return Err(PyValueError::new_err( - "database and table names must not be empty", - )); - } - let id = Identifier::new(&database, &object_name); + let id = catalog_table_identifier(identifier)?; let parsed = id.parsed_object_name().map_err(to_py_err)?; - if parsed.system_table().is_some() { - return Err(PyValueError::new_err( - "get_table() does not support system-table identifiers", - )); - } - let base_id = Identifier::new(&database, parsed.table()); + let base_id = Identifier::new(id.database(), parsed.table()); let branch = parsed.branch().map(str::to_string); let catalog = Arc::clone(&self.catalog); let table = py.detach(|| { @@ -372,6 +382,60 @@ impl PaimonCatalog { })?; Ok(PyTable::new(Arc::new(table))) } + + /// Create an empty branch or one initialized from a tag through the REST catalog. + #[pyo3(signature = (identifier, branch_name, from_tag=None))] + fn create_branch( + &self, + py: Python<'_>, + identifier: &Bound<'_, PyAny>, + branch_name: &str, + from_tag: Option<&str>, + ) -> PyResult<()> { + let identifier = catalog_table_identifier(identifier)?; + py.detach(|| { + runtime().block_on( + self.catalog + .create_branch(&identifier, branch_name, from_tag), + ) + }) + .map_err(to_py_err) + } + + /// Delete branch metadata through the REST catalog. + fn drop_branch( + &self, + py: Python<'_>, + identifier: &Bound<'_, PyAny>, + branch_name: &str, + ) -> PyResult<()> { + let identifier = catalog_table_identifier(identifier)?; + py.detach(|| runtime().block_on(self.catalog.drop_branch(&identifier, branch_name))) + .map_err(to_py_err) + } + + /// Fast-forward the identifier's current branch through the REST catalog. + fn fast_forward( + &self, + py: Python<'_>, + identifier: &Bound<'_, PyAny>, + branch_name: &str, + ) -> PyResult<()> { + let identifier = catalog_table_identifier(identifier)?; + py.detach(|| runtime().block_on(self.catalog.fast_forward(&identifier, branch_name))) + .map_err(to_py_err) + } + + /// List branch names in REST server order. + fn list_branches( + &self, + py: Python<'_>, + identifier: &Bound<'_, PyAny>, + ) -> PyResult> { + let identifier = catalog_table_identifier(identifier)?; + py.detach(|| runtime().block_on(self.catalog.list_branches(&identifier))) + .map_err(to_py_err) + } } #[pyclass(name = "SQLOptions", module = "pypaimon_rust.datafusion")] diff --git a/bindings/python/tests/test_catalog_branches.py b/bindings/python/tests/test_catalog_branches.py new file mode 100644 index 000000000..e58751543 --- /dev/null +++ b/bindings/python/tests/test_catalog_branches.py @@ -0,0 +1,154 @@ +# 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. + +"""Branch operations use core REST APIs and release the GIL during requests.""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.parse import unquote_plus + +import pytest + +from pypaimon_rust.datafusion import PaimonCatalog + + +@pytest.fixture +def catalog(monkeypatch): + requests = [] + monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1") + monkeypatch.setenv("no_proxy", "localhost,127.0.0.1") + + class Handler(BaseHTTPRequestHandler): + def handle_operation(self): + assert self.headers["Authorization"] == "Bearer test-token" + size = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(size)) if size else None + requests.append((self.command, self.path, body)) + parts = [unquote_plus(part) for part in self.path.split("/")] + if self.path.startswith("/v1/config?"): + payload, status = {"defaults": {"prefix": "test"}}, 200 + elif "missing" in parts: + payload, status = {"resourceType": "TABLE", "message": "missing", "code": 404}, 404 + elif "forbidden" in parts: + payload, status = {"message": "denied", "code": 403}, 403 + elif "conflict" in parts: + payload, status = {"message": "exists", "code": 409}, 409 + elif self.command == "GET": + payload, status = {"branches": ["dev", "release"]}, 200 + else: + payload, status = None, 204 + output = json.dumps(payload).encode() if status != 204 else b"" + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(output))) + self.end_headers() + self.wfile.write(output) + + do_GET = handle_operation + do_POST = handle_operation + do_DELETE = handle_operation + + def log_message(self, *args): + pass + + server = HTTPServer(("localhost", 0), Handler) + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + native = PaimonCatalog({ + "metastore": "rest", "uri": "http://localhost:%d" % server.server_port, + "warehouse": "warehouse", "token.provider": "bear", "token": "test-token", + }) + requests.clear() + yield native, requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +def test_rest_branch_operations_use_core_protocol(catalog): + native, requests = catalog + identifier = ("db.name", "table.name$branch_dev") + native.create_branch(identifier, "empty") + native.create_branch(identifier, "tagged", from_tag="tag 一") + assert native.list_branches(identifier) == ["dev", "release"] + native.fast_forward(identifier, "tagged") + native.drop_branch(identifier, "empty") + base = "/v1/test/databases/db.name/tables/table.name%24branch_dev/branches" + assert requests == [ + ("POST", base, {"branch": "empty", "fromTag": None}), + ("POST", base, {"branch": "tagged", "fromTag": "tag 一"}), + ("GET", base, None), + ("POST", base + "/tagged/forward", {}), + ("DELETE", base + "/empty", None), + ] + + +def test_string_identifier_and_branch_path_encoding(catalog): + native, requests = catalog + native.drop_branch("db.table", "name /?#一") + assert requests == [( + "DELETE", "/v1/test/databases/db/tables/table/branches/name+%2F%3F%23%E4%B8%80", None, + )] + + +@pytest.mark.parametrize("operation", ["drop_branch", "fast_forward"]) +@pytest.mark.parametrize("name", ["", " ", ".", ".."]) +def test_branch_paths_do_not_mutate_parent_resources(catalog, operation, name): + native, requests = catalog + with pytest.raises(ValueError): + getattr(native, operation)("db.table", name) + assert requests == [] + + +@pytest.mark.parametrize("operation", ["create_branch", "drop_branch", "fast_forward", "list_branches"]) +@pytest.mark.parametrize("identifier", ["db", ".table", "db.", ("", "t"), ("db", ""), "db.table$snapshots"]) +def test_invalid_identifiers_fail_before_sending_requests(catalog, operation, identifier): + native, requests = catalog + args = () if operation == "list_branches" else ("dev",) + with pytest.raises(ValueError): + getattr(native, operation)(identifier, *args) + assert requests == [] + + +@pytest.mark.parametrize("operation", ["create_branch", "drop_branch", "fast_forward", "list_branches"]) +def test_rest_branch_errors_are_not_retried(catalog, operation): + native, requests = catalog + args = () if operation == "list_branches" else ("dev",) + with pytest.raises(ValueError, match="does not exist"): + getattr(native, operation)("db.missing", *args) + assert len(requests) == 1 + with pytest.raises(ValueError, match="Forbidden"): + getattr(native, operation)("db.forbidden", *args) + assert len(requests) == 2 + + +def test_duplicate_branch_is_reported(catalog): + native, requests = catalog + with pytest.raises(ValueError, match="Branch dev of table db.conflict already exists"): + native.create_branch("db.conflict", "dev") + assert len(requests) == 1 + + +@pytest.mark.parametrize("operation", ["create_branch", "drop_branch", "fast_forward", "list_branches"]) +def test_catalog_without_branch_management_reports_unsupported(tmp_path, operation): + native = PaimonCatalog({"warehouse": str(tmp_path)}) + args = () if operation == "list_branches" else ("dev",) + with pytest.raises(NotImplementedError, match="branch management"): + getattr(native, operation)("db.table", *args) diff --git a/crates/paimon/src/api/api_request.rs b/crates/paimon/src/api/api_request.rs index 679ddeb8c..b46300c31 100644 --- a/crates/paimon/src/api/api_request.rs +++ b/crates/paimon/src/api/api_request.rs @@ -75,6 +75,18 @@ pub struct RenameTableRequest { pub destination: Identifier, } +/// Request to create a table branch, optionally from a tag. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateBranchRequest { + pub branch: String, + pub from_tag: Option, +} + +/// Empty JSON object sent by Java's fast-forward operation. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ForwardBranchRequest {} + /// Request to create a table tag. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/crates/paimon/src/api/api_response.rs b/crates/paimon/src/api/api_response.rs index 0a595f021..aa489f622 100644 --- a/crates/paimon/src/api/api_response.rs +++ b/crates/paimon/src/api/api_response.rs @@ -114,6 +114,12 @@ impl AuditRESTResponse { } } +/// Branch names returned by a REST catalog, in server order. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ListBranchesResponse { + pub branches: Option>, +} + /// Latest snapshot and table statistics returned by a REST catalog. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/crates/paimon/src/api/mod.rs b/crates/paimon/src/api/mod.rs index cffb8865c..55527ae01 100644 --- a/crates/paimon/src/api/mod.rs +++ b/crates/paimon/src/api/mod.rs @@ -33,19 +33,20 @@ mod user_agent; // Re-export request types pub use api_request::{ - AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateDatabaseRequest, - CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, CreateTagRequest, - CreateViewRequest, DropPartitionsRequest, DropPolicyRequest, ListPartitionsByFilterRequest, - ListPartitionsByNamesRequest, PolicyRequest, RenameTableRequest, RevokePermissionRequest, + AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateBranchRequest, + CreateDatabaseRequest, CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, + CreateTagRequest, CreateViewRequest, DropPartitionsRequest, DropPolicyRequest, + ForwardBranchRequest, ListPartitionsByFilterRequest, ListPartitionsByNamesRequest, + PolicyRequest, RenameTableRequest, RevokePermissionRequest, }; // Re-export response types pub use api_response::{ AuditRESTResponse, AuthTableQueryResponse, ConfigResponse, ErrorResponse, GetDatabaseResponse, GetFunctionResponse, GetTableResponse, GetTableSnapshotResponse, GetTableTokenResponse, - GetTagResponse, GetViewResponse, ListDatabasesResponse, ListFunctionsResponse, - ListPartitionsResponse, ListPermissionsResponse, ListPoliciesResponse, ListTablesResponse, - ListViewsResponse, PagedList, TableSnapshot, + GetTagResponse, GetViewResponse, ListBranchesResponse, ListDatabasesResponse, + ListFunctionsResponse, ListPartitionsResponse, ListPermissionsResponse, ListPoliciesResponse, + ListTablesResponse, ListViewsResponse, PagedList, TableSnapshot, }; // Re-export management types diff --git a/crates/paimon/src/api/resource_paths.rs b/crates/paimon/src/api/resource_paths.rs index 62f1bcf40..14e973226 100644 --- a/crates/paimon/src/api/resource_paths.rs +++ b/crates/paimon/src/api/resource_paths.rs @@ -121,6 +121,33 @@ impl ResourcePaths { format!("{}/snapshot", self.table(database_name, table_name)) } + /// Get the branch collection for a table. + pub fn branches(&self, database_name: &str, table_name: &str) -> String { + format!("{}/branches", self.table(database_name, table_name)) + } + + /// Get one branch, encoding its name as a single path component. + pub fn branch(&self, database_name: &str, table_name: &str, branch_name: &str) -> String { + format!( + "{}/{}", + self.branches(database_name, table_name), + RESTUtil::encode_string(branch_name) + ) + } + + /// Get the fast-forward endpoint for a branch. + pub fn forward_branch( + &self, + database_name: &str, + table_name: &str, + branch_name: &str, + ) -> String { + format!( + "{}/forward", + self.branch(database_name, table_name, branch_name) + ) + } + /// Get the views endpoint path for a database. pub fn views(&self, database_name: &str) -> String { format!( diff --git a/crates/paimon/src/api/rest_api.rs b/crates/paimon/src/api/rest_api.rs index 57828e34d..a69018b22 100644 --- a/crates/paimon/src/api/rest_api.rs +++ b/crates/paimon/src/api/rest_api.rs @@ -29,16 +29,17 @@ use crate::spec::{Partition, PartitionStatistics, Schema, SchemaChange, Snapshot use crate::Result; use super::api_request::{ - AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateDatabaseRequest, - CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, CreateTagRequest, - CreateViewRequest, DropPartitionsRequest, DropPolicyRequest, ListPartitionsByFilterRequest, - ListPartitionsByNamesRequest, PolicyRequest, RenameTableRequest, RevokePermissionRequest, + AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateBranchRequest, + CreateDatabaseRequest, CreateFunctionRequest, CreatePartitionsRequest, CreateTableRequest, + CreateTagRequest, CreateViewRequest, DropPartitionsRequest, DropPolicyRequest, + ForwardBranchRequest, ListPartitionsByFilterRequest, ListPartitionsByNamesRequest, + PolicyRequest, RenameTableRequest, RevokePermissionRequest, }; use super::api_response::{ AuthTableQueryResponse, ConfigResponse, ErrorResponse, GetDatabaseResponse, - GetFunctionResponse, GetTableResponse, GetTagResponse, GetViewResponse, ListDatabasesResponse, - ListFunctionsResponse, ListPartitionsResponse, ListPermissionsResponse, ListPoliciesResponse, - ListTablesResponse, ListViewsResponse, PagedList, + GetFunctionResponse, GetTableResponse, GetTagResponse, GetViewResponse, ListBranchesResponse, + ListDatabasesResponse, ListFunctionsResponse, ListPartitionsResponse, ListPermissionsResponse, + ListPoliciesResponse, ListTablesResponse, ListViewsResponse, PagedList, }; use super::auth::{AuthProviderFactory, RESTAuthFunction}; use super::management::{ @@ -80,6 +81,20 @@ fn validate_non_empty_multi(values: &[(&str, &str)]) -> Result<()> { Ok(()) } +/// `url` normalizes dot segments even when their dots are percent encoded. +/// Reject them before a branch mutation can be redirected to a parent resource. +fn validate_resource_path(values: &[(&str, &str)]) -> Result<()> { + validate_non_empty_multi(values)?; + for (value, field_name) in values { + if matches!(*value, "." | "..") { + return Err(crate::Error::ConfigInvalid { + message: format!("{field_name} must not be a dot path segment"), + }); + } + } + Ok(()) +} + /// The canonical `(database, table)` of a policy resource, shared by all three calls. fn policy_table(resource: PermissionResource) -> Result<(String, String)> { resource.validate_policy_attachment()?; @@ -523,6 +538,74 @@ impl RESTApi { Ok(()) } + // ==================== Branch Operations ==================== + + /// Create an empty branch or create one from a tag, as in Java `RESTApi`. + pub async fn create_branch( + &self, + identifier: &Identifier, + branch_name: &str, + from_tag: Option<&str>, + ) -> Result<()> { + validate_resource_path(&[ + (identifier.database(), "database name"), + (identifier.object(), "table name"), + ])?; + let path = self + .resource_paths + .branches(identifier.database(), identifier.object()); + let request = CreateBranchRequest { + branch: branch_name.to_string(), + from_tag: from_tag.map(str::to_string), + }; + let _: serde_json::Value = self.client.post(&path, &request).await?; + Ok(()) + } + + /// Drop branch metadata through the catalog's publication protocol. + pub async fn drop_branch(&self, identifier: &Identifier, branch_name: &str) -> Result<()> { + validate_resource_path(&[ + (identifier.database(), "database name"), + (identifier.object(), "table name"), + (branch_name, "branch name"), + ])?; + let path = + self.resource_paths + .branch(identifier.database(), identifier.object(), branch_name); + let _: serde_json::Value = self.client.delete(&path, None::<&[(&str, &str)]>).await?; + Ok(()) + } + + /// Ask the server to fast-forward from a branch. No files are changed locally. + pub async fn fast_forward(&self, identifier: &Identifier, branch_name: &str) -> Result<()> { + validate_resource_path(&[ + (identifier.database(), "database name"), + (identifier.object(), "table name"), + (branch_name, "branch name"), + ])?; + let path = self.resource_paths.forward_branch( + identifier.database(), + identifier.object(), + branch_name, + ); + let _: serde_json::Value = self.client.post(&path, &ForwardBranchRequest {}).await?; + Ok(()) + } + + /// List the server's branches, preserving its order and null-list behavior. + pub async fn list_branches(&self, identifier: &Identifier) -> Result> { + validate_resource_path(&[ + (identifier.database(), "database name"), + (identifier.object(), "table name"), + ])?; + let path = self + .resource_paths + .branches(identifier.database(), identifier.object()); + let response: ListBranchesResponse = + self.client.get(&path, None::<&[(&str, &str)]>).await?; + Ok(response.branches.unwrap_or_default()) + } + // ==================== Function Operations ==================== /// Create a persistent function. diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 5c48f454b..12e73a849 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -570,6 +570,43 @@ pub trait Catalog: Send + Sync { }) } + // ======================= branch methods ============================= + + /// Create an empty branch, or create one from a tag when `from_tag` is set. + /// Branch names and table capabilities are validated by the catalog. + async fn create_branch( + &self, + _identifier: &Identifier, + _branch_name: &str, + _from_tag: Option<&str>, + ) -> Result<()> { + Err(Error::Unsupported { + message: "branch management is not supported by this catalog".to_string(), + }) + } + + /// Delete a branch's metadata. Shared data files must remain intact. + async fn drop_branch(&self, _identifier: &Identifier, _branch_name: &str) -> Result<()> { + Err(Error::Unsupported { + message: "branch management is not supported by this catalog".to_string(), + }) + } + + /// Replace the identifier's current branch metadata with that of `branch_name`. + /// The REST catalog delegates the operation and its validation to the server. + async fn fast_forward(&self, _identifier: &Identifier, _branch_name: &str) -> Result<()> { + Err(Error::Unsupported { + message: "branch management is not supported by this catalog".to_string(), + }) + } + + /// List branches in the order supplied by the catalog, excluding no names locally. + async fn list_branches(&self, _identifier: &Identifier) -> Result> { + Err(Error::Unsupported { + message: "branch management is not supported by this catalog".to_string(), + }) + } + // ======================= view methods =============================== /// Create a persistent view. diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 1097f799f..3f4f623d2 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -487,6 +487,78 @@ impl Catalog for RESTCatalog { }) } + async fn create_branch( + &self, + identifier: &Identifier, + branch_name: &str, + from_tag: Option<&str>, + ) -> Result<()> { + self.api + .create_branch(identifier, branch_name, from_tag) + .await + .map_err(|error| match error { + Error::RestApi { + source: + RestError::NoSuchResource { + resource_type: Some(kind), + .. + }, + } if kind == "TABLE" => Error::TableNotExist { + full_name: identifier.full_name(), + }, + Error::RestApi { + source: + RestError::NoSuchResource { + resource_type: Some(kind), + .. + }, + } if kind == "TAG" => Error::TagNotExist { + tag_name: from_tag.unwrap_or_default().to_string(), + }, + Error::RestApi { + source: RestError::AlreadyExists { .. }, + } => Error::BranchAlreadyExist { + full_name: identifier.full_name(), + branch_name: branch_name.to_string(), + }, + Error::RestApi { + source: RestError::BadRequest { message }, + } => Error::DataInvalid { + message, + source: None, + }, + other => other, + }) + } + + async fn drop_branch(&self, identifier: &Identifier, branch_name: &str) -> Result<()> { + self.api + .drop_branch(identifier, branch_name) + .await + .map_err(|error| map_rest_error_for_branch(error, identifier, branch_name)) + } + + async fn fast_forward(&self, identifier: &Identifier, branch_name: &str) -> Result<()> { + self.api + .fast_forward(identifier, branch_name) + .await + .map_err(|error| map_rest_error_for_branch(error, identifier, branch_name)) + } + + async fn list_branches(&self, identifier: &Identifier) -> Result> { + self.api + .list_branches(identifier) + .await + .map_err(|error| match error { + Error::RestApi { + source: RestError::NoSuchResource { .. }, + } => Error::TableNotExist { + full_name: identifier.full_name(), + }, + other => other, + }) + } + async fn list_views(&self, database_name: &str) -> Result> { self.api .list_views(database_name) @@ -711,6 +783,18 @@ impl Catalog for RESTCatalog { // Error mapping helpers // ============================================================================ +fn map_rest_error_for_branch(error: Error, identifier: &Identifier, branch_name: &str) -> Error { + match error { + Error::RestApi { + source: RestError::NoSuchResource { .. }, + } => Error::BranchNotExist { + full_name: identifier.full_name(), + branch_name: branch_name.to_string(), + }, + other => other, + } +} + /// Map a REST API error to a catalog-level database error. /// /// Converts `RestError::NoSuchResource` -> `Error::DatabaseNotExist`, diff --git a/crates/paimon/src/error.rs b/crates/paimon/src/error.rs index 0087f16d5..01f877aa5 100644 --- a/crates/paimon/src/error.rs +++ b/crates/paimon/src/error.rs @@ -118,6 +118,16 @@ pub enum Error { TagAlreadyExist { tag_name: String }, #[snafu(display("Tag {} does not exist.", tag_name))] TagNotExist { tag_name: String }, + #[snafu(display("Branch {} of table {} already exists.", branch_name, full_name))] + BranchAlreadyExist { + full_name: String, + branch_name: String, + }, + #[snafu(display("Branch {} of table {} does not exist.", branch_name, full_name))] + BranchNotExist { + full_name: String, + branch_name: String, + }, #[snafu(display("View {} already exists.", full_name))] ViewAlreadyExist { full_name: String }, #[snafu(display("View {} does not exist.", full_name))] diff --git a/crates/paimon/tests/rest_branch_test.rs b/crates/paimon/tests/rest_branch_test.rs new file mode 100644 index 000000000..88eb9174a --- /dev/null +++ b/crates/paimon/tests/rest_branch_test.rs @@ -0,0 +1,327 @@ +// 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. + +//! Java REST branch protocol, exercised through the public Catalog contract. + +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +use axum::body::Bytes; +use axum::extract::State; +use axum::http::{HeaderMap, Method, StatusCode, Uri}; +use axum::response::{IntoResponse, Response}; +use axum::{Json, Router}; +use paimon::api::RestError; +use paimon::catalog::{Catalog, Identifier, RESTCatalog}; +use paimon::common::Options; +use paimon::Error; +use serde_json::{json, Value}; + +#[derive(Debug, PartialEq)] +struct Request { + method: Method, + uri: String, + body: Option, +} + +#[derive(Default)] +struct ServerState { + responses: VecDeque<(StatusCode, Value)>, + requests: Vec, +} + +struct TestCatalog { + catalog: RESTCatalog, + state: Arc>, + server: tokio::task::JoinHandle<()>, +} + +impl Drop for TestCatalog { + fn drop(&mut self) { + self.server.abort(); + } +} + +impl TestCatalog { + async fn new(responses: Vec<(StatusCode, Value)>) -> Self { + async fn handle( + State(state): State>>, + method: Method, + uri: Uri, + headers: HeaderMap, + bytes: Bytes, + ) -> Response { + assert_eq!(headers["authorization"], "Bearer test-token"); + let mut state = state.lock().unwrap(); + state.requests.push(Request { + method, + uri: uri.to_string(), + body: (!bytes.is_empty()).then(|| serde_json::from_slice(&bytes).unwrap()), + }); + let (status, value) = state.responses.pop_front().expect("unexpected request"); + if status == StatusCode::NO_CONTENT { + status.into_response() + } else { + (status, Json(value)).into_response() + } + } + let state = Arc::new(Mutex::new(ServerState { + responses: responses.into(), + requests: vec![], + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let router = Router::new().fallback(handle).with_state(state.clone()); + let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + let mut options = Options::new(); + options.set("uri", format!("http://{address}")); + options.set("warehouse", "warehouse"); + options.set("prefix", "tenant /一"); + options.set("token.provider", "bear"); + options.set("token", "test-token"); + let catalog = RESTCatalog::new(options, false).await.unwrap(); + Self { + catalog, + state, + server, + } + } + + fn requests(&self) -> Vec { + let mut state = self.state.lock().unwrap(); + assert!( + state.responses.is_empty(), + "not all expected requests were sent" + ); + std::mem::take(&mut state.requests) + } +} + +#[tokio::test] +async fn branch_protocol_preserves_names_tags_and_server_order() { + let ctx = TestCatalog::new(vec![ + (StatusCode::NO_CONTENT, Value::Null), + (StatusCode::OK, json!({})), + (StatusCode::OK, json!({"branches": ["z", "a", "main"]})), + (StatusCode::NO_CONTENT, Value::Null), + (StatusCode::NO_CONTENT, Value::Null), + ]) + .await; + let catalog: &dyn Catalog = &ctx.catalog; + let identifier = Identifier::new("db /", "table?$branch_dev"); + catalog + .create_branch(&identifier, "branch /?#", None) + .await + .unwrap(); + catalog + .create_branch(&identifier, "from-tag", Some("tag /一")) + .await + .unwrap(); + assert_eq!( + catalog.list_branches(&identifier).await.unwrap(), + ["z", "a", "main"] + ); + catalog + .fast_forward(&identifier, "branch /?#") + .await + .unwrap(); + catalog + .drop_branch(&identifier, "branch /?#") + .await + .unwrap(); + let base = "/v1/tenant+%2F%E4%B8%80/databases/db+%2F/tables/table%3F%24branch_dev/branches"; + assert_eq!( + ctx.requests(), + vec![ + Request { + method: Method::POST, + uri: base.into(), + body: Some(json!({"branch": "branch /?#", "fromTag": null})) + }, + Request { + method: Method::POST, + uri: base.into(), + body: Some(json!({"branch": "from-tag", "fromTag": "tag /一"})) + }, + Request { + method: Method::GET, + uri: base.into(), + body: None + }, + Request { + method: Method::POST, + uri: format!("{base}/branch+%2F%3F%23/forward"), + body: Some(json!({})) + }, + Request { + method: Method::DELETE, + uri: format!("{base}/branch+%2F%3F%23"), + body: None + }, + ] + ); +} + +#[tokio::test] +async fn null_or_absent_branches_are_empty_like_java() { + let ctx = TestCatalog::new(vec![ + (StatusCode::OK, json!({"branches": null})), + (StatusCode::OK, json!({})), + ]) + .await; + for _ in 0..2 { + assert!(ctx + .catalog + .list_branches(&Identifier::new("db", "t")) + .await + .unwrap() + .is_empty()); + } + assert_eq!(ctx.requests().len(), 2); +} + +#[tokio::test] +async fn create_branch_maps_only_java_catalog_errors() { + let ctx = TestCatalog::new(vec![ + ( + StatusCode::NOT_FOUND, + json!({"code": 404, "resourceType": "TABLE", "message": "table missing"}), + ), + ( + StatusCode::NOT_FOUND, + json!({"code": 404, "resourceType": "TAG", "message": "tag missing"}), + ), + ( + StatusCode::CONFLICT, + json!({"code": 409, "message": "branch exists"}), + ), + ( + StatusCode::BAD_REQUEST, + json!({"code": 400, "message": "invalid branch"}), + ), + ( + StatusCode::NOT_FOUND, + json!({"code": 404, "resourceType": "OTHER", "message": "unknown resource"}), + ), + ]) + .await; + let id = Identifier::new("db", "t"); + assert!( + matches!(ctx.catalog.create_branch(&id, "dev", None).await, Err(Error::TableNotExist { full_name }) if full_name == "db.t") + ); + assert!( + matches!(ctx.catalog.create_branch(&id, "dev", Some("base")).await, Err(Error::TagNotExist { tag_name }) if tag_name == "base") + ); + assert!( + matches!(ctx.catalog.create_branch(&id, "dev", None).await, Err(Error::BranchAlreadyExist { branch_name, .. }) if branch_name == "dev") + ); + assert!( + matches!(ctx.catalog.create_branch(&id, "dev", None).await, Err(Error::DataInvalid { message, .. }) if message == "invalid branch") + ); + assert!(matches!( + ctx.catalog.create_branch(&id, "dev", None).await, + Err(Error::RestApi { + source: RestError::NoSuchResource { .. } + }) + )); + assert_eq!(ctx.requests().len(), 5); +} + +#[tokio::test] +async fn missing_branch_and_table_errors_are_distinct() { + let missing = ( + StatusCode::NOT_FOUND, + json!({"code": 404, "message": "missing"}), + ); + let ctx = TestCatalog::new(vec![missing.clone(), missing.clone(), missing]).await; + let id = Identifier::new("db", "t"); + assert!( + matches!(ctx.catalog.drop_branch(&id, "dev").await, Err(Error::BranchNotExist { full_name, branch_name }) if full_name == "db.t" && branch_name == "dev") + ); + assert!( + matches!(ctx.catalog.fast_forward(&id, "dev").await, Err(Error::BranchNotExist { full_name, branch_name }) if full_name == "db.t" && branch_name == "dev") + ); + assert!( + matches!(ctx.catalog.list_branches(&id).await, Err(Error::TableNotExist { full_name }) if full_name == "db.t") + ); + assert_eq!(ctx.requests().len(), 3); +} + +#[tokio::test] +async fn branch_operations_preserve_forbidden_and_unsupported_responses() { + for status in [StatusCode::FORBIDDEN, StatusCode::NOT_IMPLEMENTED] { + let ctx = TestCatalog::new(vec![ + ( + status, + json!({"code": status.as_u16(), "message": "denied"}) + ); + 4 + ]) + .await; + let id = Identifier::new("db", "t"); + let results = [ + ctx.catalog.create_branch(&id, "dev", None).await, + ctx.catalog.drop_branch(&id, "dev").await, + ctx.catalog.fast_forward(&id, "dev").await, + ctx.catalog.list_branches(&id).await.map(|_| ()), + ]; + for result in results { + assert!(matches!( + result, + Err(Error::RestApi { + source: RestError::Forbidden { .. } | RestError::NotImplemented { .. } + }) + )); + } + assert_eq!(ctx.requests().len(), 4); + } +} + +#[tokio::test] +async fn branch_paths_cannot_address_parent_resources() { + let ctx = TestCatalog::new(vec![]).await; + let id = Identifier::new("db", "t"); + for name in ["", " ", ".", ".."] { + assert!(matches!( + ctx.catalog.drop_branch(&id, name).await, + Err(Error::ConfigInvalid { .. }) + )); + assert!(matches!( + ctx.catalog.fast_forward(&id, name).await, + Err(Error::ConfigInvalid { .. }) + )); + } + for id in [Identifier::new("..", "t"), Identifier::new("db", ".")] { + assert!(matches!( + ctx.catalog.create_branch(&id, "dev", None).await, + Err(Error::ConfigInvalid { .. }) + )); + assert!(matches!( + ctx.catalog.drop_branch(&id, "dev").await, + Err(Error::ConfigInvalid { .. }) + )); + assert!(matches!( + ctx.catalog.fast_forward(&id, "dev").await, + Err(Error::ConfigInvalid { .. }) + )); + assert!(matches!( + ctx.catalog.list_branches(&id).await, + Err(Error::ConfigInvalid { .. }) + )); + } + assert!(ctx.requests().is_empty()); +} From dba13c2af14669c0fc0a4e4294cc82424639e31c Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 17 Oct 2026 09:51:53 +0800 Subject: [PATCH 2/5] [python] Preserve REST branch configuration and structured errors --- bindings/python/README.md | 14 +++ .../python/pypaimon_rust/datafusion.pyi | 10 ++- bindings/python/src/context.rs | 38 ++++++-- bindings/python/src/error.rs | 90 ++++++++++++++++++- .../python/tests/test_catalog_branches.py | 79 ++++++++++++++-- crates/paimon/src/api/rest_error.rs | 47 +++++++++- 6 files changed, 258 insertions(+), 20 deletions(-) diff --git a/bindings/python/README.md b/bindings/python/README.md index 6244c9439..01c9dff6a 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -359,6 +359,20 @@ without reading a branch schema file. Omit it to retain the original branch. Cached time-travel resolution is discarded so the supplied options select the snapshot, with the externally resolved fields preserved. +REST catalogs also expose `create_branch(identifier, branch, from_tag=None)`, +`drop_branch(identifier, branch)`, `fast_forward(identifier, branch)` and +`list_branches(identifier)`. These delegate to core REST catalog operations and +release the GIL while waiting for requests. Creation without a tag creates an +empty branch, as in Java. REST branch rename is unsupported. + +`PaimonCatalog(options, config_required=False)` reuses already resolved REST +options without fetching `/config` again. Branch operation failures raise +`CatalogError`, a `ValueError` subclass with `code`, `resource_type`, +`resource_name` and `message` attributes. `code=0` identifies transport or +response-decoding failures without an HTTP error status. Malformed identifiers +remain `ValueError`, unsupported catalogs raise `NotImplementedError`, and +inherited Native state still raises `ForkSafetyError`. + REST tables load the latest snapshot through the catalog, including empty results and branch-scoped requests. Permission and service failures (including HTTP 501) are propagated as in Java, and the diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index cf0525721..44cee5906 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -749,8 +749,16 @@ class StreamWriteBuilder: def new_update(self) -> StreamTableUpdate: ... def new_commit(self) -> StreamTableCommit: ... +class CatalogError(ValueError): + code: int + resource_type: Optional[str] + resource_name: Optional[str] + message: str + class PaimonCatalog: - def __init__(self, catalog_options: Dict[str, str]) -> None: ... + def __init__(self, catalog_options: Dict[str, str], *, config_required: bool = True) -> None: + """Set config_required=False when the REST configuration has already been resolved.""" + ... def __datafusion_catalog_provider__(self, session: Any) -> object: ... def list_databases(self) -> List[str]: ... def list_tables(self, database_name: str) -> List[str]: ... diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 0954c5022..d72233505 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -41,7 +41,7 @@ use pyo3::prelude::*; use pyo3::types::PyCapsule; use crate::blob::PyBlobReaderRegistry; -use crate::error::{df_to_py_err, to_py_err}; +use crate::error::{df_to_py_err, to_catalog_py_err, to_py_err}; use crate::read::PyRecordBatchReader; use crate::table::PyTable; use crate::udf::{build_python_scalar_udf, udf, PyPythonScalarUDFObject}; @@ -189,12 +189,26 @@ pub(crate) fn complete_storage_options( ) } -fn build_paimon_catalog(catalog_options: HashMap) -> PyResult> { +fn build_paimon_catalog( + catalog_options: HashMap, + config_required: bool, +) -> PyResult> { let catalog_options = complete_storage_options(catalog_options); let rt = runtime(); rt.block_on(async { let options = Options::from_map(catalog_options); - let catalog = CatalogFactory::create(options).await.map_err(to_py_err)?; + let catalog: Arc = + if !config_required && options.get("metastore").map(String::as_str) == Some("rest") { + Arc::new( + paimon::catalog::RESTCatalog::new(options, false) + .await + .map_err(to_catalog_py_err)?, + ) + } else { + CatalogFactory::create(options) + .await + .map_err(to_catalog_py_err)? + }; Ok::<_, PyErr>(catalog) }) } @@ -312,8 +326,13 @@ fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult, catalog_options: HashMap) -> PyResult { - let catalog = py.detach(|| build_paimon_catalog(catalog_options))?; + #[pyo3(signature = (catalog_options, *, config_required=true))] + fn new( + py: Python<'_>, + catalog_options: HashMap, + config_required: bool, + ) -> PyResult { + let catalog = py.detach(|| build_paimon_catalog(catalog_options, config_required))?; let provider = Arc::new( PaimonCatalogProvider::new( None, @@ -399,7 +418,7 @@ impl PaimonCatalog { .create_branch(&identifier, branch_name, from_tag), ) }) - .map_err(to_py_err) + .map_err(to_catalog_py_err) } /// Delete branch metadata through the REST catalog. @@ -411,7 +430,7 @@ impl PaimonCatalog { ) -> PyResult<()> { let identifier = catalog_table_identifier(identifier)?; py.detach(|| runtime().block_on(self.catalog.drop_branch(&identifier, branch_name))) - .map_err(to_py_err) + .map_err(to_catalog_py_err) } /// Fast-forward the identifier's current branch through the REST catalog. @@ -423,7 +442,7 @@ impl PaimonCatalog { ) -> PyResult<()> { let identifier = catalog_table_identifier(identifier)?; py.detach(|| runtime().block_on(self.catalog.fast_forward(&identifier, branch_name))) - .map_err(to_py_err) + .map_err(to_catalog_py_err) } /// List branch names in REST server order. @@ -434,7 +453,7 @@ impl PaimonCatalog { ) -> PyResult> { let identifier = catalog_table_identifier(identifier)?; py.detach(|| runtime().block_on(self.catalog.list_branches(&identifier))) - .map_err(to_py_err) + .map_err(to_catalog_py_err) } } @@ -858,6 +877,7 @@ impl PySQLContext { pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { let this = PyModule::new(py, "datafusion")?; + this.add("CatalogError", py.get_type::())?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs index cb3199db0..b94999cb9 100644 --- a/bindings/python/src/error.rs +++ b/bindings/python/src/error.rs @@ -25,8 +25,96 @@ pyo3::create_exception!( "Raised when native state inherited across process fork cannot be used safely." ); +pyo3::create_exception!( + pypaimon_rust, + CatalogError, + PyValueError, + "A catalog failure with code, resource_type, resource_name and message attributes." +); + pub fn register_module(py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add("ForkSafetyError", py.get_type::()) + module.add("ForkSafetyError", py.get_type::())?; + module.add("CatalogError", py.get_type::()) +} + +/// Preserve the catalog error contract across the FFI without parsing display strings. +pub(crate) fn to_catalog_py_err(error: paimon::Error) -> PyErr { + use paimon::api::RestError; + use paimon::Error; + + if error.is_process_fork_unsupported() { + return to_py_err(error); + } + let (code, resource_type, resource_name, message) = match &error { + Error::BranchAlreadyExist { branch_name, .. } => ( + 409, + Some("BRANCH"), + Some(branch_name.as_str()), + error.to_string(), + ), + Error::BranchNotExist { branch_name, .. } => ( + 404, + Some("BRANCH"), + Some(branch_name.as_str()), + error.to_string(), + ), + Error::TableNotExist { full_name } => ( + 404, + Some("TABLE"), + Some(full_name.as_str()), + error.to_string(), + ), + Error::TagNotExist { tag_name } => { + (404, Some("TAG"), Some(tag_name.as_str()), error.to_string()) + } + Error::DataInvalid { message, .. } => (400, None, None, message.clone()), + Error::RestApi { source } => { + let (code, resource_type, resource_name, message) = match source { + RestError::BadRequest { message } => (400, None, None, message), + RestError::NotAuthorized { message } => (401, None, None, message), + RestError::Forbidden { message } => (403, None, None, message), + RestError::NoSuchResource { + resource_type, + resource_name, + message, + } => ( + 404, + resource_type.as_deref(), + resource_name.as_deref(), + message, + ), + RestError::AlreadyExists { + resource_type, + resource_name, + message, + } => ( + 409, + resource_type.as_deref(), + resource_name.as_deref(), + message, + ), + RestError::ServiceFailure { message } => (500, None, None, message), + RestError::NotImplemented { message } => (501, None, None, message), + RestError::ServiceUnavailable { message } => (503, None, None, message), + RestError::Unexpected { code, message } => (*code, None, None, message), + }; + (code, resource_type, resource_name, message.clone()) + } + Error::UnexpectedError { .. } | Error::IoUnexpected { .. } => { + (0, None, None, error.to_string()) + } + _ => return to_py_err(error), + }; + Python::attach(|py| { + let exception = CatalogError::new_err(error.to_string()); + let value = exception.value(py); + value.setattr("code", code)?; + value.setattr("resource_type", resource_type)?; + value.setattr("resource_name", resource_name)?; + value.setattr("message", message)?; + Ok::<_, PyErr>(exception) + }) + .unwrap_or_else(|error| error) } pub fn to_py_err(err: paimon::Error) -> PyErr { diff --git a/bindings/python/tests/test_catalog_branches.py b/bindings/python/tests/test_catalog_branches.py index e58751543..a46ddcace 100644 --- a/bindings/python/tests/test_catalog_branches.py +++ b/bindings/python/tests/test_catalog_branches.py @@ -24,11 +24,11 @@ import pytest -from pypaimon_rust.datafusion import PaimonCatalog +from pypaimon_rust.datafusion import CatalogError, PaimonCatalog @pytest.fixture -def catalog(monkeypatch): +def rest_server(monkeypatch): requests = [] monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1") monkeypatch.setenv("no_proxy", "localhost,127.0.0.1") @@ -48,6 +48,15 @@ def handle_operation(self): payload, status = {"message": "denied", "code": 403}, 403 elif "conflict" in parts: payload, status = {"message": "exists", "code": 409}, 409 + elif "nocode" in parts: + payload, status = {"message": "denied"}, 403 + elif any(part.startswith("status_") for part in parts): + status = int(next(part for part in parts if part.startswith("status_")).split("_")[1]) + payload = {"message": "unmapped status"} + elif "invalidjson" in parts: + payload, status = "not a branches response", 200 + elif body and body.get("fromTag") == "missing": + payload, status = {"resourceType": "TAG", "message": "tag missing", "code": 404}, 404 elif self.command == "GET": payload, status = {"branches": ["dev", "release"]}, 200 else: @@ -70,18 +79,31 @@ def log_message(self, *args): thread = threading.Thread(target=server.serve_forever) thread.start() try: - native = PaimonCatalog({ + yield { "metastore": "rest", "uri": "http://localhost:%d" % server.server_port, "warehouse": "warehouse", "token.provider": "bear", "token": "test-token", - }) - requests.clear() - yield native, requests + }, requests finally: server.shutdown() server.server_close() thread.join() +@pytest.fixture +def catalog(rest_server): + options, requests = rest_server + native = PaimonCatalog(options) + requests.clear() + return native, requests + + +def test_resolved_rest_configuration_is_not_fetched_again(rest_server): + options, requests = rest_server + native = PaimonCatalog(dict(options, prefix="resolved-prefix"), config_required=False) + assert native.list_branches("db.t") == ["dev", "release"] + assert requests == [("GET", "/v1/resolved-prefix/databases/db/tables/t/branches", None)] + + def test_rest_branch_operations_use_core_protocol(catalog): native, requests = catalog identifier = ("db.name", "table.name$branch_dev") @@ -146,6 +168,51 @@ def test_duplicate_branch_is_reported(catalog): assert len(requests) == 1 +@pytest.mark.parametrize("operation,identifier,args,code,resource_type,resource_name,message", [ + ("create_branch", "db.conflict", ("dev",), 409, "BRANCH", "dev", "already exists"), + ("create_branch", "db.missing", ("dev",), 404, "TABLE", "db.missing", "does not exist"), + ("create_branch", "db.table", ("dev", "missing"), 404, "TAG", "missing", "does not exist"), + ("drop_branch", "db.missing", ("dev",), 404, "BRANCH", "dev", "does not exist"), + ("fast_forward", "db.missing", ("dev",), 404, "BRANCH", "dev", "does not exist"), + ("list_branches", "db.missing", (), 404, "TABLE", "db.missing", "does not exist"), + ("list_branches", "db.forbidden", (), 403, None, None, "denied"), + ("list_branches", "db.nocode", (), 403, None, None, "denied"), + ("list_branches", "db.status_429", (), 429, None, None, "unmapped status"), + ("list_branches", "db.status_502", (), 502, None, None, "unmapped status"), + ("list_branches", "db.status_504", (), 504, None, None, "unmapped status"), + ("list_branches", "db.invalidjson", (), 0, None, None, "failed to parse json"), +]) +def test_catalog_errors_preserve_structured_details( + catalog, operation, identifier, args, code, resource_type, resource_name, message): + native, requests = catalog + with pytest.raises(CatalogError) as raised: + getattr(native, operation)(identifier, *args) + error = raised.value + assert error.code == code + assert error.resource_type == resource_type + assert error.resource_name == resource_name + assert message in error.message + assert isinstance(error, ValueError) + assert len(requests) == 1 + + +def test_transport_failure_has_structured_catalog_error(): + import socket + + with socket.socket() as listener: + listener.bind(("localhost", 0)) + native = PaimonCatalog({ + "metastore": "rest", "uri": "http://localhost:%d" % listener.getsockname()[1], + "warehouse": "warehouse", "token.provider": "bear", "token": "test-token", + }, config_required=False) + with pytest.raises(CatalogError) as raised: + native.drop_branch("db.t", "dev") + assert raised.value.code == 0 + assert raised.value.resource_type is None + assert raised.value.resource_name is None + assert "http delete failed" in raised.value.message + + @pytest.mark.parametrize("operation", ["create_branch", "drop_branch", "fast_forward", "list_branches"]) def test_catalog_without_branch_management_reports_unsupported(tmp_path, operation): native = PaimonCatalog({"warehouse": str(tmp_path)}) diff --git a/crates/paimon/src/api/rest_error.rs b/crates/paimon/src/api/rest_error.rs index 8b4477dde..0ea3af75c 100644 --- a/crates/paimon/src/api/rest_error.rs +++ b/crates/paimon/src/api/rest_error.rs @@ -77,7 +77,7 @@ pub enum RestError { /// Unexpected error #[snafu(display("Unexpected error: {}", message))] - Unexpected { message: String }, + Unexpected { code: i32, message: String }, } use super::api_response::ErrorResponse; @@ -89,7 +89,8 @@ impl RestError { /// creates an ErrorResponse with the raw response body as the message and the /// HTTP status code as the code. pub fn parse_error_response(text: &str, status_code: u16) -> ErrorResponse { - if let Ok(error) = serde_json::from_str::(text) { + if let Ok(mut error) = serde_json::from_str::(text) { + error.code = error.code.or(Some(status_code as i32)); if error.message.is_some() { return error; } @@ -140,7 +141,47 @@ impl RestError { 500 => RestError::ServiceFailure { message }, 501 => RestError::NotImplemented { message }, 503 => RestError::ServiceUnavailable { message }, - _ => RestError::Unexpected { message }, + _ => RestError::Unexpected { code, message }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn missing_json_code_uses_http_status() { + for code in [400, 401, 403, 404, 409, 500, 501, 503] { + let response = RestError::parse_error_response( + r#"{"resourceType":"TABLE","resourceName":"t","message":"denied"}"#, + code, + ); + assert_eq!(response.code, Some(i32::from(code))); + assert_eq!(response.message.as_deref(), Some("denied")); + assert_eq!(response.resource_type.as_deref(), Some("TABLE")); + assert_eq!(response.resource_name.as_deref(), Some("t")); + } + } + + #[test] + fn explicit_json_code_takes_precedence() { + let response = RestError::parse_error_response(r#"{"message":"denied","code":403}"#, 500); + assert!(matches!( + RestError::from_error_response(response), + RestError::Forbidden { message } if message == "denied" + )); + } + + #[test] + fn unknown_status_preserves_code_and_raw_message() { + for code in [429, 502, 504] { + let response = RestError::parse_error_response("gateway unavailable", code); + assert!(matches!( + RestError::from_error_response(response), + RestError::Unexpected { code: actual, message } + if actual == i32::from(code) && message == "gateway unavailable" + )); } } } From fc11cf539108f87e5da2d3af36424d09bd620d71 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 17 Oct 2026 10:31:38 +0800 Subject: [PATCH 3/5] [core] Align Native REST branch transport and OpenAPI V4 with Java --- bindings/python/README.md | 6 + .../python/tests/test_catalog_branches.py | 2 +- crates/paimon/src/api/auth/base.rs | 10 + crates/paimon/src/api/auth/bearer_provider.rs | 16 +- .../src/api/auth/dlf_openapi_actions.rs | 346 ++++++++++++++++++ .../src/api/auth/dlf_openapi_v4_signer.rs | 279 ++++++++++++++ crates/paimon/src/api/auth/dlf_provider.rs | 16 +- crates/paimon/src/api/auth/dlf_signer.rs | 49 +-- crates/paimon/src/api/auth/mod.rs | 2 + crates/paimon/src/api/rest_api.rs | 11 +- crates/paimon/src/api/rest_client.rs | 177 +++++++-- crates/paimon/src/common/options.rs | 2 +- crates/paimon/tests/rest_branch_test.rs | 117 +++++- 13 files changed, 949 insertions(+), 84 deletions(-) create mode 100644 crates/paimon/src/api/auth/dlf_openapi_actions.rs create mode 100644 crates/paimon/src/api/auth/dlf_openapi_v4_signer.rs diff --git a/bindings/python/README.md b/bindings/python/README.md index 01c9dff6a..5ce313ee6 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -365,6 +365,12 @@ REST catalogs also expose `create_branch(identifier, branch, from_tag=None)`, release the GIL while waiting for requests. Creation without a tag creates an empty branch, as in Java. REST branch rename is unsupported. +REST authentication also supports explicit `dlf.signing-algorithm=openapi-v4`. +Rust core signs ACS4-HMAC-SHA256 requests with the same canonical request, +credential scope and action headers as Java. GET requests retry HTTP 429/503 up +to five times and honor Retry-After; branch mutations are sent once and ignore +successful response payloads. + `PaimonCatalog(options, config_required=False)` reuses already resolved REST options without fetching `/config` again. Branch operation failures raise `CatalogError`, a `ValueError` subclass with `code`, `resource_type`, diff --git a/bindings/python/tests/test_catalog_branches.py b/bindings/python/tests/test_catalog_branches.py index a46ddcace..a5cef1e9c 100644 --- a/bindings/python/tests/test_catalog_branches.py +++ b/bindings/python/tests/test_catalog_branches.py @@ -177,7 +177,7 @@ def test_duplicate_branch_is_reported(catalog): ("list_branches", "db.missing", (), 404, "TABLE", "db.missing", "does not exist"), ("list_branches", "db.forbidden", (), 403, None, None, "denied"), ("list_branches", "db.nocode", (), 403, None, None, "denied"), - ("list_branches", "db.status_429", (), 429, None, None, "unmapped status"), + ("create_branch", "db.status_429", ("dev",), 429, None, None, "unmapped status"), ("list_branches", "db.status_502", (), 502, None, None, "unmapped status"), ("list_branches", "db.status_504", (), 504, None, None, "unmapped status"), ("list_branches", "db.invalidjson", (), 0, None, None, "failed to parse json"), diff --git a/crates/paimon/src/api/auth/base.rs b/crates/paimon/src/api/auth/base.rs index edf97c284..45fd2a50d 100644 --- a/crates/paimon/src/api/auth/base.rs +++ b/crates/paimon/src/api/auth/base.rs @@ -99,6 +99,16 @@ pub trait AuthProvider: Send + Sync { } /// Authorization header key. pub const AUTHORIZATION_HEADER_KEY: &str = "Authorization"; + +/// Authentication headers override caller headers regardless of spelling case. +pub(super) fn merge_auth_headers( + mut base: HashMap, + generated: HashMap, +) -> HashMap { + base.retain(|key, _| !generated.keys().any(|name| key.eq_ignore_ascii_case(name))); + base.extend(generated); + base +} /// Function wrapper for REST authentication. /// /// This struct combines an initial set of headers with an authentication provider diff --git a/crates/paimon/src/api/auth/bearer_provider.rs b/crates/paimon/src/api/auth/bearer_provider.rs index ba6d4dd41..b2c3ddd30 100644 --- a/crates/paimon/src/api/auth/bearer_provider.rs +++ b/crates/paimon/src/api/auth/bearer_provider.rs @@ -21,7 +21,7 @@ use std::collections::HashMap; use async_trait::async_trait; -use super::base::{AuthProvider, RESTAuthParameter}; +use super::base::{merge_auth_headers, AuthProvider, RESTAuthParameter}; /// Authentication provider using Bearer token. /// @@ -47,14 +47,16 @@ impl BearerTokenAuthProvider { impl AuthProvider for BearerTokenAuthProvider { async fn merge_auth_header( &self, - mut base_header: HashMap, + base_header: HashMap, _parameter: &RESTAuthParameter, ) -> crate::Result> { - base_header.insert( - "Authorization".to_string(), - format!("Bearer {}", self.token), - ); - Ok(base_header) + Ok(merge_auth_headers( + base_header, + HashMap::from([( + "Authorization".to_string(), + format!("Bearer {}", self.token), + )]), + )) } } diff --git a/crates/paimon/src/api/auth/dlf_openapi_actions.rs b/crates/paimon/src/api/auth/dlf_openapi_actions.rs new file mode 100644 index 000000000..a3309d347 --- /dev/null +++ b/crates/paimon/src/api/auth/dlf_openapi_actions.rs @@ -0,0 +1,346 @@ +// 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. + +//! Java's DlfNext 2026-01-18 REST route to OpenAPI action mapping. + +const OPERATIONS: &[(&str, &str, &str)] = &[ + ("GET", "/v1/config", "GetConfig"), + ("GET", "/v1/{prefix}/databases", "ListDatabases"), + ("POST", "/v1/{prefix}/databases", "CreateDatabase"), + ("GET", "/v1/{prefix}/databases/{database}", "GetDatabase"), + ("POST", "/v1/{prefix}/databases/{database}", "AlterDatabase"), + ( + "DELETE", + "/v1/{prefix}/databases/{database}", + "DropDatabase", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/register", + "RegisterTable", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables", + "ListTables", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables", + "CreateTable", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/table-details", + "ListTableDetails", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}", + "GetTable", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}", + "AlterTable", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/tables/{table}", + "DropTable", + ), + ("POST", "/v1/{prefix}/tables/rename", "RenameTable"), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/auth", + "AuthTableQuery", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/commit", + "CommitTable", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/token", + "GetTableToken", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/rollback", + "RollbackToSnapshot", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/snapshot", + "GetTableSnapshot", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/snapshots", + "ListSnapshots", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/snapshots/{version}", + "GetVersionSnapshot", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions", + "ListPartitions", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions", + "CreatePartitions", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions/drop", + "DropPartitions", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions/mark", + "MarkDonePartitions", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions/list-by-names", + "ListPartitionsByNames", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/partitions/list-by-filter", + "ListPartitionsByFilter", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/branches", + "ListBranches", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/branches", + "CreateBranch", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/tables/{table}/branches/{branch}", + "DropBranch", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/branches/{branch}/forward", + "FastForwardBranch", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/tags", + "ListTags", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/tags", + "CreateTag", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/tags/{tag}", + "GetTag", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/tables/{table}/tags/{tag}", + "DropTag", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/tables/{table}/policies", + "ListPolicies", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/policies", + "CreatePolicy", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/tables/{table}/policies/drop", + "DropPolicy", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/views", + "ListViews", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/views", + "CreateView", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/view-details", + "ListViewDetails", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/views/{view}", + "GetView", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/views/{view}", + "AlterView", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/views/{view}", + "DropView", + ), + ("POST", "/v1/{prefix}/views/rename", "RenameView"), + ( + "GET", + "/v1/{prefix}/databases/{database}/semantic-views", + "ListSemanticViews", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/semantic-views/{semanticView}", + "GetSemanticView", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/semantic-views/{semanticView}", + "UpsertSemanticView", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/semantic-views/{semanticView}", + "DeleteSemanticView", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/functions", + "ListFunctions", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/functions", + "CreateFunction", + ), + ( + "GET", + "/v1/{prefix}/databases/{database}/functions/{function}", + "GetFunction", + ), + ( + "POST", + "/v1/{prefix}/databases/{database}/functions/{function}", + "AlterFunction", + ), + ( + "DELETE", + "/v1/{prefix}/databases/{database}/functions/{function}", + "DropFunction", + ), + ( + "GET", + "/v1/{prefix}/permissions", + "ListPermissionAssignments", + ), + ( + "POST", + "/v1/{prefix}/permissions/grant", + "GrantPermissionAssignment", + ), + ( + "POST", + "/v1/{prefix}/permissions/revoke", + "RevokePermissionAssignment", + ), + ("GET", "/v1/{prefix}/labels", "ListLabels"), + ("POST", "/v1/{prefix}/labels/get", "GetLabel"), + ("POST", "/v1/{prefix}/labels/upsert", "UpsertLabel"), + ("POST", "/v1/{prefix}/labels/delete", "DeleteLabel"), +]; + +pub(super) fn resolve_action(method: &str, path: &str) -> Option<&'static str> { + OPERATIONS.iter().find_map(|(verb, template, action)| { + (method.eq_ignore_ascii_case(verb) + && template.split('/').count() == path.split('/').count() + && template + .split('/') + .zip(path.split('/')) + .all(|(part, segment)| { + if part.starts_with('{') { + !segment.is_empty() + } else { + part == segment + } + })) + .then_some(*action) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_java_operation_resolves_without_collisions() { + for (method, template, action) in OPERATIONS { + let path = template + .split('/') + .map(|part| if part.starts_with('{') { "x1" } else { part }) + .collect::>() + .join("/"); + assert_eq!( + resolve_action(method, &path), + Some(*action), + "{method} {path}" + ); + } + } + + #[test] + fn action_resolution_preserves_encoded_names_and_literal_boundaries() { + assert_eq!( + resolve_action("get", "/v1/rename/databases/tables"), + Some("GetDatabase") + ); + assert_eq!( + resolve_action("GET", "/v1/p/databases/a%2Fb/tables/c%20d"), + Some("GetTable") + ); + assert_eq!( + resolve_action("DELETE", "/v1/p/databases/db/tables/t/branches/forward"), + Some("DropBranch") + ); + for (method, path) in [ + ("GET", "/v1/p/tables"), + ("GET", "/v1/config/"), + ("GET", "/v1//databases"), + ("PUT", "/v1/p/databases/db/tables/t"), + ("POST", "/v1/p/databases/db/tables/t/branches/b/rename"), + ] { + assert_eq!(resolve_action(method, path), None, "{method} {path}"); + } + } +} diff --git a/crates/paimon/src/api/auth/dlf_openapi_v4_signer.rs b/crates/paimon/src/api/auth/dlf_openapi_v4_signer.rs new file mode 100644 index 000000000..98dfdb873 --- /dev/null +++ b/crates/paimon/src/api/auth/dlf_openapi_v4_signer.rs @@ -0,0 +1,279 @@ +// 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. + +//! Java's Aliyun OpenAPI V4 (ACS4-HMAC-SHA256) request signer. + +use std::collections::{BTreeMap, HashMap}; + +use chrono::{DateTime, Utc}; +use uuid::Uuid; + +use super::base::RESTAuthParameter; +use super::dlf_openapi_actions::resolve_action; +use super::dlf_provider::DLFToken; +use super::dlf_signer::{hmac_sha256, sha256_hex, DLFRequestSigner}; +use crate::api::rest_util::RESTUtil; + +pub(super) struct DLFOpenApiV4Signer { + region: String, +} + +impl DLFOpenApiV4Signer { + pub const IDENTIFIER: &'static str = "openapi-v4"; + const ALGORITHM: &'static str = "ACS4-HMAC-SHA256"; + const PRODUCT: &'static str = "DlfNext"; + const REQUEST_TYPE: &'static str = "aliyun_v4_request"; + + pub(super) fn new(region: &str) -> Self { + Self { + region: region.to_string(), + } + } + + fn percent_encode(value: &str) -> String { + RESTUtil::encode_string(value) + .replace('+', "%20") + .replace('*', "%2A") + .replace("%7E", "~") + } + + fn canonical_query(parameters: &HashMap) -> String { + let sorted: BTreeMap<_, _> = parameters.iter().collect(); + sorted + .into_iter() + .map(|(key, value)| { + // RESTAuthParameter has already encoded each value once, as in Java. + format!( + "{}={}", + Self::percent_encode(key), + Self::percent_encode(&RESTUtil::decode_string(value)) + ) + }) + .collect::>() + .join("&") + } + + fn canonical_headers(headers: &HashMap) -> (String, String) { + let sorted: BTreeMap<_, _> = headers + .iter() + .filter_map(|(key, value)| { + let key = key.to_ascii_lowercase(); + (key.starts_with("x-acs-") || key == "host" || key == "content-type") + .then(|| (key, value.trim())) + }) + .collect(); + let canonical = sorted + .iter() + .map(|(key, value)| format!("{key}:{value}\n")) + .collect::(); + let signed = sorted.keys().cloned().collect::>().join(";"); + (canonical, signed) + } +} + +impl DLFRequestSigner for DLFOpenApiV4Signer { + fn sign_headers( + &self, + body: Option<&str>, + now: &DateTime, + security_token: Option<&str>, + host: &str, + ) -> HashMap { + let mut headers = HashMap::from([ + ( + "x-acs-date".to_string(), + now.format("%Y-%m-%dT%H:%M:%SZ").to_string(), + ), + ("host".to_string(), host.to_string()), + ( + "x-acs-content-sha256".to_string(), + sha256_hex(body.unwrap_or_default()), + ), + ( + "x-acs-signature-nonce".to_string(), + Uuid::new_v4().to_string(), + ), + ("x-acs-version".to_string(), "2026-01-18".to_string()), + ]); + if body.is_some_and(|body| !body.is_empty()) { + headers.insert("content-type".to_string(), "application/json".to_string()); + } + if let Some(token) = security_token { + headers.insert("x-acs-security-token".to_string(), token.to_string()); + } + headers + } + + fn authorization( + &self, + parameter: &RESTAuthParameter, + token: &DLFToken, + _host: &str, + headers: &mut HashMap, + ) -> String { + // The action has to be both signed and included in the outgoing headers. + if let Some(action) = resolve_action(¶meter.method, ¶meter.path) { + headers.insert("x-acs-action".to_string(), action.to_string()); + } + let (canonical_headers, signed_headers) = Self::canonical_headers(headers); + let uri = if parameter.path.trim().is_empty() { + "/" + } else { + ¶meter.path + }; + let empty_hash = sha256_hex(""); + let payload = headers.get("x-acs-content-sha256").unwrap_or(&empty_hash); + let canonical = format!( + "{}\n{}\n{}\n{}\n{}\n{}", + parameter.method, + uri, + Self::canonical_query(¶meter.parameters), + canonical_headers, + signed_headers, + payload + ); + let to_sign = format!("{}\n{}", Self::ALGORITHM, sha256_hex(&canonical)); + let date = headers["x-acs-date"][..10].replace('-', ""); + let date_key = hmac_sha256( + format!("aliyun_v4{}", token.access_key_secret).as_bytes(), + &date, + ); + let region_key = hmac_sha256(&date_key, &self.region); + let product_key = hmac_sha256(®ion_key, Self::PRODUCT); + let signing_key = hmac_sha256(&product_key, Self::REQUEST_TYPE); + let signature = hex::encode(hmac_sha256(&signing_key, &to_sign)); + format!( + "{} Credential={}/{}/{}/{}/{},SignedHeaders={},Signature={}", + Self::ALGORITHM, + token.access_key_id, + date, + self.region, + Self::PRODUCT, + Self::REQUEST_TYPE, + signed_headers, + signature + ) + } + + fn identifier(&self) -> &str { + Self::IDENTIFIER + } +} + +#[cfg(test)] +mod tests { + use super::super::dlf_signer::DLFSignerFactory; + use super::*; + + const HOST: &str = "dlfnext.cn-hangzhou.aliyuncs.com"; + + fn now() -> DateTime { + DateTime::parse_from_rfc3339("2025-04-16T03:44:46Z") + .unwrap() + .with_timezone(&Utc) + } + + #[test] + fn java_known_signatures_cover_encoded_queries_and_actions() { + let signer = DLFSignerFactory::create_signer("openapi-v4", "cn-hangzhou"); + assert_eq!(signer.identifier(), "openapi-v4"); + let token = DLFToken::new("TestAKId", "TestAKSecret", None, None, None); + let parameter = RESTAuthParameter::for_get( + "/v1/prefix/functions", + HashMap::from([ + ("pageToken".to_string(), "a b/c".to_string()), + ("functionNamePattern".to_string(), "func%".to_string()), + ]), + ); + let mut headers = signer.sign_headers(None, &now(), None, HOST); + headers.insert( + "x-acs-signature-nonce".to_string(), + "fixed-nonce-for-test".to_string(), + ); + assert_eq!(signer.authorization(¶meter, &token, HOST, &mut headers), concat!( + "ACS4-HMAC-SHA256 Credential=TestAKId/20250416/cn-hangzhou/DlfNext/aliyun_v4_request,", + "SignedHeaders=host;x-acs-content-sha256;x-acs-date;x-acs-signature-nonce;x-acs-version,", + "Signature=9812bacd92dd1152e7ea32aacfad0545835d6a18543852364f1a1ea0cba6219a")); + + let body = r#"{"identifier":{"database":"db","object":"t"}}"#; + let parameter = + RESTAuthParameter::for_post("/v1/clg-paimon-1/databases/db/tables", body.to_string()); + let mut headers = signer.sign_headers(Some(body), &now(), None, HOST); + headers.insert( + "x-acs-signature-nonce".to_string(), + "fixed-nonce-for-test".to_string(), + ); + assert_eq!( + signer.authorization(¶meter, &token, HOST, &mut headers), + concat!( + "ACS4-HMAC-SHA256 Credential=TestAKId/20250416/cn-hangzhou/DlfNext/aliyun_v4_request,", + "SignedHeaders=content-type;host;x-acs-action;x-acs-content-sha256;", + "x-acs-date;x-acs-signature-nonce;x-acs-version,", + "Signature=420206b5263536e6bc271a7a32582b4a820e23c8c58ae692b708c69f9d19e51d") + ); + assert_eq!(headers["x-acs-action"], "CreateTable"); + } + + #[test] + fn canonical_query_uses_rfc3986_and_decodes_values_once() { + let parameter = RESTAuthParameter::for_get( + "/", + HashMap::from([ + ("z".to_string(), "".to_string()), + ("a b*~".to_string(), "a+b /%一".to_string()), + ]), + ); + assert_eq!( + DLFOpenApiV4Signer::canonical_query(¶meter.parameters), + "a%20b%2A~=a%2Bb%20%2F%25%E4%B8%80&z=" + ); + } + + #[test] + fn payload_token_and_region_are_covered_by_signature() { + let signer = DLFOpenApiV4Signer::new("cn-hangzhou"); + let token = DLFToken::new("key", "secret", Some("sts".to_string()), None, None); + let parameter = RESTAuthParameter::for_get("/test", HashMap::new()); + let mut headers = signer.sign_headers(None, &now(), Some("sts"), HOST); + assert_eq!(headers["x-acs-date"], "2025-04-16T03:44:46Z"); + assert_eq!( + headers["x-acs-content-sha256"], + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ); + assert!(!headers.contains_key("content-type")); + let original = signer.authorization(¶meter, &token, HOST, &mut headers); + assert!(original.contains(";x-acs-security-token;")); + headers.insert("x-acs-security-token".to_string(), "changed".to_string()); + assert_ne!( + original, + signer.authorization(¶meter, &token, HOST, &mut headers) + ); + headers.insert("x-acs-security-token".to_string(), "sts".to_string()); + let elsewhere = DLFOpenApiV4Signer::new("cn-beijing"); + assert_ne!( + original, + elsewhere.authorization(¶meter, &token, HOST, &mut headers) + ); + let body = signer.sign_headers(Some("一"), &now(), None, HOST); + assert_eq!(body["x-acs-content-sha256"], sha256_hex("一")); + assert_eq!(body["content-type"], "application/json"); + assert_ne!( + headers["x-acs-signature-nonce"], + body["x-acs-signature-nonce"] + ); + } +} diff --git a/crates/paimon/src/api/auth/dlf_provider.rs b/crates/paimon/src/api/auth/dlf_provider.rs index 18b0ccc61..486d66dae 100644 --- a/crates/paimon/src/api/auth/dlf_provider.rs +++ b/crates/paimon/src/api/auth/dlf_provider.rs @@ -25,7 +25,7 @@ use chrono::Utc; use reqwest::Client; use serde::{Deserialize, Serialize}; -use super::base::{AuthProvider, RESTAuthParameter, AUTHORIZATION_HEADER_KEY}; +use super::base::{merge_auth_headers, AuthProvider, RESTAuthParameter, AUTHORIZATION_HEADER_KEY}; use super::dlf_signer::{DLFRequestSigner, DLFSignerFactory}; use crate::api::user_agent; use crate::common::{CatalogOptions, Options}; @@ -334,7 +334,7 @@ const TOKEN_EXPIRATION_SAFE_TIME_MILLIS: i64 = 3_600_000; /// /// This provider implements authentication for Alibaba Cloud DLF service, /// supporting both VPC endpoints (DLF4-HMAC-SHA256) and public endpoints -/// (ROA v2 HMAC-SHA1). +/// (ROA v2 HMAC-SHA1 or explicit OpenAPI V4 ACS4-HMAC-SHA256). pub struct DLFAuthProvider { uri: String, token: tokio::sync::Mutex>, @@ -426,7 +426,7 @@ impl DLFAuthProvider { impl AuthProvider for DLFAuthProvider { async fn merge_auth_header( &self, - mut base_header: HashMap, + base_header: HashMap, rest_auth_parameter: &RESTAuthParameter, ) -> crate::Result> { // Get token (will auto-refresh if needed via token_loader) @@ -436,7 +436,7 @@ impl AuthProvider for DLFAuthProvider { let host = Self::extract_host(&self.uri); // Generate signature headers - let sign_headers = self.signer.sign_headers( + let mut sign_headers = self.signer.sign_headers( rest_auth_parameter.data.as_deref(), &now, token.security_token.as_deref(), @@ -446,13 +446,11 @@ impl AuthProvider for DLFAuthProvider { // Generate authorization header let authorization = self.signer - .authorization(rest_auth_parameter, &token, &host, &sign_headers); + .authorization(rest_auth_parameter, &token, &host, &mut sign_headers); // Merge all headers - base_header.extend(sign_headers); - base_header.insert(AUTHORIZATION_HEADER_KEY.to_string(), authorization); - - Ok(base_header) + sign_headers.insert(AUTHORIZATION_HEADER_KEY.to_string(), authorization); + Ok(merge_auth_headers(base_header, sign_headers)) } } diff --git a/crates/paimon/src/api/auth/dlf_signer.rs b/crates/paimon/src/api/auth/dlf_signer.rs index 4873b8fdf..8fc602568 100644 --- a/crates/paimon/src/api/auth/dlf_signer.rs +++ b/crates/paimon/src/api/auth/dlf_signer.rs @@ -17,7 +17,7 @@ //! DLF Request Signer implementations for Alibaba Cloud Data Lake Formation. //! -//! This module provides two signature algorithms for authenticating requests +//! This module provides signature algorithms for authenticating requests //! to Alibaba Cloud Data Lake Formation (DLF) service: //! //! # Signature Algorithms @@ -72,6 +72,7 @@ //! The signer is automatically selected based on the endpoint URI: //! - VPC endpoints → `DLFDefaultSigner` (DLF4-HMAC-SHA256) //! - Public endpoints with "dlfnext" or "openapi" → `DLFOpenApiSigner` (HMAC-SHA1) +//! - Explicit `openapi-v4` → `DLFOpenApiV4Signer` (ACS4-HMAC-SHA256) use std::collections::HashMap; @@ -84,12 +85,23 @@ use sha2::{Digest, Sha256}; use uuid::Uuid; use super::base::RESTAuthParameter; +use super::dlf_openapi_v4_signer::DLFOpenApiV4Signer; use super::dlf_provider::DLFToken; use crate::api::rest_util::RESTUtil; type HmacSha256 = Hmac; type HmacSha1 = Hmac; +pub(super) fn hmac_sha256(key: &[u8], data: &str) -> Vec { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC can take key of any size"); + mac.update(data.as_bytes()); + mac.finalize().into_bytes().to_vec() +} + +pub(super) fn sha256_hex(raw: &str) -> String { + hex::encode(Sha256::digest(raw.as_bytes())) +} + /// Trait for DLF request signers. /// /// Different signers implement different signature algorithms for @@ -99,6 +111,7 @@ type HmacSha1 = Hmac; /// /// - [`DLFDefaultSigner`]: Uses DLF4-HMAC-SHA256 for VPC endpoints /// - [`DLFOpenApiSigner`]: Uses HMAC-SHA1 for public endpoints +/// - [`DLFOpenApiV4Signer`]: Uses ACS4-HMAC-SHA256 when explicitly configured pub trait DLFRequestSigner: Send + Sync { /// Generate signature headers for the request. fn sign_headers( @@ -115,7 +128,7 @@ pub trait DLFRequestSigner: Send + Sync { rest_auth_parameter: &RESTAuthParameter, token: &DLFToken, host: &str, - sign_headers: &HashMap, + sign_headers: &mut HashMap, ) -> String; #[allow(dead_code)] /// Get the identifier for this signer. @@ -201,18 +214,6 @@ impl DLFDefaultSigner { BASE64_STANDARD.encode(hash) } - fn hmac_sha256(key: &[u8], data: &str) -> Vec { - let mut mac = HmacSha256::new_from_slice(key).expect("HMAC can take key of any size"); - mac.update(data.as_bytes()); - mac.finalize().into_bytes().to_vec() - } - - fn sha256_hex(raw: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(raw.as_bytes()); - hex::encode(hasher.finalize()) - } - fn hex_encode(raw: &[u8]) -> String { hex::encode(raw) } @@ -340,7 +341,7 @@ impl DLFRequestSigner for DLFDefaultSigner { rest_auth_parameter: &RESTAuthParameter, token: &DLFToken, _host: &str, - sign_headers: &HashMap, + sign_headers: &mut HashMap, ) -> String { let date_time = sign_headers.get(Self::DLF_DATE_HEADER_KEY).unwrap(); let date = &date_time[..8]; @@ -357,20 +358,20 @@ impl DLFRequestSigner for DLFDefaultSigner { Self::PRODUCT, Self::REQUEST_TYPE ), - Self::sha256_hex(&canonical_request), + sha256_hex(&canonical_request), ] .join(Self::NEW_LINE); // Derive signing key - let date_key = Self::hmac_sha256( + let date_key = hmac_sha256( format!("aliyun_v4{}", token.access_key_secret).as_bytes(), date, ); - let date_region_key = Self::hmac_sha256(&date_key, &self.region); - let date_region_service_key = Self::hmac_sha256(&date_region_key, Self::PRODUCT); - let signing_key = Self::hmac_sha256(&date_region_service_key, Self::REQUEST_TYPE); + let date_region_key = hmac_sha256(&date_key, &self.region); + let date_region_service_key = hmac_sha256(&date_region_key, Self::PRODUCT); + let signing_key = hmac_sha256(&date_region_service_key, Self::REQUEST_TYPE); - let signature_bytes = Self::hmac_sha256(&signing_key, &string_to_sign); + let signature_bytes = hmac_sha256(&signing_key, &string_to_sign); let signature = Self::hex_encode(&signature_bytes); format!( @@ -610,7 +611,7 @@ impl DLFRequestSigner for DLFOpenApiSigner { rest_auth_parameter: &RESTAuthParameter, token: &DLFToken, _host: &str, - sign_headers: &HashMap, + sign_headers: &mut HashMap, ) -> String { let canonicalized_headers = self.build_canonicalized_headers(sign_headers); let canonicalized_resource = self.build_canonicalized_resource(rest_auth_parameter); @@ -642,7 +643,9 @@ pub struct DLFSignerFactory; impl DLFSignerFactory { /// Create a signer based on the signing algorithm. pub fn create_signer(signing_algorithm: &str, region: &str) -> Box { - if signing_algorithm == DLFOpenApiSigner::IDENTIFIER { + if signing_algorithm == DLFOpenApiV4Signer::IDENTIFIER { + Box::new(DLFOpenApiV4Signer::new(region)) + } else if signing_algorithm == DLFOpenApiSigner::IDENTIFIER { Box::new(DLFOpenApiSigner) } else { Box::new(DLFDefaultSigner::new(region)) diff --git a/crates/paimon/src/api/auth/mod.rs b/crates/paimon/src/api/auth/mod.rs index 816c36236..23e3ccd6b 100644 --- a/crates/paimon/src/api/auth/mod.rs +++ b/crates/paimon/src/api/auth/mod.rs @@ -19,6 +19,8 @@ mod base; mod bearer_provider; +mod dlf_openapi_actions; +mod dlf_openapi_v4_signer; mod dlf_provider; mod dlf_signer; mod factory; diff --git a/crates/paimon/src/api/rest_api.rs b/crates/paimon/src/api/rest_api.rs index a69018b22..45ec5a861 100644 --- a/crates/paimon/src/api/rest_api.rs +++ b/crates/paimon/src/api/rest_api.rs @@ -558,8 +558,7 @@ impl RESTApi { branch: branch_name.to_string(), from_tag: from_tag.map(str::to_string), }; - let _: serde_json::Value = self.client.post(&path, &request).await?; - Ok(()) + self.client.post_without_response(&path, &request).await } /// Drop branch metadata through the catalog's publication protocol. @@ -572,8 +571,7 @@ impl RESTApi { let path = self.resource_paths .branch(identifier.database(), identifier.object(), branch_name); - let _: serde_json::Value = self.client.delete(&path, None::<&[(&str, &str)]>).await?; - Ok(()) + self.client.delete_without_response(&path).await } /// Ask the server to fast-forward from a branch. No files are changed locally. @@ -588,8 +586,9 @@ impl RESTApi { identifier.object(), branch_name, ); - let _: serde_json::Value = self.client.post(&path, &ForwardBranchRequest {}).await?; - Ok(()) + self.client + .post_without_response(&path, &ForwardBranchRequest {}) + .await } /// List the server's branches, preserving its order and null-list behavior. diff --git a/crates/paimon/src/api/rest_client.rs b/crates/paimon/src/api/rest_client.rs index 5448d47cf..236b12910 100644 --- a/crates/paimon/src/api/rest_client.rs +++ b/crates/paimon/src/api/rest_client.rs @@ -22,7 +22,7 @@ use super::rest_error::RestError; use super::user_agent; use crate::Error; use crate::Result; -use reqwest::header::HeaderValue; +use reqwest::header::{HeaderMap, HeaderValue, RETRY_AFTER}; use serde::de::DeserializeOwned; use std::collections::HashMap; use std::time::Duration; @@ -146,12 +146,61 @@ impl HttpClient { } } - let request = Self::apply_headers(request, &headers); - let resp = request.send().await.map_err(|e| Error::UnexpectedError { - message: "http get failed".to_string(), - source: Some(Box::new(e)), - })?; - self.parse_response(resp).await + let request = Self::apply_headers(request, &headers)? + .build() + .map_err(|e| Error::UnexpectedError { + message: "http get failed".to_string(), + source: Some(Box::new(e)), + })?; + // Retry only idempotent reads and only on Java's retryable HTTP statuses. + // A mutation may have been published before an intermediary returned an + // error, so POST/DELETE must never enter this loop. + for attempt in 0..=5 { + let pending = request + .try_clone() + .expect("GET requests have no streaming body"); + let response = + self.client + .execute(pending) + .await + .map_err(|e| Error::UnexpectedError { + message: "http get failed".to_string(), + source: Some(Box::new(e)), + })?; + if attempt == 5 || !matches!(response.status().as_u16(), 429 | 503) { + return self.parse_response(response).await; + } + let retry_after = response + .headers() + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()); + let delay = Self::retry_delay(retry_after, attempt, chrono::Utc::now()); + drop(response); + tokio::time::sleep(delay).await; + } + unreachable!("the final attempt returns its response") + } + + fn retry_delay( + retry_after: Option<&str>, + attempt: u32, + now: chrono::DateTime, + ) -> Duration { + if let Some(value) = retry_after { + let delay = value + .parse::() + .ok() + .map(Duration::from_secs) + .or_else(|| { + chrono::DateTime::parse_from_rfc2822(value) + .ok() + .and_then(|date| date.signed_duration_since(now).to_std().ok()) + }); + if let Some(delay) = delay.filter(|delay| !delay.is_zero()) { + return delay; + } + } + Duration::from_secs(1 << attempt.min(6)) } /// Perform a POST request with a JSON body. @@ -167,18 +216,35 @@ impl HttpClient { path: &str, body: &B, ) -> Result { + self.parse_response(self.send_post(path, body).await?).await + } + + /// Send a mutation whose successful response has no defined payload. + pub(crate) async fn post_without_response( + &self, + path: &str, + body: &B, + ) -> Result<()> { + Self::response_text(self.send_post(path, body).await?).await?; + Ok(()) + } + + async fn send_post( + &self, + path: &str, + body: &B, + ) -> Result { let url = self.request_url(path); let body_str = serde_json::to_string(body).ok(); let headers = self .build_auth_headers("POST", path, body_str.as_deref(), HashMap::new()) .await?; let request = self.client.post(&url).json(body); - let request = Self::apply_headers(request, &headers); - let resp = request.send().await.map_err(|e| Error::UnexpectedError { + let request = Self::apply_headers(request, &headers)?; + request.send().await.map_err(|e| Error::UnexpectedError { message: "http post failed".to_string(), source: Some(Box::new(e)), - })?; - self.parse_response(resp).await + }) } /// Perform a DELETE request with optional query parameters. @@ -194,6 +260,21 @@ impl HttpClient { path: &str, params: Option<&[(impl AsRef, impl AsRef)]>, ) -> Result { + self.parse_response(self.send_delete(path, params).await?) + .await + } + + /// Delete metadata, ignoring the successful payload as Java does. + pub(crate) async fn delete_without_response(&self, path: &str) -> Result<()> { + Self::response_text(self.send_delete(path, None::<&[(&str, &str)]>).await?).await?; + Ok(()) + } + + async fn send_delete( + &self, + path: &str, + params: Option<&[(impl AsRef, impl AsRef)]>, + ) -> Result { let url = self.request_url(path); let params_map: HashMap = match params { @@ -215,12 +296,11 @@ impl HttpClient { } } - let request = Self::apply_headers(request, &headers); - let resp = request.send().await.map_err(|e| Error::UnexpectedError { + let request = Self::apply_headers(request, &headers)?; + request.send().await.map_err(|e| Error::UnexpectedError { message: "http delete failed".to_string(), source: Some(Box::new(e)), - })?; - self.parse_response(resp).await + }) } /// Set the authentication function for this client. @@ -249,12 +329,13 @@ impl HttpClient { fn apply_headers( request: reqwest::RequestBuilder, headers: &HashMap, - ) -> reqwest::RequestBuilder { - let mut request = request; - for (key, value) in headers { - request = request.header(key, value); - } - request + ) -> Result { + let headers = HeaderMap::try_from(headers).map_err(|error| Error::UnexpectedError { + message: "Invalid REST request headers".to_string(), + source: Some(Box::new(error)), + })?; + // Replace JSON's default Content-Type rather than append a second value. + Ok(request.headers(headers)) } fn request_url(&self, path: &str) -> String { @@ -267,27 +348,24 @@ impl HttpClient { } } - async fn parse_response(&self, resp: reqwest::Response) -> Result { + async fn response_text(resp: reqwest::Response) -> Result { let status = resp.status(); - + let text = resp.text().await.map_err(|e| Error::UnexpectedError { + message: "failed to read response".to_string(), + source: Some(Box::new(e)), + })?; if !status.is_success() { - let text = resp.text().await.map_err(|e| Error::UnexpectedError { - message: "failed to read response".to_string(), - source: Some(Box::new(e)), - })?; - // Parse error response as ErrorResponse and map code to corresponding error let error_response: super::ErrorResponse = RestError::parse_error_response(&text, status.as_u16()); let rest_error: RestError = RestError::from_error_response(error_response); return Err(Error::from(rest_error)); } + Ok(text) + } - // Parse successful response - let text = resp.text().await.map_err(|e| Error::UnexpectedError { - message: "failed to read response".to_string(), - source: Some(Box::new(e)), - })?; + async fn parse_response(&self, resp: reqwest::Response) -> Result { + let text = Self::response_text(resp).await?; // Handle empty response body - return null as default for types like serde_json::Value if text.trim().is_empty() { @@ -317,6 +395,39 @@ mod tests { use super::*; use crate::api::auth::AuthProvider; + #[test] + fn retry_after_or_exponential_backoff_matches_java() { + let now = chrono::DateTime::parse_from_rfc3339("2025-04-16T03:44:46Z") + .unwrap() + .with_timezone(&chrono::Utc); + for (header, attempt, seconds) in [ + (Some("5"), 0, 5), + (Some("Wed, 16 Apr 2025 03:44:49 GMT"), 0, 3), + (Some("Wed, 16 Apr 2025 03:44:45 GMT"), 2, 4), + (None, 0, 1), + (Some("invalid"), 3, 8), + (Some("0"), 1, 2), + ] { + assert_eq!( + HttpClient::retry_delay(header, attempt, now), + Duration::from_secs(seconds) + ); + } + } + + #[tokio::test] + async fn invalid_get_headers_fail_before_sending() { + let auth = RESTAuthFunction::new( + HashMap::from([("X-Invalid".to_string(), "line\nbreak".to_string())]), + Box::new(RecordingProvider(Arc::new(Mutex::new(None)))), + ); + let client = HttpClient::new("http://localhost:0", Some(auth)).unwrap(); + assert!( + matches!(client.get::("/", None::<&[(&str, &str)]>).await, + Err(Error::UnexpectedError { message, .. }) if message == "Invalid REST request headers") + ); + } + struct RecordingProvider(Arc>>>); #[async_trait] diff --git a/crates/paimon/src/common/options.rs b/crates/paimon/src/common/options.rs index 3eed389ea..669b6938c 100644 --- a/crates/paimon/src/common/options.rs +++ b/crates/paimon/src/common/options.rs @@ -58,7 +58,7 @@ impl CatalogOptions { /// DLF security token (optional, for temporary credentials). pub const DLF_ACCESS_SECURITY_TOKEN: &'static str = "dlf.security-token"; - /// DLF signing algorithm (default or openapi). + /// DLF signing algorithm (default, openapi or openapi-v4). pub const DLF_SIGNING_ALGORITHM: &'static str = "dlf.signing-algorithm"; /// DLF token loader type ("ecs" or "local_file"). diff --git a/crates/paimon/tests/rest_branch_test.rs b/crates/paimon/tests/rest_branch_test.rs index 88eb9174a..1b5a4f005 100644 --- a/crates/paimon/tests/rest_branch_test.rs +++ b/crates/paimon/tests/rest_branch_test.rs @@ -22,9 +22,9 @@ use std::sync::{Arc, Mutex}; use axum::body::Bytes; use axum::extract::State; -use axum::http::{HeaderMap, Method, StatusCode, Uri}; +use axum::http::{header, HeaderMap, Method, StatusCode, Uri}; use axum::response::{IntoResponse, Response}; -use axum::{Json, Router}; +use axum::Router; use paimon::api::RestError; use paimon::catalog::{Catalog, Identifier, RESTCatalog}; use paimon::common::Options; @@ -40,7 +40,7 @@ struct Request { #[derive(Default)] struct ServerState { - responses: VecDeque<(StatusCode, Value)>, + responses: VecDeque<(StatusCode, String)>, requests: Vec, } @@ -58,6 +58,16 @@ impl Drop for TestCatalog { impl TestCatalog { async fn new(responses: Vec<(StatusCode, Value)>) -> Self { + Self::new_raw( + responses + .into_iter() + .map(|(status, body)| (status, body.to_string())) + .collect(), + ) + .await + } + + async fn new_raw(responses: Vec<(StatusCode, String)>) -> Self { async fn handle( State(state): State>>, method: Method, @@ -75,8 +85,10 @@ impl TestCatalog { let (status, value) = state.responses.pop_front().expect("unexpected request"); if status == StatusCode::NO_CONTENT { status.into_response() + } else if matches!(status.as_u16(), 429 | 503) { + (status, [(header::RETRY_AFTER, "1")], value).into_response() } else { - (status, Json(value)).into_response() + (status, value).into_response() } } let state = Arc::new(Mutex::new(ServerState { @@ -111,6 +123,103 @@ impl TestCatalog { } } +#[tokio::test] +async fn branch_mutations_ignore_success_response_bodies() { + let bodies = ["OK", "", "not JSON", "{}"]; + let ctx = TestCatalog::new_raw( + bodies + .into_iter() + .flat_map(|body| vec![(StatusCode::OK, body.to_string()); 3]) + .collect(), + ) + .await; + let id = Identifier::new("db", "t"); + for _ in bodies { + ctx.catalog.create_branch(&id, "dev", None).await.unwrap(); + ctx.catalog.drop_branch(&id, "dev").await.unwrap(); + ctx.catalog.fast_forward(&id, "dev").await.unwrap(); + } + assert_eq!(ctx.requests().len(), bodies.len() * 3); +} + +#[tokio::test] +async fn branch_listing_retries_only_retryable_statuses() { + for status in [ + StatusCode::TOO_MANY_REQUESTS, + StatusCode::SERVICE_UNAVAILABLE, + ] { + let ctx = TestCatalog::new(vec![ + ( + status, + json!({"code": status.as_u16(), "message": "try again"}), + ), + (StatusCode::OK, json!({"branches": ["dev"]})), + ]) + .await; + assert_eq!( + ctx.catalog + .list_branches(&Identifier::new("db", "t")) + .await + .unwrap(), + ["dev"] + ); + assert_eq!(ctx.requests().len(), 2); + } + let ctx = TestCatalog::new(vec![( + StatusCode::BAD_GATEWAY, + json!({"code": 502, "message": "gateway failed"}), + )]) + .await; + assert!(matches!( + ctx.catalog.list_branches(&Identifier::new("db", "t")).await, + Err(Error::RestApi { + source: RestError::Unexpected { code: 502, .. } + }) + )); + assert_eq!(ctx.requests().len(), 1); +} + +#[tokio::test] +async fn branch_mutations_do_not_retry_temporary_failures() { + let id = Identifier::new("db", "t"); + for status in [ + StatusCode::TOO_MANY_REQUESTS, + StatusCode::SERVICE_UNAVAILABLE, + ] { + let ctx = TestCatalog::new(vec![ + ( + status, + json!({"code": status.as_u16(), "message": "uncertain result"}) + ); + 3 + ]) + .await; + assert!(ctx.catalog.create_branch(&id, "dev", None).await.is_err()); + assert!(ctx.catalog.drop_branch(&id, "dev").await.is_err()); + assert!(ctx.catalog.fast_forward(&id, "dev").await.is_err()); + assert_eq!(ctx.requests().len(), 3); + } +} + +#[tokio::test] +async fn branch_listing_stops_after_five_retries() { + let ctx = TestCatalog::new(vec![ + ( + StatusCode::SERVICE_UNAVAILABLE, + json!({"code": 503, "message": "unavailable"}) + ); + 6 + ]) + .await; + assert!(matches!( + ctx.catalog.list_branches(&Identifier::new("db", "t")).await, + Err(Error::RestApi { + source: RestError::ServiceUnavailable { .. } + }) + )); + assert_eq!(ctx.requests().len(), 6); +} + #[tokio::test] async fn branch_protocol_preserves_names_tags_and_server_order() { let ctx = TestCatalog::new(vec![ From c16a89febca8535d07918236fb40978d887f5904 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 17 Oct 2026 11:02:09 +0800 Subject: [PATCH 4/5] [core] Preserve REST cookie sessions and Java request timeouts --- Cargo.lock | 62 +++++++++ .../python/pypaimon_rust/datafusion.pyi | 6 +- bindings/python/src/context.rs | 46 +++++-- .../python/tests/test_catalog_branches.py | 25 ++++ crates/paimon/Cargo.toml | 2 +- crates/paimon/src/api/mod.rs | 3 + crates/paimon/src/api/rest_api.rs | 4 + crates/paimon/src/api/rest_client.rs | 57 ++++++++- crates/paimon/src/api/rest_cookies.rs | 119 ++++++++++++++++++ .../paimon/src/catalog/rest/rest_catalog.rs | 7 ++ 10 files changed, 312 insertions(+), 19 deletions(-) create mode 100644 crates/paimon/src/api/rest_cookies.rs diff --git a/Cargo.lock b/Cargo.lock index 61d3be37f..534b51ea5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1309,6 +1309,35 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "cookie" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + +[[package]] +name = "cookie_store" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" +dependencies = [ + "cookie", + "document-features", + "idna", + "log", + "publicsuffix", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -2419,6 +2448,15 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + [[package]] name = "downcast-rs" version = "2.0.2" @@ -3835,6 +3873,12 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "lock_api" version = "0.4.14" @@ -5310,6 +5354,12 @@ dependencies = [ "prost 0.14.4", ] +[[package]] +name = "psl-types" +version = "2.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" + [[package]] name = "psm" version = "0.1.31" @@ -5320,6 +5370,16 @@ dependencies = [ "cc", ] +[[package]] +name = "publicsuffix" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f42ea446cab60335f76979ec15e12619a2165b5ae2c12166bef27d283a9fadf" +dependencies = [ + "idna", + "psl-types", +] + [[package]] name = "pyo3" version = "0.28.3" @@ -5869,6 +5929,8 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", + "cookie", + "cookie_store", "encoding_rs", "futures-core", "h2", diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 44cee5906..9fa29dfc5 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -756,7 +756,8 @@ class CatalogError(ValueError): message: str class PaimonCatalog: - def __init__(self, catalog_options: Dict[str, str], *, config_required: bool = True) -> None: + def __init__(self, catalog_options: Dict[str, str], *, config_required: bool = True, + _capture_response_cookies: bool = False) -> None: """Set config_required=False when the REST configuration has already been resolved.""" ... def __datafusion_catalog_provider__(self, session: Any) -> object: ... @@ -772,6 +773,9 @@ class PaimonCatalog: """Replace the identifier's current branch metadata through the REST server.""" ... def list_branches(self, identifier: Union[str, Tuple[str, str]]) -> List[str]: ... + def _take_response_cookies(self) -> List[Tuple[str, List[str], int]]: + """Drain response cookie headers, origin URLs and receipt times (Unix seconds).""" + ... class PythonScalarUDF: def __init__( diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index d72233505..e653108eb 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -192,24 +192,28 @@ pub(crate) fn complete_storage_options( fn build_paimon_catalog( catalog_options: HashMap, config_required: bool, -) -> PyResult> { + capture_response_cookies: bool, +) -> PyResult<(Arc, Option)> { let catalog_options = complete_storage_options(catalog_options); let rt = runtime(); rt.block_on(async { let options = Options::from_map(catalog_options); - let catalog: Arc = - if !config_required && options.get("metastore").map(String::as_str) == Some("rest") { - Arc::new( - paimon::catalog::RESTCatalog::new(options, false) + let (catalog, cookies): (Arc, _) = + if options.get("metastore").map(String::as_str) == Some("rest") { + let catalog = paimon::catalog::RESTCatalog::new(options, config_required) + .await + .map_err(to_catalog_py_err)?; + let cookies = capture_response_cookies.then(|| catalog.observe_response_cookies()); + (Arc::new(catalog), cookies) + } else { + ( + CatalogFactory::create(options) .await .map_err(to_catalog_py_err)?, + None, ) - } else { - CatalogFactory::create(options) - .await - .map_err(to_catalog_py_err)? }; - Ok::<_, PyErr>(catalog) + Ok::<_, PyErr>((catalog, cookies)) }) } @@ -290,6 +294,7 @@ fn ffi_logical_codec_from_pycapsule(obj: Bound<'_, PyAny>) -> PyResult, provider: Arc, + response_cookies: Option, } fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult { @@ -326,13 +331,16 @@ fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult, catalog_options: HashMap, config_required: bool, + _capture_response_cookies: bool, ) -> PyResult { - let catalog = py.detach(|| build_paimon_catalog(catalog_options, config_required))?; + let (catalog, response_cookies) = py.detach(|| { + build_paimon_catalog(catalog_options, config_required, _capture_response_cookies) + })?; let provider = Arc::new( PaimonCatalogProvider::new( None, @@ -343,7 +351,19 @@ impl PaimonCatalog { ) .with_schema_force_view_types(false), ); - Ok(Self { catalog, provider }) + Ok(Self { + catalog, + provider, + response_cookies, + }) + } + + /// Return response cookie updates to the owning PyPaimon REST session. + fn _take_response_cookies(&self) -> Vec<(String, Vec, i64)> { + self.response_cookies + .as_ref() + .map(|cookies| cookies.take()) + .unwrap_or_default() } /// Export this catalog as a DataFusion catalog provider PyCapsule. diff --git a/bindings/python/tests/test_catalog_branches.py b/bindings/python/tests/test_catalog_branches.py index a5cef1e9c..e9582c508 100644 --- a/bindings/python/tests/test_catalog_branches.py +++ b/bindings/python/tests/test_catalog_branches.py @@ -19,6 +19,7 @@ import json import threading +import time from http.server import BaseHTTPRequestHandler, HTTPServer from urllib.parse import unquote_plus @@ -65,6 +66,10 @@ def handle_operation(self): self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(output))) + if "cookies" in parts: + if len(requests) > 1: + assert self.headers["Cookie"] == "sid=accepted" + self.send_header("Set-Cookie", "sid=accepted; Path=/v1/") self.end_headers() self.wfile.write(output) @@ -104,6 +109,26 @@ def test_resolved_rest_configuration_is_not_fetched_again(rest_server): assert requests == [("GET", "/v1/resolved-prefix/databases/db/tables/t/branches", None)] +@pytest.mark.parametrize("capture", [False, True]) +def test_response_cookie_capture_is_opt_in_but_jar_always_persists(rest_server, capture): + options, requests = rest_server + native = PaimonCatalog(options, config_required=False, _capture_response_cookies=capture) + for _ in range(2): + assert native.list_branches("db.cookies") == ["dev", "release"] + updates = native._take_response_cookies() + if capture: + assert len(updates) == 1 + assert updates[0][:2] == ( + options["uri"] + "/v1/databases/db/tables/cookies/branches", + ["sid=accepted; Path=/v1/"], + ) + assert abs(time.time() - updates[0][2]) < 2 + else: + assert updates == [] + assert native._take_response_cookies() == [] + assert len(requests) == 2 + + def test_rest_branch_operations_use_core_protocol(catalog): native, requests = catalog identifier = ("db.name", "table.name$branch_dev") diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index 31302b4a8..1876ed28c 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -135,7 +135,7 @@ tokio-util = { workspace = true, features = ["compat", "io-util"] } parquet = { workspace = true, features = ["async", "zstd", "lz4", "snap"] } orc-rust = "0.8.0" async-stream = "0.3.6" -reqwest = { version = "0.12", features = ["json", "gzip", "deflate"] } +reqwest = { version = "0.12", features = ["json", "gzip", "deflate", "cookies"] } # Already in the tree via `reqwest`; direct dep for the object storage User-Agent transport. http = "1" # DLF authentication dependencies diff --git a/crates/paimon/src/api/mod.rs b/crates/paimon/src/api/mod.rs index 55527ae01..169e3d2e0 100644 --- a/crates/paimon/src/api/mod.rs +++ b/crates/paimon/src/api/mod.rs @@ -29,8 +29,11 @@ pub mod rest_error; pub mod rest_util; mod api_response; +mod rest_cookies; mod user_agent; +pub use rest_cookies::ResponseCookies; + // Re-export request types pub use api_request::{ AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateBranchRequest, diff --git a/crates/paimon/src/api/rest_api.rs b/crates/paimon/src/api/rest_api.rs index 45ec5a861..335d30c5b 100644 --- a/crates/paimon/src/api/rest_api.rs +++ b/crates/paimon/src/api/rest_api.rs @@ -220,6 +220,10 @@ impl RESTApi { &self.options } + pub(crate) fn observe_response_cookies(&self) -> super::ResponseCookies { + self.client.observe_response_cookies() + } + // ==================== Database Operations ==================== /// List all databases. diff --git a/crates/paimon/src/api/rest_client.rs b/crates/paimon/src/api/rest_client.rs index 236b12910..67c02ac91 100644 --- a/crates/paimon/src/api/rest_client.rs +++ b/crates/paimon/src/api/rest_client.rs @@ -18,13 +18,16 @@ //! Asynchronous HTTP client for REST API calls. use super::auth::{RESTAuthFunction, RESTAuthParameter}; +use super::rest_cookies::CookieStore; use super::rest_error::RestError; use super::user_agent; +use super::ResponseCookies; use crate::Error; use crate::Result; use reqwest::header::{HeaderMap, HeaderValue, RETRY_AFTER}; use serde::de::DeserializeOwned; use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; /// Asynchronous HTTP client for REST API calls. @@ -32,6 +35,7 @@ pub struct HttpClient { client: reqwest::Client, base_url: String, auth_function: Option, + cookies: Arc, } impl HttpClient { @@ -57,27 +61,37 @@ impl HttpClient { auth_function: Option, agent: &str, ) -> Result { + let cookies = Arc::new(CookieStore::default()); Ok(HttpClient { - client: Self::build_client(agent)?, + client: Self::build_client(agent, Arc::clone(&cookies))?, base_url: Self::normalize_uri(base_url)?, auth_function, + cookies, }) } /// Replace the default User-Agent, e.g. after options are merged with the server config. pub(crate) fn set_user_agent(&mut self, agent: &str) -> Result<()> { - self.client = Self::build_client(agent)?; + self.client = Self::build_client(agent, Arc::clone(&self.cookies))?; Ok(()) } - fn build_client(agent: &str) -> Result { + pub(crate) fn observe_response_cookies(&self) -> ResponseCookies { + self.cookies.observe() + } + + fn build_client(agent: &str, cookies: Arc) -> Result { let agent = HeaderValue::from_str(agent).unwrap_or_else(|_| { log::warn!("Invalid REST User-Agent {agent:?}, using the default"); HeaderValue::from_str(&user_agent::default_rest_user_agent()) .expect("the default User-Agent is visible ASCII") }); reqwest::Client::builder() - .timeout(Duration::from_secs(30)) + // Java and PyPaimon allow three minutes for connection and response + // reads. A short total deadline aborts otherwise valid REST calls. + .connect_timeout(Duration::from_secs(180)) + .read_timeout(Duration::from_secs(180)) + .cookie_provider(cookies) .user_agent(agent) .build() .map_err(|e| Error::ConfigInvalid { @@ -555,6 +569,41 @@ mod tests { ); } + #[tokio::test] + async fn config_cookies_survive_user_agent_reconfiguration() { + let app = Router::new() + .route( + "/v1/config", + get(|| async { + ( + [(header::SET_COOKIE, "sid=accepted; Path=/v1/")], + Json(serde_json::json!({})), + ) + }), + ) + .route( + "/v1/databases", + get(|headers: HeaderMap| async move { + assert_eq!(headers.get(header::COOKIE).unwrap(), "sid=accepted"); + Json(serde_json::json!({"databases": ["db"]})) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let options = crate::common::Options::from_map(HashMap::from([ + ("uri".to_string(), format!("http://{address}")), + ("warehouse".to_string(), "warehouse".to_string()), + ("token.provider".to_string(), "bear".to_string()), + ("token".to_string(), "token".to_string()), + ])); + let api = crate::api::rest_api::RESTApi::new(options, true) + .await + .unwrap(); + assert_eq!(api.list_databases().await.unwrap(), vec!["db"]); + server.abort(); + } + #[tokio::test] async fn test_user_agent_header_option_wins_over_server_config() { let config = serde_json::json!({ diff --git a/crates/paimon/src/api/rest_cookies.rs b/crates/paimon/src/api/rest_cookies.rs new file mode 100644 index 000000000..0eda515f5 --- /dev/null +++ b/crates/paimon/src/api/rest_cookies.rs @@ -0,0 +1,119 @@ +// 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. + +//! Cookie persistence and response observation for REST transport interoperability. + +use std::sync::{Arc, Mutex, Weak}; + +use reqwest::cookie::Jar; +use reqwest::header::HeaderValue; +use reqwest::Url; + +type Updates = Vec<(String, Vec, i64)>; + +/// Observes response Set-Cookie headers, their origin URLs and receipt times. +/// +/// A language binding can feed these into its owning HTTP session's cookie jar, +/// preserving domain, path, expiry and deletion semantics across transports. +/// Dropping the observer stops collecting updates without retaining credentials. +#[derive(Clone, Default)] +pub struct ResponseCookies(Arc>); + +impl ResponseCookies { + /// Drain updates as (URL, Set-Cookie headers, Unix seconds received), in + /// response order. Receipt time preserves Max-Age across delayed handoff. + pub fn take(&self) -> Vec<(String, Vec, i64)> { + std::mem::take(&mut *self.0.lock().unwrap()) + } +} + +/// The HTTP client owns its jar; an optional, weak observer copies wire updates. +#[derive(Default)] +pub(super) struct CookieStore { + jar: Jar, + observer: Mutex>>, +} + +impl CookieStore { + pub(super) fn observe(&self) -> ResponseCookies { + let observer = ResponseCookies::default(); + *self.observer.lock().unwrap() = Arc::downgrade(&observer.0); + observer + } +} + +impl reqwest::cookie::CookieStore for CookieStore { + fn set_cookies(&self, cookies: &mut dyn Iterator, url: &Url) { + let received_at = chrono::Utc::now().timestamp(); + let cookies: Vec<_> = cookies.collect(); + self.jar.set_cookies(&mut cookies.iter().copied(), url); + if let Some(observer) = self.observer.lock().unwrap().upgrade() { + let values = cookies + .into_iter() + .filter_map(|value| value.to_str().ok().map(str::to_string)) + .collect(); + observer + .lock() + .unwrap() + .push((url.to_string(), values, received_at)); + } + } + + fn cookies(&self, url: &Url) -> Option { + self.jar.cookies(url) + } +} + +#[cfg(test)] +mod tests { + use reqwest::cookie::CookieStore as _; + + use super::*; + + #[test] + fn cookies_keep_scope_and_observers_drain_without_clearing_the_jar() { + let store = CookieStore::default(); + let url = Url::parse("https://catalog.example/v1/branches").unwrap(); + let observer = store.observe(); + let cookie = HeaderValue::from_static("sid=accepted; Path=/v1/; Secure; HttpOnly"); + store.set_cookies(&mut [&cookie].into_iter(), &url); + assert_eq!(store.cookies(&url).unwrap(), "sid=accepted"); + assert!(store + .cookies(&Url::parse("http://catalog.example/v1/").unwrap()) + .is_none()); + assert!(store + .cookies(&Url::parse("https://catalog.example/other").unwrap()) + .is_none()); + assert!(store + .cookies(&Url::parse("https://other.example/v1/").unwrap()) + .is_none()); + let updates = observer.take(); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].0, url.to_string()); + assert_eq!(updates[0].1, vec![cookie.to_str().unwrap()]); + assert!((chrono::Utc::now().timestamp() - updates[0].2).abs() <= 1); + assert!(observer.take().is_empty()); + assert_eq!(store.cookies(&url).unwrap(), "sid=accepted"); + + let remove = HeaderValue::from_static("sid=; Path=/v1/; Max-Age=0"); + store.set_cookies(&mut [&remove].into_iter(), &url); + assert!(store.cookies(&url).is_none()); + assert_eq!(observer.take()[0].1, vec![remove.to_str().unwrap()]); + drop(observer); + assert!(store.observer.lock().unwrap().upgrade().is_none()); + } +} diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 3f4f623d2..0fa8b1048 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -112,6 +112,13 @@ impl RESTCatalog { &self.options } + /// Observe subsequent response cookies for integration with an owning HTTP + /// session. A new observer replaces the previous one; dropping it stops + /// collecting updates. The catalog's own cookie jar is unaffected. + pub fn observe_response_cookies(&self) -> crate::api::ResponseCookies { + self.api.observe_response_cookies() + } + /// Whether data token is enabled. pub fn data_token_enabled(&self) -> bool { self.data_token_enabled From 345f5ccd2febb2b49127ab22e0e83b36d6761b04 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sat, 17 Oct 2026 11:27:22 +0800 Subject: [PATCH 5/5] [core] Remove REST cookie session synchronization --- Cargo.lock | 62 --------- .../python/pypaimon_rust/datafusion.pyi | 6 +- bindings/python/src/context.rs | 46 ++----- .../python/tests/test_catalog_branches.py | 25 ---- crates/paimon/Cargo.toml | 2 +- crates/paimon/src/api/mod.rs | 3 - crates/paimon/src/api/rest_api.rs | 4 - crates/paimon/src/api/rest_client.rs | 52 +------- crates/paimon/src/api/rest_cookies.rs | 119 ------------------ .../paimon/src/catalog/rest/rest_catalog.rs | 7 -- 10 files changed, 18 insertions(+), 308 deletions(-) delete mode 100644 crates/paimon/src/api/rest_cookies.rs diff --git a/Cargo.lock b/Cargo.lock index 534b51ea5..61d3be37f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1309,35 +1309,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" -[[package]] -name = "cookie" -version = "0.18.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87" -dependencies = [ - "percent-encoding", - "time", - "version_check", -] - -[[package]] -name = "cookie_store" -version = "0.22.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" -dependencies = [ - "cookie", - "document-features", - "idna", - "log", - "publicsuffix", - "serde", - "serde_derive", - "serde_json", - "time", - "url", -] - [[package]] name = "core-foundation" version = "0.9.4" @@ -2448,15 +2419,6 @@ dependencies = [ "windows-sys 0.60.2", ] -[[package]] -name = "document-features" -version = "0.2.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" -dependencies = [ - "litrs", -] - [[package]] name = "downcast-rs" version = "2.0.2" @@ -3873,12 +3835,6 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" -[[package]] -name = "litrs" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" - [[package]] name = "lock_api" version = "0.4.14" @@ -5354,12 +5310,6 @@ dependencies = [ "prost 0.14.4", ] -[[package]] -name = "psl-types" -version = "2.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" - [[package]] name = "psm" version = "0.1.31" @@ -5370,16 +5320,6 @@ dependencies = [ "cc", ] -[[package]] -name = "publicsuffix" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f42ea446cab60335f76979ec15e12619a2165b5ae2c12166bef27d283a9fadf" -dependencies = [ - "idna", - "psl-types", -] - [[package]] name = "pyo3" version = "0.28.3" @@ -5929,8 +5869,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", - "cookie", - "cookie_store", "encoding_rs", "futures-core", "h2", diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 9fa29dfc5..44cee5906 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -756,8 +756,7 @@ class CatalogError(ValueError): message: str class PaimonCatalog: - def __init__(self, catalog_options: Dict[str, str], *, config_required: bool = True, - _capture_response_cookies: bool = False) -> None: + def __init__(self, catalog_options: Dict[str, str], *, config_required: bool = True) -> None: """Set config_required=False when the REST configuration has already been resolved.""" ... def __datafusion_catalog_provider__(self, session: Any) -> object: ... @@ -773,9 +772,6 @@ class PaimonCatalog: """Replace the identifier's current branch metadata through the REST server.""" ... def list_branches(self, identifier: Union[str, Tuple[str, str]]) -> List[str]: ... - def _take_response_cookies(self) -> List[Tuple[str, List[str], int]]: - """Drain response cookie headers, origin URLs and receipt times (Unix seconds).""" - ... class PythonScalarUDF: def __init__( diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index e653108eb..d72233505 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -192,28 +192,24 @@ pub(crate) fn complete_storage_options( fn build_paimon_catalog( catalog_options: HashMap, config_required: bool, - capture_response_cookies: bool, -) -> PyResult<(Arc, Option)> { +) -> PyResult> { let catalog_options = complete_storage_options(catalog_options); let rt = runtime(); rt.block_on(async { let options = Options::from_map(catalog_options); - let (catalog, cookies): (Arc, _) = - if options.get("metastore").map(String::as_str) == Some("rest") { - let catalog = paimon::catalog::RESTCatalog::new(options, config_required) - .await - .map_err(to_catalog_py_err)?; - let cookies = capture_response_cookies.then(|| catalog.observe_response_cookies()); - (Arc::new(catalog), cookies) - } else { - ( - CatalogFactory::create(options) + let catalog: Arc = + if !config_required && options.get("metastore").map(String::as_str) == Some("rest") { + Arc::new( + paimon::catalog::RESTCatalog::new(options, false) .await .map_err(to_catalog_py_err)?, - None, ) + } else { + CatalogFactory::create(options) + .await + .map_err(to_catalog_py_err)? }; - Ok::<_, PyErr>((catalog, cookies)) + Ok::<_, PyErr>(catalog) }) } @@ -294,7 +290,6 @@ fn ffi_logical_codec_from_pycapsule(obj: Bound<'_, PyAny>) -> PyResult, provider: Arc, - response_cookies: Option, } fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult { @@ -331,16 +326,13 @@ fn catalog_table_identifier(identifier: &Bound<'_, PyAny>) -> PyResult, catalog_options: HashMap, config_required: bool, - _capture_response_cookies: bool, ) -> PyResult { - let (catalog, response_cookies) = py.detach(|| { - build_paimon_catalog(catalog_options, config_required, _capture_response_cookies) - })?; + let catalog = py.detach(|| build_paimon_catalog(catalog_options, config_required))?; let provider = Arc::new( PaimonCatalogProvider::new( None, @@ -351,19 +343,7 @@ impl PaimonCatalog { ) .with_schema_force_view_types(false), ); - Ok(Self { - catalog, - provider, - response_cookies, - }) - } - - /// Return response cookie updates to the owning PyPaimon REST session. - fn _take_response_cookies(&self) -> Vec<(String, Vec, i64)> { - self.response_cookies - .as_ref() - .map(|cookies| cookies.take()) - .unwrap_or_default() + Ok(Self { catalog, provider }) } /// Export this catalog as a DataFusion catalog provider PyCapsule. diff --git a/bindings/python/tests/test_catalog_branches.py b/bindings/python/tests/test_catalog_branches.py index e9582c508..a5cef1e9c 100644 --- a/bindings/python/tests/test_catalog_branches.py +++ b/bindings/python/tests/test_catalog_branches.py @@ -19,7 +19,6 @@ import json import threading -import time from http.server import BaseHTTPRequestHandler, HTTPServer from urllib.parse import unquote_plus @@ -66,10 +65,6 @@ def handle_operation(self): self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(output))) - if "cookies" in parts: - if len(requests) > 1: - assert self.headers["Cookie"] == "sid=accepted" - self.send_header("Set-Cookie", "sid=accepted; Path=/v1/") self.end_headers() self.wfile.write(output) @@ -109,26 +104,6 @@ def test_resolved_rest_configuration_is_not_fetched_again(rest_server): assert requests == [("GET", "/v1/resolved-prefix/databases/db/tables/t/branches", None)] -@pytest.mark.parametrize("capture", [False, True]) -def test_response_cookie_capture_is_opt_in_but_jar_always_persists(rest_server, capture): - options, requests = rest_server - native = PaimonCatalog(options, config_required=False, _capture_response_cookies=capture) - for _ in range(2): - assert native.list_branches("db.cookies") == ["dev", "release"] - updates = native._take_response_cookies() - if capture: - assert len(updates) == 1 - assert updates[0][:2] == ( - options["uri"] + "/v1/databases/db/tables/cookies/branches", - ["sid=accepted; Path=/v1/"], - ) - assert abs(time.time() - updates[0][2]) < 2 - else: - assert updates == [] - assert native._take_response_cookies() == [] - assert len(requests) == 2 - - def test_rest_branch_operations_use_core_protocol(catalog): native, requests = catalog identifier = ("db.name", "table.name$branch_dev") diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index 1876ed28c..31302b4a8 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -135,7 +135,7 @@ tokio-util = { workspace = true, features = ["compat", "io-util"] } parquet = { workspace = true, features = ["async", "zstd", "lz4", "snap"] } orc-rust = "0.8.0" async-stream = "0.3.6" -reqwest = { version = "0.12", features = ["json", "gzip", "deflate", "cookies"] } +reqwest = { version = "0.12", features = ["json", "gzip", "deflate"] } # Already in the tree via `reqwest`; direct dep for the object storage User-Agent transport. http = "1" # DLF authentication dependencies diff --git a/crates/paimon/src/api/mod.rs b/crates/paimon/src/api/mod.rs index 169e3d2e0..55527ae01 100644 --- a/crates/paimon/src/api/mod.rs +++ b/crates/paimon/src/api/mod.rs @@ -29,11 +29,8 @@ pub mod rest_error; pub mod rest_util; mod api_response; -mod rest_cookies; mod user_agent; -pub use rest_cookies::ResponseCookies; - // Re-export request types pub use api_request::{ AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateBranchRequest, diff --git a/crates/paimon/src/api/rest_api.rs b/crates/paimon/src/api/rest_api.rs index 335d30c5b..45ec5a861 100644 --- a/crates/paimon/src/api/rest_api.rs +++ b/crates/paimon/src/api/rest_api.rs @@ -220,10 +220,6 @@ impl RESTApi { &self.options } - pub(crate) fn observe_response_cookies(&self) -> super::ResponseCookies { - self.client.observe_response_cookies() - } - // ==================== Database Operations ==================== /// List all databases. diff --git a/crates/paimon/src/api/rest_client.rs b/crates/paimon/src/api/rest_client.rs index 67c02ac91..22f320c00 100644 --- a/crates/paimon/src/api/rest_client.rs +++ b/crates/paimon/src/api/rest_client.rs @@ -18,16 +18,13 @@ //! Asynchronous HTTP client for REST API calls. use super::auth::{RESTAuthFunction, RESTAuthParameter}; -use super::rest_cookies::CookieStore; use super::rest_error::RestError; use super::user_agent; -use super::ResponseCookies; use crate::Error; use crate::Result; use reqwest::header::{HeaderMap, HeaderValue, RETRY_AFTER}; use serde::de::DeserializeOwned; use std::collections::HashMap; -use std::sync::Arc; use std::time::Duration; /// Asynchronous HTTP client for REST API calls. @@ -35,7 +32,6 @@ pub struct HttpClient { client: reqwest::Client, base_url: String, auth_function: Option, - cookies: Arc, } impl HttpClient { @@ -61,26 +57,20 @@ impl HttpClient { auth_function: Option, agent: &str, ) -> Result { - let cookies = Arc::new(CookieStore::default()); Ok(HttpClient { - client: Self::build_client(agent, Arc::clone(&cookies))?, + client: Self::build_client(agent)?, base_url: Self::normalize_uri(base_url)?, auth_function, - cookies, }) } /// Replace the default User-Agent, e.g. after options are merged with the server config. pub(crate) fn set_user_agent(&mut self, agent: &str) -> Result<()> { - self.client = Self::build_client(agent, Arc::clone(&self.cookies))?; + self.client = Self::build_client(agent)?; Ok(()) } - pub(crate) fn observe_response_cookies(&self) -> ResponseCookies { - self.cookies.observe() - } - - fn build_client(agent: &str, cookies: Arc) -> Result { + fn build_client(agent: &str) -> Result { let agent = HeaderValue::from_str(agent).unwrap_or_else(|_| { log::warn!("Invalid REST User-Agent {agent:?}, using the default"); HeaderValue::from_str(&user_agent::default_rest_user_agent()) @@ -91,7 +81,6 @@ impl HttpClient { // reads. A short total deadline aborts otherwise valid REST calls. .connect_timeout(Duration::from_secs(180)) .read_timeout(Duration::from_secs(180)) - .cookie_provider(cookies) .user_agent(agent) .build() .map_err(|e| Error::ConfigInvalid { @@ -569,41 +558,6 @@ mod tests { ); } - #[tokio::test] - async fn config_cookies_survive_user_agent_reconfiguration() { - let app = Router::new() - .route( - "/v1/config", - get(|| async { - ( - [(header::SET_COOKIE, "sid=accepted; Path=/v1/")], - Json(serde_json::json!({})), - ) - }), - ) - .route( - "/v1/databases", - get(|headers: HeaderMap| async move { - assert_eq!(headers.get(header::COOKIE).unwrap(), "sid=accepted"); - Json(serde_json::json!({"databases": ["db"]})) - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let options = crate::common::Options::from_map(HashMap::from([ - ("uri".to_string(), format!("http://{address}")), - ("warehouse".to_string(), "warehouse".to_string()), - ("token.provider".to_string(), "bear".to_string()), - ("token".to_string(), "token".to_string()), - ])); - let api = crate::api::rest_api::RESTApi::new(options, true) - .await - .unwrap(); - assert_eq!(api.list_databases().await.unwrap(), vec!["db"]); - server.abort(); - } - #[tokio::test] async fn test_user_agent_header_option_wins_over_server_config() { let config = serde_json::json!({ diff --git a/crates/paimon/src/api/rest_cookies.rs b/crates/paimon/src/api/rest_cookies.rs deleted file mode 100644 index 0eda515f5..000000000 --- a/crates/paimon/src/api/rest_cookies.rs +++ /dev/null @@ -1,119 +0,0 @@ -// 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. - -//! Cookie persistence and response observation for REST transport interoperability. - -use std::sync::{Arc, Mutex, Weak}; - -use reqwest::cookie::Jar; -use reqwest::header::HeaderValue; -use reqwest::Url; - -type Updates = Vec<(String, Vec, i64)>; - -/// Observes response Set-Cookie headers, their origin URLs and receipt times. -/// -/// A language binding can feed these into its owning HTTP session's cookie jar, -/// preserving domain, path, expiry and deletion semantics across transports. -/// Dropping the observer stops collecting updates without retaining credentials. -#[derive(Clone, Default)] -pub struct ResponseCookies(Arc>); - -impl ResponseCookies { - /// Drain updates as (URL, Set-Cookie headers, Unix seconds received), in - /// response order. Receipt time preserves Max-Age across delayed handoff. - pub fn take(&self) -> Vec<(String, Vec, i64)> { - std::mem::take(&mut *self.0.lock().unwrap()) - } -} - -/// The HTTP client owns its jar; an optional, weak observer copies wire updates. -#[derive(Default)] -pub(super) struct CookieStore { - jar: Jar, - observer: Mutex>>, -} - -impl CookieStore { - pub(super) fn observe(&self) -> ResponseCookies { - let observer = ResponseCookies::default(); - *self.observer.lock().unwrap() = Arc::downgrade(&observer.0); - observer - } -} - -impl reqwest::cookie::CookieStore for CookieStore { - fn set_cookies(&self, cookies: &mut dyn Iterator, url: &Url) { - let received_at = chrono::Utc::now().timestamp(); - let cookies: Vec<_> = cookies.collect(); - self.jar.set_cookies(&mut cookies.iter().copied(), url); - if let Some(observer) = self.observer.lock().unwrap().upgrade() { - let values = cookies - .into_iter() - .filter_map(|value| value.to_str().ok().map(str::to_string)) - .collect(); - observer - .lock() - .unwrap() - .push((url.to_string(), values, received_at)); - } - } - - fn cookies(&self, url: &Url) -> Option { - self.jar.cookies(url) - } -} - -#[cfg(test)] -mod tests { - use reqwest::cookie::CookieStore as _; - - use super::*; - - #[test] - fn cookies_keep_scope_and_observers_drain_without_clearing_the_jar() { - let store = CookieStore::default(); - let url = Url::parse("https://catalog.example/v1/branches").unwrap(); - let observer = store.observe(); - let cookie = HeaderValue::from_static("sid=accepted; Path=/v1/; Secure; HttpOnly"); - store.set_cookies(&mut [&cookie].into_iter(), &url); - assert_eq!(store.cookies(&url).unwrap(), "sid=accepted"); - assert!(store - .cookies(&Url::parse("http://catalog.example/v1/").unwrap()) - .is_none()); - assert!(store - .cookies(&Url::parse("https://catalog.example/other").unwrap()) - .is_none()); - assert!(store - .cookies(&Url::parse("https://other.example/v1/").unwrap()) - .is_none()); - let updates = observer.take(); - assert_eq!(updates.len(), 1); - assert_eq!(updates[0].0, url.to_string()); - assert_eq!(updates[0].1, vec![cookie.to_str().unwrap()]); - assert!((chrono::Utc::now().timestamp() - updates[0].2).abs() <= 1); - assert!(observer.take().is_empty()); - assert_eq!(store.cookies(&url).unwrap(), "sid=accepted"); - - let remove = HeaderValue::from_static("sid=; Path=/v1/; Max-Age=0"); - store.set_cookies(&mut [&remove].into_iter(), &url); - assert!(store.cookies(&url).is_none()); - assert_eq!(observer.take()[0].1, vec![remove.to_str().unwrap()]); - drop(observer); - assert!(store.observer.lock().unwrap().upgrade().is_none()); - } -} diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 0fa8b1048..3f4f623d2 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -112,13 +112,6 @@ impl RESTCatalog { &self.options } - /// Observe subsequent response cookies for integration with an owning HTTP - /// session. A new observer replaces the previous one; dropping it stops - /// collecting updates. The catalog's own cookie jar is unaffected. - pub fn observe_response_cookies(&self) -> crate::api::ResponseCookies { - self.api.observe_response_cookies() - } - /// Whether data token is enabled. pub fn data_token_enabled(&self) -> bool { self.data_token_enabled