diff --git a/Cargo.lock b/Cargo.lock index 6b3bc6796..e85d2df26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4752,6 +4752,7 @@ dependencies = [ "datafusion", "flate2", "futures", + "indexmap 2.14.0", "lexical-write-float", "log", "paimon", diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 1e285fd19..40f7cfab1 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -122,7 +122,7 @@ impl PaimonCatalog { fn new(py: Python<'_>, catalog_options: HashMap) -> PyResult { let catalog = py.detach(|| build_paimon_catalog(catalog_options))?; let provider = Arc::new( - PaimonCatalogProvider::new( + PaimonCatalogProvider::new_uninitialized( None, Arc::clone(&catalog), Default::default(), @@ -134,12 +134,22 @@ impl PaimonCatalog { Ok(Self { catalog, provider }) } + /// Refresh the metadata snapshot used by synchronous DataFusion callbacks. + fn refresh_metadata(&self, py: Python<'_>) -> PyResult<()> { + let provider = Arc::clone(&self.provider); + py.detach(|| runtime().block_on(provider.refresh_metadata())) + .map_err(df_to_py_err) + } + /// Export this catalog as a DataFusion catalog provider PyCapsule. fn __datafusion_catalog_provider__<'py>( &self, py: Python<'py>, session: Bound<'py, PyAny>, ) -> PyResult> { + let provider = Arc::clone(&self.provider); + py.detach(|| runtime().block_on(provider.initialize_metadata())) + .map_err(df_to_py_err)?; let name = cr"datafusion_catalog_provider".into(); let provider = Arc::clone(&self.provider) as Arc; let codec = ffi_logical_codec_from_pycapsule(session)?; diff --git a/bindings/python/tests/test_catalog_gil.py b/bindings/python/tests/test_catalog_gil.py index 4cb38b4ca..8a7cbbb1e 100644 --- a/bindings/python/tests/test_catalog_gil.py +++ b/bindings/python/tests/test_catalog_gil.py @@ -19,6 +19,7 @@ from http.server import BaseHTTPRequestHandler, HTTPServer import pytest +from datafusion import SessionContext from pypaimon_rust.datafusion import PaimonCatalog @@ -84,3 +85,23 @@ def test_rest_catalog_calls_release_gil(rest_server, monkeypatch): assert catalog.list_tables("db") == ["table"] with pytest.raises(ValueError, match="does not exist"): catalog.get_table("db.missing") + + +def test_rest_catalog_without_views_registers_datafusion_provider( + rest_server, monkeypatch +): + monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1") + monkeypatch.setenv("no_proxy", "localhost,127.0.0.1") + catalog = PaimonCatalog( + { + "metastore": "rest", + "uri": rest_server, + "warehouse": "warehouse", + "token.provider": "bear", + "token": "test-token", + } + ) + + SessionContext().register_catalog_provider("paimon", catalog) + with pytest.raises(ValueError, match="Resource not found"): + catalog.refresh_metadata() diff --git a/bindings/python/tests/test_datafusion.py b/bindings/python/tests/test_datafusion.py index cfaa1a0e7..db7ed6798 100644 --- a/bindings/python/tests/test_datafusion.py +++ b/bindings/python/tests/test_datafusion.py @@ -520,6 +520,37 @@ def test_query_simple_table_via_catalog_provider(): ] +def test_catalog_provider_initializes_information_schema_snapshot(): + with tempfile.TemporaryDirectory() as warehouse: + writer = SQLContext() + writer.register_catalog("paimon", {"warehouse": warehouse}) + writer.sql("CREATE TABLE paimon.default.users (id INT, name STRING)") + + catalog = PaimonCatalog({"warehouse": warehouse}) + ctx = SessionContext() + ctx.register_catalog_provider("paimon", catalog) + batches = ctx.sql( + "SELECT column_name FROM paimon.information_schema.columns " + "WHERE table_schema = 'default' AND table_name = 'users'" + ).collect() + + assert set(pa.Table.from_batches(batches)["column_name"].to_pylist()) == { + "id", + "name", + } + + writer.sql("CREATE TABLE paimon.default.orders (order_id BIGINT)") + catalog.refresh_metadata() + batches = ctx.sql( + "SELECT table_name FROM paimon.information_schema.tables " + "WHERE table_schema = 'default'" + ).collect() + assert set(pa.Table.from_batches(batches)["table_name"].to_pylist()) == { + "orders", + "users", + } + + def test_catalog_provider_returns_pyarrow_compatible_strings(): with tempfile.TemporaryDirectory() as warehouse: writer = SQLContext() diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index f7072fc36..76e062b0d 100644 --- a/crates/integrations/datafusion/Cargo.toml +++ b/crates/integrations/datafusion/Cargo.toml @@ -41,9 +41,10 @@ datafusion = { workspace = true } log = "0.4" paimon = { workspace = true } futures = "0.3" +indexmap = "2" serde = { version = "1", features = ["derive"] } serde_json = "1" -tokio = { workspace = true, features = ["rt", "time", "fs"] } +tokio = { workspace = true, features = ["rt", "time", "fs", "sync"] } lexical-write-float = "1.0.6" uuid = { version = "1", features = ["v4"] } diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 06f34c3c4..d88a2c305 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -17,16 +17,18 @@ //! Paimon catalog integration for DataFusion. -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt::Debug; -use std::sync::Arc; -use std::sync::RwLock; +use std::ops::Deref; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::Duration; use async_trait::async_trait; use datafusion::catalog::{CatalogProvider, MemorySchemaProvider, SchemaProvider}; use datafusion::common::{plan_datafusion_err, Column}; use datafusion::datasource::{TableProvider, TableType}; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::SessionState; use datafusion::logical_expr::{expr_fn::cast, Expr, LogicalPlan, LogicalPlanBuilder}; use datafusion::prelude::SessionContext; @@ -34,6 +36,8 @@ use datafusion::sql::planner::IdentNormalizer; use datafusion::sql::sqlparser::ast::{Ident, ObjectName, Query, Statement, Visit, Visitor}; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::parser::Parser; +use futures::{stream, StreamExt, TryStreamExt}; +use indexmap::IndexMap; use paimon::catalog::{Catalog, Identifier, View}; use paimon::spec::TableType as PaimonTableType; @@ -49,6 +53,512 @@ pub(crate) type SessionStateProvider = Arc Option + Se /// providers, so registrations stay visible to schemas obtained earlier. type TableEngines = Arc>>>; +#[derive(Clone, Debug, Default)] +struct CatalogMetadataSnapshot { + databases: IndexMap>, +} + +#[derive(Clone, Debug, Default)] +struct DatabaseMetadata { + objects: IndexMap, + system_table_capabilities: HashMap, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SystemTableCapability { + Paimon, + Unsupported, + Unknown, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ObjectResolution { + Paimon, + Object, + Routed, + Unavailable, +} + +#[derive(Debug)] +struct CatalogMetadataState { + snapshot: RwLock>, + next_generation: AtomicU64, + database_generations: RwLock>, + active_refreshes: Mutex>, + object_resolutions: RwLock>, +} + +impl Default for CatalogMetadataState { + fn default() -> Self { + Self { + snapshot: RwLock::new(Arc::new(CatalogMetadataSnapshot::default())), + next_generation: AtomicU64::new(0), + database_generations: RwLock::new(HashMap::new()), + active_refreshes: Mutex::new(BTreeSet::new()), + object_resolutions: RwLock::new(HashMap::new()), + } + } +} + +struct RefreshGeneration<'a> { + generation: u64, + state: &'a CatalogMetadataState, +} + +impl RefreshGeneration<'_> { + fn get(&self) -> u64 { + self.generation + } +} + +impl Drop for RefreshGeneration<'_> { + fn drop(&mut self) { + self.state.finish_refresh(self.generation); + } +} + +impl Deref for CatalogMetadataState { + type Target = RwLock>; + + fn deref(&self) -> &Self::Target { + &self.snapshot + } +} + +impl CatalogMetadataState { + fn next_generation(&self) -> u64 { + self.next_generation.fetch_add(1, Ordering::AcqRel) + 1 + } + + fn begin_refresh(&self) -> RefreshGeneration<'_> { + let generation = self.next_generation(); + self.active_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(generation); + RefreshGeneration { + generation, + state: self, + } + } + + fn finish_refresh(&self, generation: u64) { + let mut active_refreshes = self + .active_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()); + active_refreshes.remove(&generation); + let oldest_active = active_refreshes.first().copied(); + let current = self.snapshot.read().unwrap_or_else(|e| e.into_inner()); + let mut database_generations = self + .database_generations + .write() + .unwrap_or_else(|e| e.into_inner()); + Self::prune_database_generations( + current.as_ref(), + &mut database_generations, + oldest_active, + ); + } + + fn publish(&self, generation: u64, refreshed: CatalogMetadataSnapshot) -> HashSet { + let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); + let mut database_generations = self + .database_generations + .write() + .unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + let mut conflicts = HashSet::new(); + let mut published_databases = HashSet::new(); + let refreshed_names: HashSet<_> = refreshed.databases.keys().cloned().collect(); + + let removed_names: HashSet<_> = next + .databases + .keys() + .chain(database_generations.keys()) + .filter(|name| !refreshed_names.contains(*name)) + .cloned() + .collect(); + for name in removed_names { + if database_generations.get(&name).copied().unwrap_or_default() <= generation { + next.databases.shift_remove(&name); + database_generations.insert(name.clone(), generation); + published_databases.insert(name); + } else { + conflicts.insert(name); + } + } + for (name, metadata) in refreshed.databases { + if database_generations.get(&name).copied().unwrap_or_default() <= generation { + next.databases.insert(name.clone(), metadata); + database_generations.insert(name.clone(), generation); + published_databases.insert(name); + } else { + conflicts.insert(name); + } + } + *current = Arc::new(next); + self.retain_object_resolutions(current.as_ref()); + self.clear_unavailable_resolutions(&published_databases); + conflicts + } + + fn publish_database( + &self, + generation: u64, + database: String, + metadata: Arc, + ) -> bool { + let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); + let mut database_generations = self + .database_generations + .write() + .unwrap_or_else(|e| e.into_inner()); + if database_generations + .get(&database) + .copied() + .unwrap_or_default() + > generation + { + return false; + } + let mut next = (**current).clone(); + next.databases.insert(database.clone(), metadata); + *current = Arc::new(next); + self.retain_object_resolutions(current.as_ref()); + self.clear_unavailable_resolution_for_database(&database); + database_generations.insert(database, generation); + true + } + + fn mutate_database(&self, database: &str, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { + let generation = self.next_generation(); + let active_refreshes = self + .active_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()); + let oldest_active = active_refreshes.first().copied(); + let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); + let mut database_generations = self + .database_generations + .write() + .unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + update(&mut next); + database_generations.insert(database.to_string(), generation); + Self::prune_database_generations(&next, &mut database_generations, oldest_active); + *current = Arc::new(next); + } + + fn set_object_resolution(&self, database: &str, object: &str, resolution: ObjectResolution) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .insert((database.to_string(), object.to_string()), resolution); + } + + fn object_resolution(&self, database: &str, object: &str) -> Option { + self.object_resolutions + .read() + .unwrap_or_else(|e| e.into_inner()) + .get(&(database.to_string(), object.to_string())) + .copied() + } + + fn system_table_capabilities(&self, database: &str) -> HashMap { + self.snapshot + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .get(database) + .map(|metadata| metadata.system_table_capabilities.clone()) + .unwrap_or_default() + } + + fn remove_object_resolution(&self, database: &str, object: &str) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .remove(&(database.to_string(), object.to_string())); + } + + fn remove_database_resolutions(&self, database: &str) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .retain(|(resolved_database, _), _| resolved_database != database); + } + + fn rename_object_resolution(&self, database: &str, from: &str, to: &str) { + let mut resolutions = self + .object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()); + if let Some(resolution) = resolutions.remove(&(database.to_string(), from.to_string())) { + resolutions.insert((database.to_string(), to.to_string()), resolution); + } + } + + fn retain_object_resolutions(&self, snapshot: &CatalogMetadataSnapshot) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .retain(|(database, object), _| { + snapshot + .databases + .get(database) + .is_some_and(|metadata| metadata.objects.contains_key(object)) + }); + } + + fn clear_unavailable_resolutions(&self, databases: &HashSet) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .retain(|(database, _), resolution| { + *resolution != ObjectResolution::Unavailable || !databases.contains(database) + }); + } + + fn clear_unavailable_resolution_for_database(&self, database: &str) { + self.object_resolutions + .write() + .unwrap_or_else(|e| e.into_inner()) + .retain(|(resolved_database, _), resolution| { + *resolution != ObjectResolution::Unavailable || resolved_database != database + }); + } + + fn prune_database_generations( + snapshot: &CatalogMetadataSnapshot, + database_generations: &mut HashMap, + oldest_active: Option, + ) { + database_generations.retain(|database, generation| { + !snapshot.databases.contains_key(database) + || oldest_active.is_some_and(|oldest| *generation >= oldest) + }); + + let mut evictable_tombstones: Vec<_> = database_generations + .iter() + .filter(|(database, generation)| { + !snapshot.databases.contains_key(*database) + && !oldest_active.is_some_and(|oldest| **generation >= oldest) + }) + .map(|(database, generation)| (database.clone(), *generation)) + .collect(); + if evictable_tombstones.len() > MAX_RETAINED_DATABASE_TOMBSTONES { + evictable_tombstones.sort_unstable_by_key(|(_, generation)| *generation); + let remove_count = evictable_tombstones.len() - MAX_RETAINED_DATABASE_TOMBSTONES; + for (database, _) in evictable_tombstones.into_iter().take(remove_count) { + database_generations.remove(&database); + } + } + } +} + +type SharedCatalogMetadata = Arc; + +const MAX_CONCURRENT_METADATA_LISTINGS: usize = 16; +const MAX_TABLE_TYPE_CLASSIFICATION_FAILURES: usize = 16; +pub(crate) const DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS: usize = 16; +const DEFAULT_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(30); +// Recent tombstones guard against eventually consistent database listings. Tombstones still +// needed by an active older refresh are exempt from this bound until that refresh completes. +const MAX_RETAINED_DATABASE_TOMBSTONES: usize = 1024; + +#[derive(Debug)] +struct TableTypeClassificationBudget { + failure_permits: Arc, + failures: Mutex>, + skipped: AtomicUsize, +} + +impl TableTypeClassificationBudget { + fn new() -> Arc { + Arc::new(Self { + failure_permits: Arc::new(tokio::sync::Semaphore::new( + MAX_TABLE_TYPE_CLASSIFICATION_FAILURES, + )), + failures: Mutex::new(Vec::new()), + skipped: AtomicUsize::new(0), + }) + } + + async fn reserve_failure_permit(&self) -> Option { + match Arc::clone(&self.failure_permits).acquire_owned().await { + Ok(permit) => Some(permit), + Err(_) => { + self.skipped.fetch_add(1, Ordering::Relaxed); + None + } + } + } + + fn record_failure(&self, message: String, permit: tokio::sync::OwnedSemaphorePermit) { + let exhausted = { + let mut failures = self.failures.lock().unwrap_or_else(|e| e.into_inner()); + failures.push(message); + failures.len() >= MAX_TABLE_TYPE_CLASSIFICATION_FAILURES + }; + permit.forget(); + if exhausted { + self.failure_permits.close(); + } + } + + fn error(&self) -> Option { + let failures = self.failures.lock().unwrap_or_else(|e| e.into_inner()); + if failures.is_empty() { + return None; + } + let skipped = self.skipped.load(Ordering::Relaxed); + Some(DataFusionError::Execution(format!( + "table type classification failed for {} object(s); skipped {skipped} additional object(s) after the failure budget was exhausted: {}", + failures.len(), + failures.join("; ") + ))) + } +} + +fn finish_table_type_classification( + budget: &TableTypeClassificationBudget, + fail_on_error: bool, +) -> DFResult<()> { + let Some(error) = budget.error() else { + return Ok(()); + }; + if fail_on_error { + Err(error) + } else { + log::warn!("metadata snapshot contains unclassified tables: {error}"); + Ok(()) + } +} + +async fn load_database_metadata( + catalog: &dyn Catalog, + database: &str, + ignore_missing_views_endpoint: bool, + known_capabilities: HashMap, + metadata_io_semaphore: &tokio::sync::Semaphore, + classification_budget: Arc, +) -> DFResult { + let tables = async { + let _permit = metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + catalog + .list_tables(database) + .await + .map_err(to_datafusion_error) + }; + let views = async { + let _permit = metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + match catalog.list_views(database).await { + Ok(names) => Ok(names), + Err(paimon::Error::Unsupported { .. }) => Ok(vec![]), + Err( + error @ paimon::Error::RestApi { + source: paimon::api::RestError::NoSuchResource { .. }, + }, + ) if ignore_missing_views_endpoint => { + log::debug!( + "ignoring unavailable views endpoint while initializing database \ + '{database}': {error}" + ); + Ok(vec![]) + } + Err(error) => Err(to_datafusion_error(error)), + } + }; + let (mut table_names, mut view_names) = futures::try_join!(tables, views)?; + let mut seen_tables = HashSet::new(); + table_names.retain(|name| seen_tables.insert(name.clone())); + let mut seen_views = HashSet::new(); + view_names.retain(|name| seen_views.insert(name.clone())); + + let unresolved_names: Vec<_> = table_names + .iter() + .filter(|name| { + !matches!( + known_capabilities.get(*name), + Some(SystemTableCapability::Paimon | SystemTableCapability::Unsupported) + ) + }) + .cloned() + .collect(); + let declared_types: HashMap<_, _> = stream::iter(unresolved_names) + .map(|name| { + let classification_budget = Arc::clone(&classification_budget); + async move { + let Some(failure_permit) = classification_budget.reserve_failure_permit().await + else { + return Ok::<_, DataFusionError>(None); + }; + let _permit = metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + match catalog + .list_table_types(database, std::slice::from_ref(&name)) + .await + { + Ok(mut types) => match types.remove(&name) { + Some(table_type) => Ok::<_, DataFusionError>(Some((name, table_type))), + None => Ok(None), + }, + Err(error) => { + log::debug!( + "unable to classify table type for '{database}.{name}': {error}" + ); + classification_budget.record_failure( + format!("'{database}.{name}': {error}"), + failure_permit, + ); + Ok(None) + } + } + } + }) + .buffer_unordered(MAX_CONCURRENT_METADATA_LISTINGS) + .try_collect::>() + .await? + .into_iter() + .flatten() + .collect(); + + let mut objects = IndexMap::with_capacity(table_names.len() + view_names.len()); + let mut system_table_capabilities = HashMap::with_capacity(table_names.len()); + for name in table_names { + let capability = match known_capabilities.get(&name).copied() { + Some( + capability @ (SystemTableCapability::Paimon | SystemTableCapability::Unsupported), + ) => capability, + Some(SystemTableCapability::Unknown) | None => declared_types + .get(&name) + .map(|table_type| { + if table_type.requires_table_engine() { + SystemTableCapability::Unsupported + } else { + SystemTableCapability::Paimon + } + }) + .unwrap_or(SystemTableCapability::Unknown), + }; + system_table_capabilities.insert(name.clone(), capability); + objects.entry(name).or_insert(TableType::Base); + } + for name in view_names { + objects.entry(name).or_insert(TableType::View); + } + Ok(DatabaseMetadata { + objects, + system_table_capabilities, + }) +} + /// What an engine is asked to resolve. Non-exhaustive so later releases can /// carry more of the request — a snapshot selector, say — without breaking /// existing resolvers. @@ -162,17 +672,12 @@ impl TableProvider for ReadOnlyTableProvider { #[derive(Debug)] struct UnavailableEngineTableProvider { schema: datafusion::arrow::datatypes::SchemaRef, - declared: PaimonTableType, - table_name: String, + error_message: String, } impl UnavailableEngineTableProvider { fn unavailable_error(&self) -> datafusion::error::DataFusionError { - plan_datafusion_err!( - "no table engine is registered for '{}' tables ('{}')", - self.declared, - self.table_name - ) + plan_datafusion_err!("{}", self.error_message) } } @@ -232,8 +737,10 @@ pub fn register_catalog_table_engine( /// Provides an interface to manage and access multiple schemas (databases) /// within a Paimon [`Catalog`]. /// -/// This provider uses lazy loading - databases and tables are fetched -/// on-demand from the catalog, ensuring data is always fresh. +/// Database and object listings are refreshed asynchronously and served from an +/// in-memory snapshot because DataFusion's catalog discovery callbacks are +/// synchronous. Concrete tables are still loaded lazily by the async +/// [`SchemaProvider::table`] callback. /// /// # Table-version clauses /// @@ -265,6 +772,11 @@ pub struct PaimonCatalogProvider { /// Engines for table types served elsewhere, keyed by declared /// [`PaimonTableType`]. Same poison-recovery stance as `temp_tables`. table_engines: TableEngines, + /// Remotely refreshed metadata used by DataFusion's synchronous catalog callbacks. + metadata: SharedCatalogMetadata, + /// Shared bound for metadata requests issued by this session's providers. + metadata_io_semaphore: Arc, + metadata_refresh_timeout: Duration, } impl Debug for PaimonCatalogProvider { @@ -274,13 +786,37 @@ impl Debug for PaimonCatalogProvider { } impl PaimonCatalogProvider { - /// Creates a new [`PaimonCatalogProvider`]. - pub fn new( + /// Creates a provider with an empty metadata snapshot. + /// + /// Callers must refresh it before exposing synchronous discovery callbacks. + pub fn new_uninitialized( + catalog_name: Option, + catalog: Arc, + dynamic_options: DynamicOptions, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + ) -> Self { + Self::new_uninitialized_with_metadata_io_semaphore( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + Arc::new(tokio::sync::Semaphore::new( + DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, + )), + DEFAULT_METADATA_REFRESH_TIMEOUT, + ) + } + + pub(crate) fn new_uninitialized_with_metadata_io_semaphore( catalog_name: Option, catalog: Arc, dynamic_options: DynamicOptions, blob_reader_registry: BlobReaderRegistry, session_state: Option, + metadata_io_semaphore: Arc, + metadata_refresh_timeout: Duration, ) -> Self { PaimonCatalogProvider { catalog_name, @@ -291,9 +827,247 @@ impl PaimonCatalogProvider { session_state, schema_force_view_types: true, table_engines: Arc::new(RwLock::new(HashMap::new())), + metadata: Arc::new(CatalogMetadataState::default()), + metadata_io_semaphore, + metadata_refresh_timeout, } } + /// Creates a provider with an initialized metadata snapshot. + /// + /// The old synchronous `new` entry point is intentionally unavailable because it cannot + /// return a ready metadata snapshot. Use this constructor, or choose + /// [`Self::new_uninitialized`] explicitly and call [`Self::initialize_metadata`] before + /// exposing synchronous discovery callbacks. + /// + /// ```compile_fail + /// use std::sync::Arc; + /// use paimon::Catalog; + /// use paimon_datafusion::PaimonCatalogProvider; + /// + /// let catalog: Arc = todo!(); + /// let _ = PaimonCatalogProvider::new( + /// None, + /// catalog, + /// Default::default(), + /// Default::default(), + /// None, + /// ); + /// ``` + pub async fn try_new( + catalog_name: Option, + catalog: Arc, + dynamic_options: DynamicOptions, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + ) -> DFResult { + Self::try_new_with_metadata_io_semaphore( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + Arc::new(tokio::sync::Semaphore::new( + DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, + )), + DEFAULT_METADATA_REFRESH_TIMEOUT, + ) + .await + } + + pub(crate) async fn try_new_with_metadata_io_semaphore( + catalog_name: Option, + catalog: Arc, + dynamic_options: DynamicOptions, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + metadata_io_semaphore: Arc, + metadata_refresh_timeout: Duration, + ) -> DFResult { + let provider = Self::new_uninitialized_with_metadata_io_semaphore( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + metadata_io_semaphore, + metadata_refresh_timeout, + ); + provider.initialize_metadata().await?; + Ok(provider) + } + + /// Refresh the metadata consumed by DataFusion's synchronous catalog callbacks. + /// + /// Remote calls finish before the shared snapshot is replaced, so readers either + /// observe the previous complete snapshot or the new complete snapshot. + pub async fn refresh_metadata(&self) -> DFResult<()> { + self.run_metadata_refresh(self.refresh_metadata_inner(false, true)) + .await + } + + /// Initialize metadata while tolerating a REST server without the optional views endpoint. + pub async fn initialize_metadata(&self) -> DFResult<()> { + self.run_metadata_refresh(self.refresh_metadata_inner(true, false)) + .await + } + + /// Refresh metadata for SQL planning while tolerating an optional missing views endpoint. + pub(crate) async fn refresh_metadata_for_sql(&self) -> DFResult<()> { + self.run_metadata_refresh(self.refresh_metadata_inner(true, true)) + .await + } + + async fn run_metadata_refresh( + &self, + refresh: impl std::future::Future>, + ) -> DFResult { + tokio::time::timeout(self.metadata_refresh_timeout, refresh) + .await + .map_err(|_| { + DataFusionError::Execution(format!( + "metadata refresh timed out after {:?}", + self.metadata_refresh_timeout + )) + })? + } + + /// Set the maximum duration of one metadata initialization or refresh. + pub fn with_metadata_refresh_timeout(mut self, timeout: Duration) -> Self { + self.metadata_refresh_timeout = timeout; + self + } + + async fn refresh_metadata_inner( + &self, + ignore_missing_views_endpoint: bool, + fail_on_classification_error: bool, + ) -> DFResult<()> { + let generation = self.metadata.begin_refresh(); + let classification_budget = TableTypeClassificationBudget::new(); + let mut database_names = { + let _permit = self.metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + self.catalog + .list_databases() + .await + .map_err(to_datafusion_error)? + }; + let mut seen_databases = HashSet::new(); + database_names.retain(|name| seen_databases.insert(name.clone())); + + let entries: HashMap<_, _> = stream::iter(database_names.iter().cloned()) + .map(|database| { + let known_capabilities = self.metadata.system_table_capabilities(&database); + let classification_budget = Arc::clone(&classification_budget); + async move { + let metadata = load_database_metadata( + self.catalog.as_ref(), + database.as_str(), + ignore_missing_views_endpoint, + known_capabilities, + self.metadata_io_semaphore.as_ref(), + Arc::clone(&classification_budget), + ) + .await?; + Ok::<_, datafusion::error::DataFusionError>((database, metadata)) + } + }) + .buffer_unordered(MAX_CONCURRENT_METADATA_LISTINGS) + .try_collect() + .await?; + let mut entries = entries; + let databases = database_names + .into_iter() + .filter_map(|database| { + entries + .remove(&database) + .map(|metadata| (database, Arc::new(metadata))) + }) + .collect(); + + let mut conflicts = self + .metadata + .publish(generation.get(), CatalogMetadataSnapshot { databases }); + if !conflicts.is_empty() { + let current_databases: HashSet<_> = { + let _permit = self.metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + self.catalog + .list_databases() + .await + .map_err(to_datafusion_error)? + .into_iter() + .collect() + }; + conflicts.retain(|database| current_databases.contains(database)); + } + stream::iter(conflicts) + .map(|database| { + let classification_budget = Arc::clone(&classification_budget); + async move { + self.refresh_database_metadata_inner( + database.as_str(), + ignore_missing_views_endpoint, + classification_budget, + ) + .await + } + }) + .buffer_unordered(MAX_CONCURRENT_METADATA_LISTINGS) + .try_collect::>() + .await?; + finish_table_type_classification( + classification_budget.as_ref(), + fail_on_classification_error, + ) + } + + /// Refresh one database in the metadata snapshot. + pub(crate) async fn refresh_database_metadata(&self, database: &str) -> DFResult<()> { + let classification_budget = TableTypeClassificationBudget::new(); + self.run_metadata_refresh(self.refresh_database_metadata_inner( + database, + true, + Arc::clone(&classification_budget), + )) + .await?; + finish_table_type_classification(classification_budget.as_ref(), true) + } + + async fn refresh_database_metadata_inner( + &self, + database: &str, + ignore_missing_views_endpoint: bool, + classification_budget: Arc, + ) -> DFResult<()> { + const MAX_PUBLICATION_ATTEMPTS: usize = 3; + for _ in 0..MAX_PUBLICATION_ATTEMPTS { + let generation = self.metadata.begin_refresh(); + let database_metadata = load_database_metadata( + self.catalog.as_ref(), + database, + ignore_missing_views_endpoint, + self.metadata.system_table_capabilities(database), + self.metadata_io_semaphore.as_ref(), + Arc::clone(&classification_budget), + ) + .await?; + if self.metadata.publish_database( + generation.get(), + database.to_string(), + Arc::new(database_metadata), + ) { + return Ok(()); + } + } + Err(DataFusionError::Execution(format!( + "metadata for database '{database}' changed during {MAX_PUBLICATION_ATTEMPTS} refresh attempts" + ))) + } + /// Configure whether table schemas use Arrow view types when available. /// /// Disable this for consumers that cannot operate on Arrow view arrays. This changes the @@ -333,79 +1107,228 @@ impl PaimonCatalogProvider { Arc::clone(&self.table_engines) } - fn paimon_schema(&self, name: &str) -> Option> { - let catalog = Arc::clone(&self.catalog); - let table_engines = self.table_engines(); - let dynamic_options = Arc::clone(&self.dynamic_options); - let blob_reader_registry = self.blob_reader_registry.clone(); - let catalog_name = self.catalog_name.clone(); - let session_state = self.session_state.clone(); - let schema_force_view_types = self.schema_force_view_types; - let name = name.to_string(); + pub(crate) fn metadata_contains_object(&self, database: &str, name: &str) -> bool { + if self + .temp_tables + .read() + .unwrap_or_else(|e| e.into_inner()) + .get(database) + .is_some_and(|provider| provider.table_exist(name)) + { + return true; + } + let object_name = system_tables::parse_object_name_for_datafusion(name) + .map(|object| object.table().to_string()) + .unwrap_or_else(|_| name.to_string()); + self.metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .get(database) + .is_some_and(|metadata| metadata.objects.contains_key(&object_name)) + } + pub(crate) fn record_table_created( + &self, + database: &str, + name: &str, + declared_type: Option, + ) { + self.metadata.mutate_database(database, |next| { + let database = next + .databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + let database = Arc::make_mut(database); + database.objects.insert(name.to_string(), TableType::Base); + let capability = match declared_type { + Some(table_type) if table_type.requires_table_engine() => { + SystemTableCapability::Unsupported + } + Some(_) => SystemTableCapability::Paimon, + None => SystemTableCapability::Unknown, + }; + database + .system_table_capabilities + .insert(name.to_string(), capability); + }); + match declared_type { + Some(PaimonTableType::ObjectTable) => { + self.metadata + .set_object_resolution(database, name, ObjectResolution::Object); + } + Some(table_type) if !table_type.requires_table_engine() => { + self.metadata + .set_object_resolution(database, name, ObjectResolution::Paimon); + } + Some(_) | None => self.metadata.remove_object_resolution(database, name), + } + } + + pub(crate) async fn reconcile_table_created(&self, database: &str, name: &str) { + if self.metadata_contains_object(database, name) { + return; + } + let declared_type = self + .run_metadata_refresh(async { + let _permit = self.metadata_io_semaphore.acquire().await.map_err(|_| { + DataFusionError::Execution("metadata request limiter was closed".to_string()) + })?; + let mut types = self + .catalog + .list_table_types(database, &[name.to_string()]) + .await + .map_err(to_datafusion_error)?; + Ok(types.remove(name)) + }) + .await; + if self.metadata_contains_object(database, name) { + return; + } + match declared_type { + Ok(declared_type) => self.record_table_created(database, name, declared_type), + Err(error) => { + log::warn!( + "unable to reconcile created table '{database}.{name}' with remote metadata: {error}" + ); + self.record_table_created(database, name, None); + } + } + } + + pub(crate) fn record_view_created(&self, database: &str, name: &str) { + self.metadata.mutate_database(database, |next| { + let database = next + .databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + let database = Arc::make_mut(database); + database.objects.insert(name.to_string(), TableType::View); + database.system_table_capabilities.remove(name); + }); + self.metadata.remove_object_resolution(database, name); + } + + pub(crate) fn record_database_created(&self, database: &str) { + self.metadata.mutate_database(database, |next| { + next.databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + }); + } + + pub(crate) fn record_database_dropped(&self, database: &str) { + self.metadata.mutate_database(database, |next| { + next.databases.shift_remove(database); + }); + self.metadata.remove_database_resolutions(database); + } + + pub(crate) fn record_object_dropped(&self, database: &str, name: &str) { + self.metadata.mutate_database(database, |next| { + if let Some(metadata) = next.databases.get_mut(database) { + let metadata = Arc::make_mut(metadata); + metadata.objects.shift_remove(name); + metadata.system_table_capabilities.remove(name); + } + }); + self.metadata.remove_object_resolution(database, name); + } + + pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) -> bool { + let mut renamed = false; + self.metadata.mutate_database(database, |next| { + let metadata = next + .databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + let metadata = Arc::make_mut(metadata); + let objects = &mut metadata.objects; + if let Some(table_type) = objects.shift_remove(from) { + objects.insert(to.to_string(), table_type); + let capability = metadata + .system_table_capabilities + .remove(from) + .unwrap_or(SystemTableCapability::Unknown); + metadata + .system_table_capabilities + .insert(to.to_string(), capability); + renamed = true; + } else { + objects.insert(to.to_string(), TableType::Base); + metadata + .system_table_capabilities + .insert(to.to_string(), SystemTableCapability::Unknown); + } + }); + if renamed { + self.metadata.rename_object_resolution(database, from, to); + } else { + self.metadata.remove_object_resolution(database, from); + self.metadata.remove_object_resolution(database, to); + } + renamed + } + + fn paimon_schema(&self, name: &str) -> Option> { let temp_provider = { let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); - databases.get(&name).cloned() + databases.get(name).cloned() }; + let catalog_has_database = self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .contains_key(name); + if !catalog_has_database && temp_provider.is_none() { + return None; + } - block_on_with_runtime( - async move { - match catalog.get_database(&name).await { - Ok(_) => Some(Arc::new( - PaimonSchemaProvider::new( - catalog_name, - Arc::clone(&catalog), - name, - dynamic_options, - temp_provider, - blob_reader_registry, - session_state, - ) - .with_schema_force_view_types(schema_force_view_types) - .with_table_engines(Arc::clone(&table_engines)), - ) as Arc), - Err(paimon::Error::DatabaseNotExist { .. }) => { - if temp_provider.is_some() { - Some(Arc::new( - PaimonSchemaProvider::new( - catalog_name, - Arc::clone(&catalog), - name, - dynamic_options, - temp_provider, - blob_reader_registry, - session_state, - ) - .with_schema_force_view_types(schema_force_view_types) - .with_table_engines(Arc::clone(&table_engines)), - ) as Arc) - } else { - None - } - } - Err(e) => { - log::error!("failed to get database '{}': {e}", name); - None - } - } - }, - "paimon catalog access thread panicked", - ) + Some(Arc::new( + PaimonSchemaProvider::new_uninitialized( + self.catalog_name.clone(), + Arc::clone(&self.catalog), + name.to_string(), + Arc::clone(&self.dynamic_options), + temp_provider, + self.blob_reader_registry.clone(), + self.session_state.clone(), + ) + .with_metadata_io_semaphore(Arc::clone(&self.metadata_io_semaphore)) + .with_metadata_refresh_timeout(self.metadata_refresh_timeout) + .with_schema_force_view_types(self.schema_force_view_types) + .with_table_engines(self.table_engines()) + .with_metadata_snapshot(Arc::clone(&self.metadata)), + ) as Arc) } } impl CatalogProvider for PaimonCatalogProvider { fn schema_names(&self) -> Vec { - let catalog = Arc::clone(&self.catalog); - block_on_with_runtime( - async move { - catalog.list_databases().await.unwrap_or_else(|e| { - log::error!("failed to list databases: {e}"); - vec![] - }) - }, - "paimon catalog access thread panicked", - ) + let mut names: Vec = self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .keys() + .cloned() + .collect(); + let mut temp_names: Vec<_> = self + .temp_tables + .read() + .unwrap_or_else(|e| e.into_inner()) + .keys() + .cloned() + .collect(); + temp_names.sort_unstable(); + let mut seen: HashSet<_> = names.iter().cloned().collect(); + names.extend( + temp_names + .into_iter() + .filter(|name| seen.insert(name.clone())), + ); + names } fn schema(&self, name: &str) -> Option> { @@ -423,6 +1346,9 @@ impl CatalogProvider for PaimonCatalogProvider { let catalog_name = self.catalog_name.clone(); let session_state = self.session_state.clone(); let schema_force_view_types = self.schema_force_view_types; + let table_engines = self.table_engines(); + let metadata = Arc::clone(&self.metadata); + let metadata_io_semaphore = Arc::clone(&self.metadata_io_semaphore); let name = name.to_string(); block_on_with_runtime( async move { @@ -430,8 +1356,13 @@ impl CatalogProvider for PaimonCatalogProvider { .create_database(&name, false, HashMap::new()) .await .map_err(to_datafusion_error)?; + metadata.mutate_database(&name, |next| { + next.databases + .entry(name.clone()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + }); Ok(Some(Arc::new( - PaimonSchemaProvider::new( + PaimonSchemaProvider::new_uninitialized( catalog_name, Arc::clone(&catalog), name, @@ -440,7 +1371,10 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) - .with_schema_force_view_types(schema_force_view_types), + .with_metadata_io_semaphore(metadata_io_semaphore) + .with_schema_force_view_types(schema_force_view_types) + .with_table_engines(table_engines) + .with_metadata_snapshot(metadata), ) as Arc)) }, "paimon catalog access thread panicked", @@ -458,6 +1392,9 @@ impl CatalogProvider for PaimonCatalogProvider { let catalog_name = self.catalog_name.clone(); let session_state = self.session_state.clone(); let schema_force_view_types = self.schema_force_view_types; + let table_engines = self.table_engines(); + let metadata = Arc::clone(&self.metadata); + let metadata_io_semaphore = Arc::clone(&self.metadata_io_semaphore); let name = name.to_string(); block_on_with_runtime( async move { @@ -465,8 +1402,11 @@ impl CatalogProvider for PaimonCatalogProvider { .drop_database(&name, false, cascade) .await .map_err(to_datafusion_error)?; + metadata.mutate_database(&name, |next| { + next.databases.shift_remove(&name); + }); Ok(Some(Arc::new( - PaimonSchemaProvider::new( + PaimonSchemaProvider::new_uninitialized( catalog_name, Arc::clone(&catalog), name, @@ -475,7 +1415,10 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) - .with_schema_force_view_types(schema_force_view_types), + .with_metadata_io_semaphore(metadata_io_semaphore) + .with_schema_force_view_types(schema_force_view_types) + .with_table_engines(table_engines) + .with_metadata_snapshot(metadata), ) as Arc)) }, "paimon catalog access thread panicked", @@ -495,21 +1438,15 @@ impl PaimonCatalogProvider { table_name: &str, table: Arc, ) -> DFResult<()> { - // Warn if this shadows a real Paimon table (outside the lock — not critical) - let catalog = Arc::clone(&self.catalog); - let db = database.to_string(); - let tbl = table_name.to_string(); - let identifier = Identifier::new(db, tbl); - if let Ok(true) = block_on_with_runtime( - async move { - match catalog.get_table(&identifier).await { - Ok(_) => Ok::(true), - Err(paimon::Error::TableNotExist { .. }) => Ok(false), - Err(_) => Ok(false), - } - }, - "paimon catalog access thread panicked", - ) { + // The warning is best-effort and must not turn this synchronous API into remote I/O. + if self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .get(database) + .is_some_and(|metadata| metadata.objects.contains_key(table_name)) + { log::warn!( "Temporary table '{database}.{table_name}' shadows an existing Paimon table" ); @@ -575,13 +1512,16 @@ pub struct PaimonSchemaProvider { dynamic_options: DynamicOptions, /// Optional temporary in-memory provider for temp tables and views. temp_provider: Option>, - /// Table types populated together with `table_names` for metadata-only lookups. - catalog_table_types: RwLock>, + /// Shared metadata used by synchronous schema callbacks. + metadata: SharedCatalogMetadata, blob_reader_registry: BlobReaderRegistry, session_state: Option, schema_force_view_types: bool, /// Engines for table types served elsewhere; empty without routing. table_engines: TableEngines, + /// Shared bound for metadata requests issued by this session's providers. + metadata_io_semaphore: Arc, + metadata_refresh_timeout: Duration, } impl Debug for PaimonSchemaProvider { @@ -594,8 +1534,8 @@ impl Debug for PaimonSchemaProvider { } impl PaimonSchemaProvider { - /// Creates a new [`PaimonSchemaProvider`]. - pub fn new( + /// Creates a schema provider with an empty metadata snapshot. + pub fn new_uninitialized( catalog_name: Option, catalog: Arc, database: String, @@ -610,12 +1550,119 @@ impl PaimonSchemaProvider { database, dynamic_options, temp_provider, - catalog_table_types: RwLock::new(HashMap::new()), + metadata: Arc::new(CatalogMetadataState::default()), blob_reader_registry, session_state, schema_force_view_types: true, table_engines: Arc::new(RwLock::new(HashMap::new())), + metadata_io_semaphore: Arc::new(tokio::sync::Semaphore::new( + DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, + )), + metadata_refresh_timeout: DEFAULT_METADATA_REFRESH_TIMEOUT, + } + } + + /// Creates a schema provider with initialized metadata. + /// + /// The old synchronous `new` entry point is intentionally unavailable because it cannot + /// return a ready metadata snapshot. Use this constructor, or choose + /// [`Self::new_uninitialized`] explicitly and call [`Self::initialize_metadata`] before + /// exposing synchronous discovery callbacks. + /// + /// ```compile_fail + /// use std::sync::Arc; + /// use paimon::Catalog; + /// use paimon_datafusion::PaimonSchemaProvider; + /// + /// let catalog: Arc = todo!(); + /// let _ = PaimonSchemaProvider::new( + /// None, + /// catalog, + /// "default".to_string(), + /// Default::default(), + /// None, + /// Default::default(), + /// None, + /// ); + /// ``` + pub async fn try_new( + catalog_name: Option, + catalog: Arc, + database: String, + dynamic_options: DynamicOptions, + temp_provider: Option>, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + ) -> DFResult { + let provider = Self::new_uninitialized( + catalog_name, + catalog, + database, + dynamic_options, + temp_provider, + blob_reader_registry, + session_state, + ); + provider.initialize_metadata().await?; + Ok(provider) + } + + /// Refresh this database in the snapshot used by synchronous callbacks. + pub async fn refresh_metadata(&self) -> DFResult<()> { + self.run_metadata_refresh(self.refresh_metadata_inner(false, true)) + .await + } + + /// Initialize metadata while tolerating a REST server without the optional views endpoint. + pub async fn initialize_metadata(&self) -> DFResult<()> { + self.run_metadata_refresh(self.refresh_metadata_inner(true, false)) + .await + } + + async fn run_metadata_refresh( + &self, + refresh: impl std::future::Future>, + ) -> DFResult { + tokio::time::timeout(self.metadata_refresh_timeout, refresh) + .await + .map_err(|_| { + DataFusionError::Execution(format!( + "metadata refresh timed out after {:?}", + self.metadata_refresh_timeout + )) + })? + } + + async fn refresh_metadata_inner( + &self, + ignore_missing_views_endpoint: bool, + fail_on_classification_error: bool, + ) -> DFResult<()> { + let generation = self.metadata.begin_refresh(); + let classification_budget = TableTypeClassificationBudget::new(); + let database = load_database_metadata( + self.catalog.as_ref(), + &self.database, + ignore_missing_views_endpoint, + self.metadata.system_table_capabilities(&self.database), + self.metadata_io_semaphore.as_ref(), + Arc::clone(&classification_budget), + ) + .await?; + if !self.metadata.publish_database( + generation.get(), + self.database.clone(), + Arc::new(database), + ) { + return Err(DataFusionError::Execution(format!( + "metadata for database '{}' changed during refresh", + self.database + ))); } + finish_table_type_classification( + classification_budget.as_ref(), + fail_on_classification_error, + ) } fn with_schema_force_view_types(mut self, schema_force_view_types: bool) -> Self { @@ -623,56 +1670,41 @@ impl PaimonSchemaProvider { self } + fn with_metadata_io_semaphore( + mut self, + metadata_io_semaphore: Arc, + ) -> Self { + self.metadata_io_semaphore = metadata_io_semaphore; + self + } + + fn with_metadata_refresh_timeout(mut self, timeout: Duration) -> Self { + self.metadata_refresh_timeout = timeout; + self + } + pub(crate) fn with_table_engines(mut self, table_engines: TableEngines) -> Self { self.table_engines = table_engines; self } + + fn with_metadata_snapshot(mut self, metadata: SharedCatalogMetadata) -> Self { + self.metadata = metadata; + self + } } #[async_trait] impl SchemaProvider for PaimonSchemaProvider { fn table_names(&self) -> Vec { - let catalog = Arc::clone(&self.catalog); - let database = self.database.clone(); - let (mut names, views) = block_on_with_runtime( - { - let db = database.clone(); - async move { - let names = match catalog.list_tables(&db).await { - Ok(names) => names, - Err(e) => { - log::error!("failed to list tables in '{}': {e}", db); - vec![] - } - }; - let views = match catalog.list_views(&db).await { - Ok(views) => views, - Err(paimon::Error::Unsupported { .. }) => vec![], - Err(error) => { - log::error!("failed to list views in '{}': {error}", db); - vec![] - } - }; - (names, views) - } - }, - "paimon catalog access thread panicked", - ); - - let mut catalog_table_types = HashMap::with_capacity(names.len() + views.len()); - for name in &names { - catalog_table_types.insert(name.clone(), TableType::Base); - } - for view in views { - catalog_table_types - .entry(view.clone()) - .or_insert(TableType::View); - names.push(view); - } - *self - .catalog_table_types - .write() - .unwrap_or_else(|e| e.into_inner()) = catalog_table_types; + let mut names: Vec = self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .get(&self.database) + .map(|database| database.objects.keys().cloned().collect()) + .unwrap_or_default(); if let Some(temp) = &self.temp_provider { names.extend(temp.table_names()); @@ -714,8 +1746,17 @@ impl SchemaProvider for PaimonSchemaProvider { let catalog_name = self.catalog_name.clone(); let session_state = self.session_state.clone(); let schema_force_view_types = self.schema_force_view_types; + let metadata = Arc::clone(&self.metadata); let identifier = Identifier::new(self.database.clone(), object.table().to_string()); let branch = object.branch().map(str::to_string); + if branch.is_none() + && session_state + .as_ref() + .and_then(|provider| provider()) + .is_some_and(|state| state.table_functions().contains_key(identifier.object())) + { + return Ok(None); + } let table_engines: HashMap> = self .table_engines .read() @@ -724,6 +1765,11 @@ impl SchemaProvider for PaimonSchemaProvider { await_with_runtime(async move { match catalog.load_table(&identifier).await { Ok(paimon::catalog::LoadedTable::Object(table)) => { + metadata.set_object_resolution( + identifier.database(), + identifier.object(), + ObjectResolution::Object, + ); if branch.is_some() { return Err(plan_datafusion_err!( "branches are not supported for 'object-table' tables ('{}')", @@ -744,6 +1790,11 @@ impl SchemaProvider for PaimonSchemaProvider { } Ok(paimon::catalog::LoadedTable::External(external)) => { let declared = external.declared(); + metadata.set_object_resolution( + identifier.database(), + identifier.object(), + ObjectResolution::Routed, + ); if branch.is_some() { return Err(plan_datafusion_err!( "branches are not supported for '{}' tables ('{}')", @@ -761,8 +1812,10 @@ impl SchemaProvider for PaimonSchemaProvider { }; return Ok(Some(Arc::new(UnavailableEngineTableProvider { schema, - declared, - table_name: identifier.full_name(), + error_message: format!( + "no table engine is registered for '{declared}' tables ('{}')", + identifier.full_name() + ), }) as Arc)); }; // The Paimon arm below applies these; an engine would @@ -783,6 +1836,13 @@ impl SchemaProvider for PaimonSchemaProvider { declared, )) .await?; + if resolved.is_none() { + metadata.set_object_resolution( + identifier.database(), + identifier.object(), + ObjectResolution::Unavailable, + ); + } Ok(resolved.map(|inner| { Arc::new(ReadOnlyTableProvider { inner, @@ -792,6 +1852,11 @@ impl SchemaProvider for PaimonSchemaProvider { })) } Ok(paimon::catalog::LoadedTable::Paimon(table)) => { + metadata.set_object_resolution( + identifier.database(), + identifier.object(), + ObjectResolution::Paimon, + ); let mut table = *table; if let Some(branch) = branch.as_deref() { table = table @@ -827,16 +1892,6 @@ impl SchemaProvider for PaimonSchemaProvider { if branch.is_some() { return Ok(None); } - // DataFusion preloads every relation name before planning, including - // registered table functions. Do not reinterpret a missing UDTF name as - // a REST view; the planner will resolve it through the UDTF registry. - if session_state - .as_ref() - .and_then(|provider| provider()) - .is_some_and(|state| state.table_functions().contains_key(identifier.object())) - { - return Ok(None); - } let view = match catalog.get_view(&identifier).await { Ok(view) => view, Err(paimon::Error::ViewNotExist { .. }) @@ -895,12 +1950,19 @@ impl SchemaProvider for PaimonSchemaProvider { return Ok(Some(table_type)); } } + if self.metadata.object_resolution(&self.database, name) + == Some(ObjectResolution::Unavailable) + { + return Ok(None); + } if let Some(table_type) = self - .catalog_table_types + .metadata .read() .unwrap_or_else(|e| e.into_inner()) - .get(name) + .databases + .get(&self.database) + .and_then(|database| database.objects.get(name)) { return Ok(Some(*table_type)); } @@ -924,101 +1986,47 @@ impl SchemaProvider for PaimonSchemaProvider { return false; } }; - if let Some(system_name) = object.system_table() { - if !system_tables::is_registered(system_name) { - return false; - } - } - - let catalog = Arc::clone(&self.catalog); - let identifier = Identifier::new(self.database.clone(), object.table().to_string()); - let branch = object.branch().map(str::to_string); - let is_branches_table = object + if object .system_table() - .is_some_and(|name| name.eq_ignore_ascii_case("branches")); - let has_system_suffix = object.system_table().is_some(); - let engines: HashMap> = self - .table_engines + .is_some_and(|system_name| !system_tables::is_registered(system_name)) + { + return false; + } + let (table_type, snapshot_capability) = self + .metadata .read() .unwrap_or_else(|e| e.into_inner()) - .clone(); - block_on_with_runtime( - async move { - match catalog.load_table(&identifier).await { - Ok(paimon::catalog::LoadedTable::Object(_)) => { - branch.is_none() && !has_system_suffix - } - Ok(paimon::catalog::LoadedTable::External(external)) => { - let declared = external.declared(); - // Paimon-only; `table()` rejects them here too. - if branch.is_some() || has_system_suffix { - return false; - } - match engines.get(&declared) { - Some(resolver) => match resolver - .resolve_table(&EngineTableRequest::new( - identifier.database().to_string(), - identifier.object().to_string(), - declared, - )) - .await - { - Ok(table) => table.is_some(), - // Report failures as existing so `table()` - // surfaces the real error. - Err(err) => { - log::warn!( - "failed to probe engine table existence for '{}': {err}", - identifier.full_name() - ); - true - } - }, - // `table()` returns a metadata-only provider in this - // case, so the SchemaProvider existence contract - // requires the same answer here. - None => true, - } - } - Ok(paimon::catalog::LoadedTable::Paimon(table)) => { - if let Some(branch) = branch.as_deref() { - if is_branches_table { - return true; - } - (*table).copy_with_branch(branch).await.is_ok() - } else { - true - } - } - Err(paimon::Error::TableNotExist { .. }) => { - if branch.is_some() { - return false; - } - match catalog.get_view(&identifier).await { - Ok(_) => true, - Err(paimon::Error::ViewNotExist { .. }) - | Err(paimon::Error::Unsupported { .. }) => false, - Err(error) => { - log::error!("failed to check view '{}': {error}", identifier); - false - } - } - } - Err(e) => { - log::error!("failed to check table '{}': {e}", identifier); - false - } - Ok(_) => { - log::error!( - "catalog returned an unsupported loaded table kind for '{}'", - identifier - ); - false - } - } - }, - "paimon catalog access thread panicked", - ) + .databases + .get(&self.database) + .map(|database| { + ( + database.objects.get(object.table()).copied(), + database + .system_table_capabilities + .get(object.table()) + .copied(), + ) + }) + .unwrap_or((None, None)); + let resolution = self + .metadata + .object_resolution(&self.database, object.table()); + if resolution == Some(ObjectResolution::Unavailable) { + return false; + } + if object.system_table().is_some() { + let supports_system_tables = match resolution { + Some(ObjectResolution::Paimon) => true, + Some( + ObjectResolution::Object + | ObjectResolution::Routed + | ObjectResolution::Unavailable, + ) => false, + None => snapshot_capability == Some(SystemTableCapability::Paimon), + }; + return table_type == Some(TableType::Base) && supports_system_tables; + } + table_type.is_some() } fn register_table( @@ -1033,7 +2041,10 @@ impl SchemaProvider for PaimonSchemaProvider { fn deregister_table(&self, name: &str) -> DFResult>> { let catalog = Arc::clone(&self.catalog); - let identifier = Identifier::new(self.database.clone(), name); + let database = self.database.clone(); + let identifier = Identifier::new(database.clone(), name); + let metadata = Arc::clone(&self.metadata); + let name = name.to_string(); block_on_with_runtime( async move { // Try to get the table first so we can return it. @@ -1047,6 +2058,11 @@ impl SchemaProvider for PaimonSchemaProvider { .drop_table(&identifier, false) .await .map_err(to_datafusion_error)?; + metadata.mutate_database(&database, |next| { + if let Some(database) = next.databases.get_mut(&database) { + Arc::make_mut(database).objects.shift_remove(&name); + } + }); Ok(Some(Arc::new(provider) as Arc)) }, "paimon catalog access thread panicked", @@ -1332,6 +2348,57 @@ fn find_view_dependency_cycle( mod tests { use super::*; + #[test] + fn database_generation_tombstones_are_bounded() { + let metadata = CatalogMetadataState::default(); + for index in 0..2048 { + let database = format!("deleted_{index}"); + metadata.mutate_database(&database, |snapshot| { + snapshot.databases.shift_remove(&database); + }); + } + + assert_eq!( + metadata + .database_generations + .read() + .unwrap_or_else(|e| e.into_inner()) + .len(), + MAX_RETAINED_DATABASE_TOMBSTONES + ); + } + + #[test] + fn active_refresh_protects_tombstones_from_bounded_eviction() { + let metadata = CatalogMetadataState::default(); + let refresh = metadata.begin_refresh(); + for index in 0..(MAX_RETAINED_DATABASE_TOMBSTONES * 2) { + let database = format!("deleted_{index}"); + metadata.mutate_database(&database, |snapshot| { + snapshot.databases.shift_remove(&database); + }); + } + + assert_eq!( + metadata + .database_generations + .read() + .unwrap_or_else(|e| e.into_inner()) + .len(), + MAX_RETAINED_DATABASE_TOMBSTONES * 2 + ); + + drop(refresh); + assert_eq!( + metadata + .database_generations + .read() + .unwrap_or_else(|e| e.into_inner()) + .len(), + MAX_RETAINED_DATABASE_TOMBSTONES + ); + } + #[test] fn relation_identifiers_follow_datafusion_normalization() { let relation = ObjectName(vec![ diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index aba77b159..67f5c2118 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -45,7 +45,8 @@ //! - `TRUNCATE TABLE db.t PARTITION (col = val, ...)` use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{Arc, Mutex, Weak}; +use std::time::{Duration, Instant}; use datafusion::arrow::array::{ new_null_array, ArrayRef, BooleanArray, Date32Array, Float32Array, Float64Array, Int16Array, @@ -67,8 +68,9 @@ use datafusion::sql::sqlparser::ast::{ AlterColumnOperation, AlterTableOperation, BinaryLength, CharacterLength, ColumnDef, ColumnOption, CreateFunction, CreateFunctionBody, CreateTable, CreateTableOptions, CreateView, Delete, Expr as SqlExpr, FromTable, FunctionBehavior, FunctionReturnType, Ident, Insert, Merge, - ObjectName, ObjectType, RenameTableNameKind, Reset, ResetStatement, Set, ShowCreateObject, - SqlOption, Statement, TableFactor, TableObject, Truncate, Update, Use, Value as SqlValue, + ObjectName, ObjectType, RenameTableNameKind, Reset, ResetStatement, SchemaName, Set, + ShowCreateObject, SqlOption, Statement, TableFactor, TableObject, Truncate, Update, Use, + Value as SqlValue, }; use datafusion::sql::sqlparser::dialect::GenericDialect; use datafusion::sql::sqlparser::keywords::Keyword; @@ -103,6 +105,29 @@ pub struct SQLContext { /// Session-scoped dynamic options set via `SET 'paimon.key' = 'value'`. dynamic_options: DynamicOptions, blob_reader_registry: BlobReaderRegistry, + /// Last successful refresh used to resolve a missing object, keyed by database. + missing_object_refreshes: Mutex>, + /// Last successful full metadata refresh, keyed by catalog. + catalog_refreshes: Mutex>, + /// Last failed automatic full metadata refresh, keyed by catalog. + catalog_refresh_failures: Mutex>, + metadata_refresh_gates: Mutex>>>, + catalog_metadata_refresh_ttl: Duration, + catalog_metadata_refresh_timeout: Duration, + catalog_metadata_refresh_semaphore: Arc, + catalog_metadata_request_semaphore: Arc, +} + +const MISSING_OBJECT_REFRESH_TTL: Duration = Duration::from_secs(1); +const CATALOG_METADATA_REFRESH_TTL: Duration = Duration::from_secs(1); +const CATALOG_METADATA_REFRESH_FAILURE_BACKOFF: Duration = Duration::from_secs(5); +const DEFAULT_CATALOG_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(30); +const DEFAULT_MAX_CONCURRENT_CATALOG_METADATA_REFRESHES: usize = 4; + +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +enum MetadataRefreshTarget { + Catalog(String), + Database { catalog: String, database: String }, } /// Builder for [`SQLContext`]. @@ -132,6 +157,10 @@ pub struct SQLContext { #[derive(Default)] pub struct SQLContextBuilder { runtime_env: Option>, + catalog_metadata_refresh_ttl: Option, + catalog_metadata_refresh_timeout: Option, + max_concurrent_catalog_metadata_refreshes: Option, + max_concurrent_catalog_metadata_requests: Option, } impl SQLContextBuilder { @@ -146,6 +175,38 @@ impl SQLContextBuilder { self } + /// Bounds each best-effort full catalog refresh triggered by SQL metadata queries. + pub fn with_catalog_metadata_refresh_timeout(mut self, timeout: Duration) -> Self { + self.catalog_metadata_refresh_timeout = Some(timeout); + self + } + + /// Sets how long automatic SQL metadata queries reuse a successful full refresh. + pub fn with_catalog_metadata_refresh_ttl(mut self, ttl: Duration) -> Self { + self.catalog_metadata_refresh_ttl = Some(ttl); + self + } + + /// Limits concurrent full catalog refreshes triggered by SQL metadata queries. + pub fn with_max_concurrent_catalog_metadata_refreshes(mut self, maximum: usize) -> Self { + assert!( + maximum > 0, + "catalog metadata refresh concurrency must be positive" + ); + self.max_concurrent_catalog_metadata_refreshes = Some(maximum); + self + } + + /// Limits the metadata requests issued across all catalogs and databases. + pub fn with_max_concurrent_catalog_metadata_requests(mut self, maximum: usize) -> Self { + assert!( + maximum > 0, + "catalog metadata request concurrency must be positive" + ); + self.max_concurrent_catalog_metadata_requests = Some(maximum); + self + } + /// Builds a [`SQLContext`]. pub fn build(self) -> SQLContext { let mut state_builder = SessionStateBuilder::new() @@ -170,6 +231,24 @@ impl SQLContextBuilder { catalogs: HashMap::new(), dynamic_options: Default::default(), blob_reader_registry: BlobReaderRegistry::default(), + missing_object_refreshes: Mutex::new(HashMap::new()), + catalog_refreshes: Mutex::new(HashMap::new()), + catalog_refresh_failures: Mutex::new(HashMap::new()), + metadata_refresh_gates: Mutex::new(HashMap::new()), + catalog_metadata_refresh_ttl: self + .catalog_metadata_refresh_ttl + .unwrap_or(CATALOG_METADATA_REFRESH_TTL), + catalog_metadata_refresh_timeout: self + .catalog_metadata_refresh_timeout + .unwrap_or(DEFAULT_CATALOG_METADATA_REFRESH_TIMEOUT), + catalog_metadata_refresh_semaphore: Arc::new(tokio::sync::Semaphore::new( + self.max_concurrent_catalog_metadata_refreshes + .unwrap_or(DEFAULT_MAX_CONCURRENT_CATALOG_METADATA_REFRESHES), + )), + catalog_metadata_request_semaphore: Arc::new(tokio::sync::Semaphore::new( + self.max_concurrent_catalog_metadata_requests + .unwrap_or(crate::catalog::DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS), + )), } } } @@ -252,16 +331,19 @@ impl SQLContext { let weak_state = self.ctx.state_weak_ref(); let session_state: crate::catalog::SessionStateProvider = Arc::new(move || weak_state.upgrade().map(|state| state.read().clone())); - self.ctx.register_catalog( - &catalog_name, - Arc::new(crate::catalog::PaimonCatalogProvider::new( + let provider = Arc::new( + crate::catalog::PaimonCatalogProvider::try_new_with_metadata_io_semaphore( Some(catalog_name.clone()), catalog.clone(), self.dynamic_options.clone(), self.blob_reader_registry.clone(), Some(session_state), - )), + Arc::clone(&self.catalog_metadata_request_semaphore), + self.catalog_metadata_refresh_timeout, + ) + .await?, ); + self.ctx.register_catalog(&catalog_name, provider); register_table_functions( &self.ctx, &catalog, @@ -269,6 +351,10 @@ impl SQLContext { self.dynamic_options.clone(), ); self.catalogs.insert(catalog_name.clone(), catalog); + self.catalog_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(catalog_name.clone(), Instant::now()); if is_first { self.set_current_catalog(catalog_name).await?; if let Some(default_db) = default_db { @@ -473,7 +559,74 @@ impl SQLContext { )); } - match &statements[0] { + let metadata_mutation_targets = self.metadata_change_targets(&statements[0])?; + let (mut metadata_to_refresh, best_effort_targets) = + self.metadata_refresh_targets(&statements[0])?; + metadata_to_refresh.retain(|target| !metadata_mutation_targets.contains(target)); + let statement = &statements[0]; + let metadata_mutation_targets = &metadata_mutation_targets; + let best_effort_targets = &best_effort_targets; + futures::future::try_join_all(metadata_to_refresh.into_iter().map(|target| async move { + let refresh_gate = self.metadata_refresh_gate(&target); + let _refresh_guard = refresh_gate.lock().await; + let (mut current_targets, _) = self.metadata_refresh_targets(statement)?; + current_targets.retain(|target| !metadata_mutation_targets.contains(target)); + if current_targets.contains(&target) { + let best_effort = best_effort_targets.contains(&target); + let refresh = async { + let _permit = if best_effort { + Some( + self.catalog_metadata_refresh_semaphore + .acquire() + .await + .map_err(|_| { + DataFusionError::Execution( + "catalog metadata refresh coordinator closed".to_string(), + ) + })?, + ) + } else { + None + }; + self.refresh_metadata_targets([target.clone()]).await + }; + let refresh_result = if best_effort { + match tokio::time::timeout(self.catalog_metadata_refresh_timeout, refresh).await + { + Ok(result) => result, + Err(_) => Err(DataFusionError::Execution(format!( + "catalog metadata refresh timed out after {:?}", + self.catalog_metadata_refresh_timeout + ))), + } + } else { + refresh.await + }; + if let Err(error) = refresh_result { + if best_effort { + self.record_catalog_refresh_failure(&target); + log::warn!("information schema metadata refresh failed: {error}"); + } else { + return Err(error); + } + } + } + Ok::<(), DataFusionError>(()) + })) + .await?; + + let mut ordered_mutation_targets: Vec<_> = metadata_mutation_targets.iter().collect(); + ordered_mutation_targets.sort_unstable(); + let mutation_gates: Vec<_> = ordered_mutation_targets + .into_iter() + .map(|target| self.metadata_refresh_gate(target)) + .collect(); + let mut _mutation_guards = Vec::with_capacity(mutation_gates.len()); + for gate in &mutation_gates { + _mutation_guards.push(gate.lock().await); + } + + let result = match &statements[0] { Statement::ShowDatabases { terse, history, @@ -515,6 +668,29 @@ impl SQLContext { } self.handle_create_database(db_name, *if_not_exists).await } + Statement::CreateSchema { + schema_name: SchemaName::Simple(schema_name), + if_not_exists, + with, + options, + default_collate_spec, + clone, + } => { + if with.as_ref().is_some_and(|options| !options.is_empty()) + || options.as_ref().is_some_and(|options| !options.is_empty()) + || default_collate_spec.is_some() + || clone.is_some() + { + return Err(DataFusionError::Plan( + "CREATE SCHEMA options are not supported".to_string(), + )); + } + self.handle_create_database(schema_name, *if_not_exists) + .await + } + Statement::CreateSchema { .. } => Err(DataFusionError::Plan( + "CREATE SCHEMA AUTHORIZATION is not supported".to_string(), + )), Statement::Use(Use::Object(name)) => self.handle_use_database(name).await, Statement::CreateTable(create_table) => { if create_table.temporary { @@ -641,7 +817,7 @@ impl SQLContext { } } Statement::Drop { - object_type: ObjectType::Database, + object_type: object_type @ (ObjectType::Database | ObjectType::Schema), if_exists, names, cascade, @@ -650,15 +826,21 @@ impl SQLContext { temporary, table, } => { + let object_name = if *object_type == ObjectType::Database { + "DATABASE" + } else { + "SCHEMA" + }; let [name] = names.as_slice() else { - return Err(DataFusionError::Plan( - "DROP DATABASE requires exactly one database".to_string(), - )); + return Err(DataFusionError::Plan(format!( + "DROP {object_name} requires exactly one {}", + object_name.to_ascii_lowercase() + ))); }; if *restrict || *purge || *temporary || table.is_some() { - return Err(DataFusionError::Plan( - "DROP DATABASE options are not supported".to_string(), - )); + return Err(DataFusionError::Plan(format!( + "DROP {object_name} options are not supported" + ))); } self.handle_drop_database(name, *if_exists, *cascade).await } @@ -751,7 +933,391 @@ impl SQLContext { self.ctx.sql(&expanded.to_string()).await } _ => self.ctx.sql(sql).await, + }; + + if result.is_ok() { + self.apply_metadata_change(&statements[0]).await?; + } + result + } + + fn metadata_refresh_targets( + &self, + statement: &Statement, + ) -> DFResult<( + HashSet, + HashSet, + )> { + let mut targets = HashSet::new(); + let mut best_effort_targets = HashSet::new(); + if matches!(statement, Statement::ShowTables { .. }) { + let state = self.ctx.state(); + targets.insert(MetadataRefreshTarget::Database { + catalog: self.current_catalog_name(), + database: state.config_options().catalog.default_schema.clone(), + }); + return Ok((targets, best_effort_targets)); + } + if let Statement::ShowColumns { show_options, .. } = statement { + if let Some(show_in) = &show_options.show_in { + if show_in.parent_type.is_none() { + if let Some(name) = &show_in.parent_name { + let (_, catalog, identifier) = self.resolve_catalog_and_table(name)?; + targets.insert(MetadataRefreshTarget::Database { + catalog, + database: identifier.database().to_string(), + }); + } + } + } + return Ok((targets, best_effort_targets)); + } + if matches!(statement, Statement::ShowFunctions { .. }) { + return Ok((targets, best_effort_targets)); + } + + let statement = datafusion::sql::parser::Statement::Statement(Box::new(statement.clone())); + let state = self.ctx.state(); + let default_catalog = state.config_options().catalog.default_catalog.clone(); + let default_schema = state.config_options().catalog.default_schema.clone(); + for reference in state.resolve_table_references(&statement)? { + let schema = reference.schema().unwrap_or(&default_schema); + let catalog_name = reference.catalog().unwrap_or(&default_catalog); + if !self.catalogs.contains_key(catalog_name) { + continue; + } + let provider = self.ctx.catalog(catalog_name).ok_or_else(|| { + DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")) + })?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Catalog '{catalog_name}' is not a Paimon catalog" + )) + })?; + if schema.eq_ignore_ascii_case("information_schema") { + let catalogs = self + .catalogs + .keys() + .filter(|catalog| !self.catalog_refresh_is_recent(catalog)) + .cloned() + .map(MetadataRefreshTarget::Catalog); + targets.extend(catalogs.clone()); + best_effort_targets.extend(catalogs); + } else if !provider.metadata_contains_object(schema, reference.table()) { + let target = MetadataRefreshTarget::Database { + catalog: catalog_name.to_string(), + database: schema.to_string(), + }; + if !self.missing_object_refresh_is_recent(&target) { + targets.insert(target); + } + } + } + Ok((targets, best_effort_targets)) + } + + fn metadata_change_targets( + &self, + statement: &Statement, + ) -> DFResult> { + let mut targets = HashSet::new(); + let mut add_database = |name: &ObjectName| -> DFResult<()> { + let table_ref: TableReference = name.to_string().as_str().into(); + if !self.is_paimon_catalog_ref(&table_ref) { + return Ok(()); + } + let (_, catalog, identifier) = self.resolve_catalog_and_table(name)?; + targets.insert(MetadataRefreshTarget::Database { + catalog, + database: identifier.database().to_string(), + }); + Ok(()) + }; + match statement { + Statement::CreateDatabase { db_name, .. } => { + let (_, catalog, database) = self.resolve_catalog_and_database(db_name)?; + targets.insert(MetadataRefreshTarget::Database { catalog, database }); + } + Statement::CreateSchema { + schema_name: SchemaName::Simple(name), + .. + } => { + let (_, catalog, database) = self.resolve_catalog_and_database(name)?; + targets.insert(MetadataRefreshTarget::Database { catalog, database }); + } + Statement::CreateTable(create) => add_database(&create.name)?, + Statement::CreateView(create) => add_database(&create.name)?, + Statement::AlterTable(alter) => add_database(&alter.name)?, + Statement::Drop { + object_type: ObjectType::Database | ObjectType::Schema, + names, + temporary: false, + .. + } => { + for name in names { + let (_, catalog, database) = self.resolve_catalog_and_database(name)?; + targets.insert(MetadataRefreshTarget::Database { catalog, database }); + } + } + Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + names, + .. + } => { + for name in names { + add_database(name)?; + } + } + _ => {} + } + Ok(targets) + } + + async fn apply_metadata_change(&self, statement: &Statement) -> DFResult<()> { + match statement { + Statement::CreateDatabase { db_name, .. } => { + let (_, catalog_name, database) = self.resolve_catalog_and_database(db_name)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_database_created(&database) + })?; + Ok(()) + } + Statement::CreateSchema { + schema_name: SchemaName::Simple(name), + .. + } => { + let (_, catalog_name, database) = self.resolve_catalog_and_database(name)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_database_created(&database) + })?; + Ok(()) + } + Statement::CreateTable(create) if !create.temporary => { + let table_ref: TableReference = create.name.to_string().as_str().into(); + if !self.is_paimon_catalog_ref(&table_ref) { + return Ok(()); + } + let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&create.name)?; + if create.if_not_exists { + let provider = self.ctx.catalog(&catalog_name).ok_or_else(|| { + DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")) + })?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Catalog '{catalog_name}' is not a Paimon catalog" + )) + })?; + provider + .reconcile_table_created(identifier.database(), identifier.object()) + .await; + return Ok(()); + } + let options: HashMap<_, _> = extract_options(&create.table_options)? + .into_iter() + .collect(); + let declared_type = CoreOptions::new(&options) + .table_type() + .map_err(to_datafusion_error)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_table_created( + identifier.database(), + identifier.object(), + Some(declared_type), + ) + })?; + Ok(()) + } + Statement::CreateView(create) if !create.temporary => { + let table_ref: TableReference = create.name.to_string().as_str().into(); + if !self.is_paimon_catalog_ref(&table_ref) { + return Ok(()); + } + let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&create.name)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_view_created(identifier.database(), identifier.object()) + })?; + Ok(()) + } + Statement::Drop { + object_type: ObjectType::Database | ObjectType::Schema, + names, + temporary: false, + .. + } => { + for name in names { + let (_, catalog_name, database) = self.resolve_catalog_and_database(name)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_database_dropped(&database) + })?; + } + Ok(()) + } + Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + names, + temporary: false, + .. + } => { + for name in names { + let table_ref: TableReference = name.to_string().as_str().into(); + if !self.is_paimon_catalog_ref(&table_ref) { + continue; + } + let (_, catalog_name, identifier) = self.resolve_catalog_and_table(name)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_object_dropped(identifier.database(), identifier.object()) + })?; + } + Ok(()) + } + _ => Ok(()), + } + } + + fn update_catalog_metadata( + &self, + catalog_name: &str, + update: impl FnOnce(&crate::catalog::PaimonCatalogProvider) -> T, + ) -> DFResult { + let provider = self + .ctx + .catalog(catalog_name) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")))?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog_name}' is not a Paimon catalog")) + })?; + Ok(update(provider)) + } + + async fn refresh_catalog_database_metadata( + &self, + catalog_name: &str, + database: &str, + ) -> DFResult<()> { + let provider = self + .ctx + .catalog(catalog_name) + .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")))?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!("Catalog '{catalog_name}' is not a Paimon catalog")) + })?; + provider.refresh_database_metadata(database).await + } + + async fn refresh_metadata_targets( + &self, + targets: impl IntoIterator, + ) -> DFResult<()> { + let targets: HashSet<_> = targets.into_iter().collect(); + let full_catalogs: HashSet<_> = targets + .iter() + .filter_map(|target| match target { + MetadataRefreshTarget::Catalog(catalog) => Some(catalog.clone()), + MetadataRefreshTarget::Database { .. } => None, + }) + .collect(); + + for target in targets { + let (catalog_name, database) = match &target { + MetadataRefreshTarget::Catalog(catalog) => (catalog.as_str(), None), + MetadataRefreshTarget::Database { catalog, database } => { + if full_catalogs.contains(catalog.as_str()) { + continue; + } + (catalog.as_str(), Some(database.as_str())) + } + }; + let provider = self.ctx.catalog(catalog_name).ok_or_else(|| { + DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'")) + })?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Catalog '{catalog_name}' is not a Paimon catalog" + )) + })?; + if let Some(database) = database { + provider.refresh_database_metadata(database).await?; + self.missing_object_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert( + (catalog_name.to_string(), database.to_string()), + Instant::now(), + ); + } else { + provider.refresh_metadata_for_sql().await?; + self.catalog_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(catalog_name.to_string(), Instant::now()); + self.catalog_refresh_failures + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(catalog_name); + } } + Ok(()) + } + + fn missing_object_refresh_is_recent(&self, target: &MetadataRefreshTarget) -> bool { + let MetadataRefreshTarget::Database { catalog, database } = target else { + return false; + }; + self.missing_object_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&(catalog.clone(), database.clone())) + .is_some_and(|refreshed| refreshed.elapsed() < MISSING_OBJECT_REFRESH_TTL) + } + + fn catalog_refresh_is_recent(&self, catalog: &str) -> bool { + let recently_succeeded = self + .catalog_refreshes + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(catalog) + .is_some_and(|refreshed| refreshed.elapsed() < self.catalog_metadata_refresh_ttl); + recently_succeeded + || self + .catalog_refresh_failures + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(catalog) + .is_some_and(|failed| failed.elapsed() < CATALOG_METADATA_REFRESH_FAILURE_BACKOFF) + } + + fn record_catalog_refresh_failure(&self, target: &MetadataRefreshTarget) { + let MetadataRefreshTarget::Catalog(catalog) = target else { + return; + }; + self.catalog_refresh_failures + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(catalog.clone(), Instant::now()); + } + + fn metadata_refresh_gate(&self, target: &MetadataRefreshTarget) -> Arc> { + let mut gates = self + .metadata_refresh_gates + .lock() + .unwrap_or_else(|e| e.into_inner()); + gates.retain(|_, gate| gate.strong_count() > 0); + if let Some(gate) = gates.get(target).and_then(Weak::upgrade) { + return gate; + } + let gate = Arc::new(tokio::sync::Mutex::new(())); + gates.insert(target.clone(), Arc::downgrade(&gate)); + gate } /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). @@ -1304,7 +1870,7 @@ impl SQLContext { } else { Self::ensure_main_branch_write_target(name, "ALTER TABLE")?; } - let identifier = self.resolve_table_name(name)?; + let (_, catalog_name, identifier) = self.resolve_catalog_and_table(name)?; if operations.len() > 1 && operations @@ -1433,10 +1999,37 @@ impl SQLContext { } if let Some(new_identifier) = rename_to { - catalog - .rename_table(&identifier, &new_identifier, if_exists) + match catalog + .rename_table(&identifier, &new_identifier, false) .await - .map_err(to_datafusion_error)?; + { + Ok(()) => { + let updated = self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_table_renamed( + identifier.database(), + identifier.object(), + new_identifier.object(), + ) + })?; + if !updated { + if let Err(error) = self + .refresh_catalog_database_metadata(&catalog_name, identifier.database()) + .await + { + log::warn!( + "unable to reconcile metadata after renaming '{}': {error}", + identifier.full_name() + ); + } + } + } + Err(paimon::Error::TableNotExist { .. }) if if_exists => { + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_object_dropped(identifier.database(), identifier.object()) + })?; + } + Err(error) => return Err(to_datafusion_error(error)), + } } if !changes.is_empty() { @@ -3905,7 +4498,24 @@ mod tests { #[async_trait] impl Catalog for MockCatalog { async fn list_databases(&self) -> paimon::Result> { - Ok(vec![]) + let mut databases = vec!["default".to_string()]; + databases.extend( + self.functions + .lock() + .unwrap() + .keys() + .map(|identifier| identifier.database().to_string()), + ); + databases.extend( + self.views + .lock() + .unwrap() + .keys() + .map(|identifier| identifier.database().to_string()), + ); + databases.sort_unstable(); + databases.dedup(); + Ok(databases) } async fn create_database( &self, @@ -6053,11 +6663,8 @@ mod tests { assert_eq!( catalog.table_identifiers(), - vec![ - "analytics.vector_search", - "analytics.documents" - ], - "DataFusion may preload the UDTF name, then the function must resolve its table argument" + vec!["analytics.documents"], + "synchronous catalog callbacks must not probe the remote catalog for the UDTF name" ); assert_eq!( catalog.get_view_count(), @@ -7028,7 +7635,7 @@ mod tests { ignore_if_not_exists, } = &calls[0] { - assert!(ignore_if_not_exists); + assert!(!ignore_if_not_exists); assert_eq!(from.object(), "t1"); assert_eq!(to.object(), "t2"); } else { diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index f867b07d8..0b476b73f 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -757,13 +757,15 @@ async fn test_query_via_catalog_provider() { #[tokio::test] async fn test_missing_database_returns_no_schema() { let catalog = create_catalog(); - let provider = PaimonCatalogProvider::new( + let provider = PaimonCatalogProvider::try_new( None, Arc::new(catalog), Default::default(), Default::default(), None, - ); + ) + .await + .unwrap(); assert!( provider.schema("definitely_missing_database").is_none(), diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 2e550631c..aa20c5ab6 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -22,7 +22,7 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use datafusion::arrow::array::{Array, Int64Array}; -use datafusion::catalog::CatalogProvider; +use datafusion::catalog::{CatalogProvider, SchemaProvider}; use datafusion::datasource::MemTable; use paimon::catalog::{list_partitions_from_file_system, Identifier}; use paimon::spec::{ @@ -32,8 +32,9 @@ use paimon::spec::{ }; use paimon::table::{BranchManager, SnapshotManager, TagManager}; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; -use paimon_datafusion::{PaimonCatalogProvider, SQLContext}; +use paimon_datafusion::{PaimonCatalogProvider, PaimonSchemaProvider, SQLContext}; use tempfile::TempDir; +use tokio::sync::Notify; fn create_test_env() -> (TempDir, Arc) { let temp_dir = TempDir::new().expect("Failed to create temp dir"); @@ -52,24 +53,220 @@ async fn create_sql_context(catalog: Arc) -> SQLContext { struct MetadataListingCatalog { get_table_calls: AtomicUsize, + metadata_calls: AtomicUsize, + reject_remote_calls: AtomicBool, + fail_list_tables: AtomicBool, + fail_list_tables_database: Mutex>, + database_names: Mutex>, + list_tables_by_database: Mutex>, + race_refreshes: AtomicBool, + racing_list_tables_calls: AtomicUsize, + stale_refresh_started: Notify, + release_stale_refresh: Notify, + require_parallel_object_listing: AtomicBool, + view_listing_started: Notify, + require_parallel_database_listing: AtomicBool, + database_listings_started: AtomicUsize, + parallel_database_listing_started: Notify, + block_next_list_tables: AtomicBool, + block_list_tables_database: Mutex>, + blocked_list_tables_started: Notify, + release_blocked_list_tables: Notify, + table_names: Mutex>, + table_names_by_database: Mutex>>, + list_table_types_calls: AtomicUsize, + fail_next_list_table_types: AtomicBool, + fail_all_list_table_types: AtomicBool, + listing_concurrency: Option>, + block_drop_response: AtomicBool, + drop_committed: Notify, + release_drop_response: Notify, + rename_source_missing: AtomicBool, + rename_ignore_flags: Mutex>, +} + +#[derive(Default)] +struct MetadataListingConcurrency { + active: AtomicUsize, + maximum: AtomicUsize, +} + +impl MetadataListingConcurrency { + fn reset(&self) { + assert_eq!(self.active.load(Ordering::SeqCst), 0); + self.maximum.store(0, Ordering::SeqCst); + } + + fn maximum(&self) -> usize { + self.maximum.load(Ordering::SeqCst) + } + + async fn observe(&self) { + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.maximum.fetch_max(active, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + self.active.fetch_sub(1, Ordering::SeqCst); + } } impl MetadataListingCatalog { fn new() -> Self { Self { get_table_calls: AtomicUsize::new(0), + metadata_calls: AtomicUsize::new(0), + reject_remote_calls: AtomicBool::new(false), + fail_list_tables: AtomicBool::new(false), + fail_list_tables_database: Mutex::new(None), + database_names: Mutex::new(vec!["default".to_string()]), + list_tables_by_database: Mutex::new(std::collections::HashMap::new()), + race_refreshes: AtomicBool::new(false), + racing_list_tables_calls: AtomicUsize::new(0), + stale_refresh_started: Notify::new(), + release_stale_refresh: Notify::new(), + require_parallel_object_listing: AtomicBool::new(false), + view_listing_started: Notify::new(), + require_parallel_database_listing: AtomicBool::new(false), + database_listings_started: AtomicUsize::new(0), + parallel_database_listing_started: Notify::new(), + block_next_list_tables: AtomicBool::new(false), + block_list_tables_database: Mutex::new(None), + blocked_list_tables_started: Notify::new(), + release_blocked_list_tables: Notify::new(), + table_names: Mutex::new(vec!["metadata_only".to_string()]), + table_names_by_database: Mutex::new(std::collections::HashMap::new()), + list_table_types_calls: AtomicUsize::new(0), + fail_next_list_table_types: AtomicBool::new(false), + fail_all_list_table_types: AtomicBool::new(false), + listing_concurrency: None, + block_drop_response: AtomicBool::new(false), + drop_committed: Notify::new(), + release_drop_response: Notify::new(), + rename_source_missing: AtomicBool::new(false), + rename_ignore_flags: Mutex::new(Vec::new()), + } + } + + fn with_databases(databases: Vec<&str>) -> Self { + let catalog = Self::new(); + *catalog.database_names.lock().unwrap() = + databases.into_iter().map(ToString::to_string).collect(); + catalog + } + + fn with_listing_concurrency(listing_concurrency: Arc) -> Self { + Self { + listing_concurrency: Some(listing_concurrency), + ..Self::new() } } + fn set_databases(&self, databases: Vec<&str>) { + *self.database_names.lock().unwrap() = + databases.into_iter().map(ToString::to_string).collect(); + } + fn get_table_calls(&self) -> usize { self.get_table_calls.load(Ordering::SeqCst) } + + fn metadata_calls(&self) -> usize { + self.metadata_calls.load(Ordering::SeqCst) + } + + fn reject_remote_calls(&self) { + self.reject_remote_calls.store(true, Ordering::SeqCst); + } + + fn set_table_names(&self, names: Vec<&str>) { + *self.table_names.lock().unwrap() = names.into_iter().map(ToString::to_string).collect(); + } + + fn set_table_names_for(&self, database: &str, names: Vec<&str>) { + self.table_names_by_database.lock().unwrap().insert( + database.to_string(), + names.into_iter().map(ToString::to_string).collect(), + ); + } + + fn fail_list_tables(&self) { + self.fail_list_tables.store(true, Ordering::SeqCst); + } + + fn fail_list_tables_for(&self, database: &str) { + *self.fail_list_tables_database.lock().unwrap() = Some(database.to_string()); + } + + fn list_tables_calls_for(&self, database: &str) -> usize { + self.list_tables_by_database + .lock() + .unwrap() + .get(database) + .copied() + .unwrap_or_default() + } + + fn list_table_types_calls(&self) -> usize { + self.list_table_types_calls.load(Ordering::SeqCst) + } + + fn fail_next_list_table_types(&self) { + self.fail_next_list_table_types + .store(true, Ordering::SeqCst); + } + + fn fail_all_list_table_types(&self) { + self.fail_all_list_table_types.store(true, Ordering::SeqCst); + } + + fn race_next_refreshes(&self) { + self.racing_list_tables_calls.store(0, Ordering::SeqCst); + self.race_refreshes.store(true, Ordering::SeqCst); + } + + fn require_parallel_object_listing(&self) { + self.require_parallel_object_listing + .store(true, Ordering::SeqCst); + } + + fn require_parallel_database_listing(&self) { + self.require_parallel_database_listing + .store(true, Ordering::SeqCst); + } + + fn block_next_list_tables(&self) { + self.block_next_list_tables.store(true, Ordering::SeqCst); + } + + fn block_next_list_tables_for(&self, database: &str) { + *self.block_list_tables_database.lock().unwrap() = Some(database.to_string()); + } + + fn block_next_drop_response(&self) { + self.block_drop_response.store(true, Ordering::SeqCst); + } + + fn mark_rename_source_missing(&self) { + self.rename_source_missing.store(true, Ordering::SeqCst); + } + + fn rename_ignore_flags(&self) -> Vec { + self.rename_ignore_flags.lock().unwrap().clone() + } + + fn record_remote_call(&self) { + assert!( + !self.reject_remote_calls.load(Ordering::SeqCst), + "synchronous provider callback accessed the remote catalog" + ); + self.metadata_calls.fetch_add(1, Ordering::SeqCst); + } } #[async_trait] impl Catalog for MetadataListingCatalog { async fn list_databases(&self) -> paimon::Result> { - Ok(vec!["default".to_string()]) + self.record_remote_call(); + Ok(self.database_names.lock().unwrap().clone()) } async fn create_database( @@ -82,6 +279,7 @@ impl Catalog for MetadataListingCatalog { } async fn get_database(&self, name: &str) -> paimon::Result { + self.record_remote_call(); Ok(paimon::catalog::Database::new( name.to_string(), std::collections::HashMap::new(), @@ -95,21 +293,134 @@ impl Catalog for MetadataListingCatalog { _ignore_if_not_exists: bool, _cascade: bool, ) -> paimon::Result<()> { + if self.block_drop_response.swap(false, Ordering::SeqCst) { + self.drop_committed.notify_one(); + self.release_drop_response.notified().await; + } Ok(()) } async fn get_table(&self, _identifier: &Identifier) -> paimon::Result { + self.record_remote_call(); self.get_table_calls.fetch_add(1, Ordering::SeqCst); Err(paimon::Error::Unsupported { message: "table loading is unavailable".to_string(), }) } - async fn list_tables(&self, _database_name: &str) -> paimon::Result> { - Ok(vec!["metadata_only".to_string()]) + async fn list_tables(&self, database_name: &str) -> paimon::Result> { + self.record_remote_call(); + if let Some(listing_concurrency) = &self.listing_concurrency { + listing_concurrency.observe().await; + } + *self + .list_tables_by_database + .lock() + .unwrap() + .entry(database_name.to_string()) + .or_default() += 1; + if self.fail_list_tables.load(Ordering::SeqCst) + || self.fail_list_tables_database.lock().unwrap().as_deref() == Some(database_name) + { + return Err(paimon::Error::Unsupported { + message: "simulated metadata refresh failure".to_string(), + }); + } + if self.race_refreshes.load(Ordering::SeqCst) { + let call = self.racing_list_tables_calls.fetch_add(1, Ordering::SeqCst); + if call == 0 { + self.stale_refresh_started.notify_one(); + self.release_stale_refresh.notified().await; + return Ok(vec!["stale".to_string()]); + } + return Ok(vec!["fresh".to_string()]); + } + if self.require_parallel_object_listing.load(Ordering::SeqCst) { + tokio::time::timeout( + std::time::Duration::from_millis(100), + self.view_listing_started.notified(), + ) + .await + .map_err(|_| paimon::Error::Unsupported { + message: "list_views did not run concurrently".to_string(), + })?; + } + if self + .require_parallel_database_listing + .load(Ordering::SeqCst) + { + let started = self + .database_listings_started + .fetch_add(1, Ordering::SeqCst) + + 1; + if started == 1 { + tokio::time::timeout( + std::time::Duration::from_millis(100), + self.parallel_database_listing_started.notified(), + ) + .await + .map_err(|_| paimon::Error::Unsupported { + message: "databases were not listed concurrently".to_string(), + })?; + } else { + self.parallel_database_listing_started.notify_waiters(); + } + } + let table_names = self + .table_names_by_database + .lock() + .unwrap() + .get(database_name) + .cloned() + .unwrap_or_else(|| self.table_names.lock().unwrap().clone()); + let block_target = { + let mut block_database = self.block_list_tables_database.lock().unwrap(); + if block_database.as_deref() == Some(database_name) { + block_database.take(); + true + } else { + false + } + }; + let should_block = + block_target || self.block_next_list_tables.swap(false, Ordering::SeqCst); + if should_block { + self.blocked_list_tables_started.notify_one(); + self.release_blocked_list_tables.notified().await; + } + Ok(table_names) + } + + async fn list_table_types( + &self, + _database_name: &str, + table_names: &[String], + ) -> paimon::Result> { + self.list_table_types_calls.fetch_add(1, Ordering::SeqCst); + if let Some(listing_concurrency) = &self.listing_concurrency { + listing_concurrency.observe().await; + } + if self.fail_all_list_table_types.load(Ordering::SeqCst) + || self + .fail_next_list_table_types + .swap(false, Ordering::SeqCst) + { + return Err(paimon::Error::Unsupported { + message: "simulated table type classification failure".to_string(), + }); + } + Ok(table_names + .iter() + .cloned() + .map(|name| (name, paimon::spec::TableType::Table)) + .collect()) } async fn list_views(&self, _database_name: &str) -> paimon::Result> { + self.record_remote_call(); + if self.require_parallel_object_listing.load(Ordering::SeqCst) { + self.view_listing_started.notify_one(); + } Ok(vec!["metadata_view".to_string()]) } @@ -127,15 +438,43 @@ impl Catalog for MetadataListingCatalog { _identifier: &Identifier, _ignore_if_not_exists: bool, ) -> paimon::Result<()> { + if self.block_drop_response.swap(false, Ordering::SeqCst) { + self.drop_committed.notify_one(); + self.release_drop_response.notified().await; + } Ok(()) } async fn rename_table( &self, - _from: &Identifier, - _to: &Identifier, - _ignore_if_not_exists: bool, + from: &Identifier, + to: &Identifier, + ignore_if_not_exists: bool, ) -> paimon::Result<()> { + self.rename_ignore_flags + .lock() + .unwrap() + .push(ignore_if_not_exists); + if self.rename_source_missing.load(Ordering::SeqCst) { + if ignore_if_not_exists { + return Ok(()); + } + return Err(paimon::Error::TableNotExist { + full_name: from.full_name(), + }); + } + let mut table_names_by_database = self.table_names_by_database.lock().unwrap(); + if let Some(names) = table_names_by_database.get_mut(from.database()) { + if let Some(index) = names.iter().position(|name| name == from.object()) { + names[index] = to.object().to_string(); + } + return Ok(()); + } + drop(table_names_by_database); + let mut names = self.table_names.lock().unwrap(); + if let Some(index) = names.iter().position(|name| name == from.object()) { + names[index] = to.object().to_string(); + } Ok(()) } @@ -149,279 +488,1563 @@ impl Catalog for MetadataListingCatalog { } } -struct PartitionCatalog { - inner: Arc, - fail_list_partitions: AtomicBool, - unknown_statistics: AtomicBool, - partition_identifiers: Mutex>, +#[tokio::test] +async fn test_refreshed_catalog_callbacks_do_not_access_remote_catalog() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new_uninitialized( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ); + + provider.refresh_metadata().await.unwrap(); + let calls_after_refresh = catalog.metadata_calls(); + catalog.reject_remote_calls(); + + assert_eq!(provider.schema_names(), vec!["default"]); + let schema = provider.schema("default").unwrap(); + assert_eq!(schema.table_names(), vec!["metadata_only", "metadata_view"]); + assert!(schema.table_exist("metadata_only")); + assert!(schema.table_exist("metadata_view")); + assert!(!schema.table_exist("missing")); + assert_eq!(catalog.metadata_calls(), calls_after_refresh); } -impl PartitionCatalog { - fn new(inner: Arc) -> Self { - Self { - inner, - fail_list_partitions: AtomicBool::new(false), - unknown_statistics: AtomicBool::new(false), - partition_identifiers: Mutex::new(Vec::new()), - } - } +#[tokio::test] +async fn test_public_catalog_constructor_returns_ready_provider() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); - fn set_fail_list_partitions(&self, fail: bool) { - self.fail_list_partitions.store(fail, Ordering::SeqCst); - } + assert_eq!(provider.schema_names(), vec!["default"]); + assert_eq!( + provider.schema("default").unwrap().table_names(), + vec!["metadata_only", "metadata_view"] + ); +} - /// Report every statistic as never measured, the way a catalog does for a partition that was - /// registered but never had statistics reported to it. - fn set_unknown_statistics(&self, unknown: bool) { - self.unknown_statistics.store(unknown, Ordering::SeqCst); - } +#[tokio::test] +async fn test_metadata_refresh_classifies_only_new_objects() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["existing"]); + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + assert_eq!(catalog.list_table_types_calls(), 1); - fn take_partition_identifiers(&self) -> Vec { - std::mem::take(&mut *self.partition_identifiers.lock().unwrap()) - } + provider.refresh_metadata().await.unwrap(); + assert_eq!(catalog.list_table_types_calls(), 1); + + catalog.set_table_names(vec!["existing", "discovered"]); + provider.refresh_metadata().await.unwrap(); + assert_eq!(catalog.list_table_types_calls(), 2); + provider.refresh_metadata().await.unwrap(); + assert_eq!(catalog.list_table_types_calls(), 2); } -#[async_trait] -impl Catalog for PartitionCatalog { - async fn list_databases(&self) -> paimon::Result> { - self.inner.list_databases().await - } +#[tokio::test] +async fn test_metadata_refresh_deduplicates_object_names_before_classification() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["duplicate", "duplicate"]); - async fn create_database( - &self, - name: &str, - ignore_if_exists: bool, - properties: std::collections::HashMap, - ) -> paimon::Result<()> { - self.inner - .create_database(name, ignore_if_exists, properties) - .await - } + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); - async fn get_database(&self, name: &str) -> paimon::Result { - self.inner.get_database(name).await - } + assert_eq!(catalog.list_table_types_calls(), 1); + let names = provider.schema("default").unwrap().table_names(); + assert_eq!(names.iter().filter(|name| *name == "duplicate").count(), 1); +} - async fn drop_database( - &self, - name: &str, - ignore_if_not_exists: bool, - cascade: bool, - ) -> paimon::Result<()> { - self.inner - .drop_database(name, ignore_if_not_exists, cascade) - .await - } +#[tokio::test] +async fn test_metadata_refresh_retries_unknown_table_types() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["existing"]); + catalog.fail_next_list_table_types(); + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(!schema.table_exist("existing$snapshots")); + assert_eq!(catalog.list_table_types_calls(), 1); - async fn get_table(&self, identifier: &Identifier) -> paimon::Result { - self.inner.get_table(identifier).await - } + provider.refresh_metadata().await.unwrap(); - async fn list_tables(&self, database_name: &str) -> paimon::Result> { - self.inner.list_tables(database_name).await - } + assert_eq!(catalog.list_table_types_calls(), 2); + let schema = provider.schema("default").unwrap(); + assert!(schema.table_exist("existing$snapshots")); +} - async fn create_table( - &self, - identifier: &Identifier, - creation: paimon::spec::Schema, - ignore_if_exists: bool, - ) -> paimon::Result<()> { - self.inner - .create_table(identifier, creation, ignore_if_exists) - .await - } +#[tokio::test] +async fn test_metadata_refresh_limits_and_reports_persistent_classification_failures() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let table_names: Vec<_> = (0..100).map(|index| format!("table_{index}")).collect(); + *catalog.table_names.lock().unwrap() = table_names; + catalog.fail_all_list_table_types(); - async fn drop_table( - &self, - identifier: &Identifier, - ignore_if_not_exists: bool, - ) -> paimon::Result<()> { - self.inner - .drop_table(identifier, ignore_if_not_exists) - .await - } + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); - async fn rename_table( - &self, - from: &Identifier, - to: &Identifier, - ignore_if_not_exists: bool, - ) -> paimon::Result<()> { - self.inner - .rename_table(from, to, ignore_if_not_exists) - .await - } + assert_eq!(provider.schema("default").unwrap().table_names().len(), 101); + assert!(catalog.list_table_types_calls() <= 16); - async fn alter_table( - &self, - identifier: &Identifier, - changes: Vec, - ignore_if_not_exists: bool, - ) -> paimon::Result<()> { - self.inner - .alter_table(identifier, changes, ignore_if_not_exists) - .await - } + let calls_before_refresh = catalog.list_table_types_calls(); + let error = provider.refresh_metadata().await.unwrap_err(); + assert!(error.to_string().contains("table type classification")); + assert!(catalog.list_table_types_calls() - calls_before_refresh <= 16); +} - async fn list_partitions( - &self, - identifier: &Identifier, - ) -> paimon::Result> { - self.partition_identifiers - .lock() +#[tokio::test] +async fn test_explicit_uninitialized_providers_require_metadata_initialization() { + let catalog: Arc = Arc::new(MetadataListingCatalog::new()); + let catalog_provider = PaimonCatalogProvider::new_uninitialized( + Some("paimon".to_string()), + Arc::clone(&catalog), + Default::default(), + Default::default(), + None, + ); + assert!(catalog_provider.schema_names().is_empty()); + catalog_provider.initialize_metadata().await.unwrap(); + assert_eq!(catalog_provider.schema_names(), vec!["default"]); + + let schema_provider = PaimonSchemaProvider::new_uninitialized( + Some("paimon".to_string()), + catalog, + "default".to_string(), + Default::default(), + None, + Default::default(), + None, + ); + assert!(schema_provider.table_names().is_empty()); + schema_provider.initialize_metadata().await.unwrap(); + assert_eq!( + schema_provider.table_names(), + vec!["metadata_only", "metadata_view"] + ); +} + +#[tokio::test] +async fn test_explicit_metadata_refresh_times_out_pending_remote_listing() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new_uninitialized( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .with_metadata_refresh_timeout(std::time::Duration::from_millis(20)); + catalog.block_next_list_tables(); + + let error = provider.refresh_metadata().await.unwrap_err(); + + assert!(error.to_string().contains("metadata refresh timed out")); +} + +#[tokio::test] +async fn test_failed_catalog_refresh_preserves_last_good_snapshot() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + + catalog.set_table_names(vec!["not_committed"]); + catalog.fail_list_tables(); + assert!(provider.refresh_metadata().await.is_err()); + + assert_eq!(provider.schema_names(), vec!["default"]); + let schema = provider.schema("default").unwrap(); + assert_eq!(schema.table_names(), vec!["metadata_only", "metadata_view"]); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_older_refresh_cannot_overwrite_newer_snapshot() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = Arc::new( + PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(), + ); + catalog.race_next_refreshes(); + + let stale_provider = Arc::clone(&provider); + let stale_refresh = tokio::spawn(async move { stale_provider.refresh_metadata().await }); + catalog.stale_refresh_started.notified().await; + + let fresh_provider = Arc::clone(&provider); + tokio::spawn(async move { fresh_provider.refresh_metadata().await }) + .await + .unwrap() + .unwrap(); + + catalog.release_stale_refresh.notify_one(); + stale_refresh.await.unwrap().unwrap(); + + let schema = provider.schema("default").unwrap(); + assert_eq!(schema.table_names(), vec!["fresh", "metadata_view"]); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_full_refresh_merges_non_conflicting_database_updates() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "first", "second", + ])); + catalog.set_table_names_for("first", vec!["first_old"]); + catalog.set_table_names_for("second", vec!["second_old"]); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.set_table_names_for("first", vec!["first_new"]); + catalog.set_table_names_for("second", vec!["second_stale"]); + catalog.block_next_list_tables_for("second"); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let full_refresh = tokio::spawn(async move { + provider + .downcast_ref::() .unwrap() - .push(identifier.clone()); + .refresh_metadata() + .await + }); + catalog.blocked_list_tables_started.notified().await; - let Some(branch) = identifier.branch_name()? else { - let mut partitions = self.inner.list_partitions(identifier).await?; - if self.unknown_statistics.load(Ordering::SeqCst) { - for partition in &mut partitions { - partition.record_count = paimon::spec::Partition::UNKNOWN; - partition.file_size_in_bytes = paimon::spec::Partition::UNKNOWN; - partition.file_count = paimon::spec::Partition::UNKNOWN; - partition.last_file_creation_time = paimon::spec::Partition::UNKNOWN; - partition.total_buckets = paimon::spec::Partition::UNKNOWN_TOTAL_BUCKETS; - } - } - return Ok(partitions); - }; - if self.fail_list_partitions.load(Ordering::SeqCst) { - return Err(paimon::Error::Unsupported { - message: "injected list_partitions failure".to_string(), - }); - } + catalog.set_table_names_for("second", vec!["second_new"]); + assert!(sql_context + .sql("SELECT * FROM paimon.second.second_new") + .await + .is_err()); + catalog.release_blocked_list_tables.notify_one(); + full_refresh.await.unwrap().unwrap(); - let base = Identifier::new(identifier.database(), identifier.table_name()?); - let table = self.inner.get_table(&base).await?; - let table = table.copy_with_branch(&branch).await?; - let mut partitions = list_partitions_from_file_system(&table).await?; - for partition in &mut partitions { - partition.created_by = Some("catalog".to_string()); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert_eq!( + provider.schema("first").unwrap().table_names(), + vec!["first_new", "metadata_view"] + ); + assert_eq!( + provider.schema("second").unwrap().table_names(), + vec!["second_new", "metadata_view"] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_full_refresh_retries_database_rejected_by_ddl_delta() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["existing"]); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.set_table_names(vec!["existing", "discovered"]); + catalog.block_next_list_tables_for("default"); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let full_refresh = tokio::spawn(async move { + provider + .downcast_ref::() + .unwrap() + .refresh_metadata() + .await + }); + catalog.blocked_list_tables_started.notified().await; + + sql_context + .sql("CREATE TABLE created (id BIGINT)") + .await + .unwrap(); + catalog.set_table_names(vec!["existing", "discovered", "created"]); + catalog.release_blocked_list_tables.notify_one(); + full_refresh.await.unwrap().unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(schema.table_exist("created")); + assert!(schema.table_exist("discovered")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_newer_full_refresh_advances_database_tombstone() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "first", "second", + ])); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + sql_context.sql("DROP DATABASE second").await.unwrap(); + + catalog.block_next_list_tables_for("second"); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let stale_refresh = tokio::spawn(async move { + provider + .downcast_ref::() + .unwrap() + .refresh_metadata() + .await + }); + catalog.blocked_list_tables_started.notified().await; + + catalog.set_databases(vec!["first"]); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + provider + .downcast_ref::() + .unwrap() + .refresh_metadata() + .await + .unwrap(); + catalog.release_blocked_list_tables.notify_one(); + stale_refresh.await.unwrap().unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert!(provider.schema("second").is_none()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_stale_refresh_cannot_overwrite_local_schema_registration() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = Arc::new( + PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(), + ); + catalog.race_next_refreshes(); + + let stale_provider = Arc::clone(&provider); + let stale_refresh = tokio::spawn(async move { stale_provider.refresh_metadata().await }); + catalog.stale_refresh_started.notified().await; + + provider + .register_schema( + "local", + Arc::new(datafusion::catalog::MemorySchemaProvider::new()), + ) + .unwrap(); + assert!(provider.schema("local").is_some()); + + catalog.release_stale_refresh.notify_one(); + stale_refresh.await.unwrap().unwrap(); + + assert!(provider.schema("local").is_some()); +} + +#[tokio::test] +async fn test_refresh_lists_tables_and_views_concurrently() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.require_parallel_object_listing(); + + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + + assert_eq!( + provider.schema("default").unwrap().table_names(), + vec!["metadata_only", "metadata_view"] + ); +} + +#[tokio::test] +async fn test_refresh_lists_databases_concurrently() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "first", "second", + ])); + catalog.require_parallel_database_listing(); + + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + + assert_eq!(provider.schema_names(), vec!["first", "second"]); +} + +#[tokio::test] +async fn test_temp_table_registration_does_not_access_remote_catalog() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new_uninitialized( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ); + provider.refresh_metadata().await.unwrap(); + let calls_after_refresh = catalog.metadata_calls(); + catalog.reject_remote_calls(); + + let table = MemTable::try_new(Arc::new(Schema::empty()), vec![vec![]]).unwrap(); + provider + .register_temp_table("default", "metadata_only", Arc::new(table)) + .unwrap(); + + assert_eq!(catalog.metadata_calls(), calls_after_refresh); +} + +struct PartitionCatalog { + inner: Arc, + fail_list_partitions: AtomicBool, + unknown_statistics: AtomicBool, + partition_identifiers: Mutex>, +} + +impl PartitionCatalog { + fn new(inner: Arc) -> Self { + Self { + inner, + fail_list_partitions: AtomicBool::new(false), + unknown_statistics: AtomicBool::new(false), + partition_identifiers: Mutex::new(Vec::new()), } - Ok(partitions) + } + + fn set_fail_list_partitions(&self, fail: bool) { + self.fail_list_partitions.store(fail, Ordering::SeqCst); + } + + /// Report every statistic as never measured, the way a catalog does for a partition that was + /// registered but never had statistics reported to it. + fn set_unknown_statistics(&self, unknown: bool) { + self.unknown_statistics.store(unknown, Ordering::SeqCst); + } + + fn take_partition_identifiers(&self) -> Vec { + std::mem::take(&mut *self.partition_identifiers.lock().unwrap()) } } -async fn collect_ids(sql_context: &SQLContext, sql: &str) -> Vec { - let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - let mut ids = Vec::new(); - for batch in batches { - let id_array = batch - .column_by_name("id") - .and_then(|c| c.as_any().downcast_ref::()) - .expect("id column"); - for row in 0..batch.num_rows() { - ids.push(id_array.value(row)); - } +#[async_trait] +impl Catalog for PartitionCatalog { + async fn list_databases(&self) -> paimon::Result> { + self.inner.list_databases().await } - ids.sort_unstable(); - ids + + async fn create_database( + &self, + name: &str, + ignore_if_exists: bool, + properties: std::collections::HashMap, + ) -> paimon::Result<()> { + self.inner + .create_database(name, ignore_if_exists, properties) + .await + } + + async fn get_database(&self, name: &str) -> paimon::Result { + self.inner.get_database(name).await + } + + async fn drop_database( + &self, + name: &str, + ignore_if_not_exists: bool, + cascade: bool, + ) -> paimon::Result<()> { + self.inner + .drop_database(name, ignore_if_not_exists, cascade) + .await + } + + async fn get_table(&self, identifier: &Identifier) -> paimon::Result { + self.inner.get_table(identifier).await + } + + async fn list_tables(&self, database_name: &str) -> paimon::Result> { + self.inner.list_tables(database_name).await + } + + async fn create_table( + &self, + identifier: &Identifier, + creation: paimon::spec::Schema, + ignore_if_exists: bool, + ) -> paimon::Result<()> { + self.inner + .create_table(identifier, creation, ignore_if_exists) + .await + } + + async fn drop_table( + &self, + identifier: &Identifier, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .drop_table(identifier, ignore_if_not_exists) + .await + } + + async fn rename_table( + &self, + from: &Identifier, + to: &Identifier, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .rename_table(from, to, ignore_if_not_exists) + .await + } + + async fn alter_table( + &self, + identifier: &Identifier, + changes: Vec, + ignore_if_not_exists: bool, + ) -> paimon::Result<()> { + self.inner + .alter_table(identifier, changes, ignore_if_not_exists) + .await + } + + async fn list_partitions( + &self, + identifier: &Identifier, + ) -> paimon::Result> { + self.partition_identifiers + .lock() + .unwrap() + .push(identifier.clone()); + + let Some(branch) = identifier.branch_name()? else { + let mut partitions = self.inner.list_partitions(identifier).await?; + if self.unknown_statistics.load(Ordering::SeqCst) { + for partition in &mut partitions { + partition.record_count = paimon::spec::Partition::UNKNOWN; + partition.file_size_in_bytes = paimon::spec::Partition::UNKNOWN; + partition.file_count = paimon::spec::Partition::UNKNOWN; + partition.last_file_creation_time = paimon::spec::Partition::UNKNOWN; + partition.total_buckets = paimon::spec::Partition::UNKNOWN_TOTAL_BUCKETS; + } + } + return Ok(partitions); + }; + if self.fail_list_partitions.load(Ordering::SeqCst) { + return Err(paimon::Error::Unsupported { + message: "injected list_partitions failure".to_string(), + }); + } + + let base = Identifier::new(identifier.database(), identifier.table_name()?); + let table = self.inner.get_table(&base).await?; + let table = table.copy_with_branch(&branch).await?; + let mut partitions = list_partitions_from_file_system(&table).await?; + for partition in &mut partitions { + partition.created_by = Some("catalog".to_string()); + } + Ok(partitions) + } +} + +async fn collect_ids(sql_context: &SQLContext, sql: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut ids = Vec::new(); + for batch in batches { + let id_array = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::()) + .expect("id column"); + for row in 0..batch.num_rows() { + ids.push(id_array.value(row)); + } + } + ids.sort_unstable(); + ids +} + +async fn collect_i64_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column_by_name(column) + .and_then(|c| c.as_any().downcast_ref::()) + .expect(column); + for row in 0..batch.num_rows() { + values.push(array.value(row)); + } + } + values.sort_unstable(); + values +} + +async fn collect_string_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { + let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column_by_name(column) + .and_then(|c| c.as_any().downcast_ref::()) + .expect(column); + for row in 0..batch.num_rows() { + if !array.is_null(row) { + values.push(array.value(row).to_string()); + } + } + } + values.sort_unstable(); + values +} + +async fn assert_sql_error_contains(sql_context: &SQLContext, sql: &str, expected: &str) { + let err = match sql_context.sql(sql).await { + Ok(df) => df + .collect() + .await + .expect_err("SQL should fail but succeeded") + .to_string(), + Err(err) => err.to_string(), + }; + assert!( + err.contains(expected), + "expected error containing '{expected}', got: {err}" + ); +} + +#[tokio::test] +async fn test_show_tables_is_enabled() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog).await; + + sql_context + .sql("SHOW TABLES") + .await + .expect("SHOW TABLES should be planned when information_schema is enabled") + .collect() + .await + .expect("SHOW TABLES should execute"); +} + +#[tokio::test] +async fn test_show_tables_does_not_load_table_providers() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await; + + assert!(table_names.contains(&"metadata_only".to_string())); + assert_eq!(catalog.get_table_calls(), 0); + + assert!(sql_context + .sql("SELECT * FROM metadata_only") + .await + .is_err()); + assert!(catalog.get_table_calls() > 0); +} + +#[tokio::test] +async fn test_show_tables_refreshes_catalog_metadata() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.set_table_names(vec!["added_after_registration"]); + let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await; + + assert!(table_names.contains(&"added_after_registration".to_string())); + assert!(!table_names.contains(&"metadata_only".to_string())); +} + +#[tokio::test] +async fn test_show_tables_does_not_refresh_unrelated_catalogs() { + let healthy = Arc::new(MetadataListingCatalog::new()); + let unavailable = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("healthy", healthy.clone()) + .await + .unwrap(); + sql_context + .register_catalog("unavailable", unavailable.clone()) + .await + .unwrap(); + + healthy.set_table_names(vec!["current"]); + unavailable.fail_list_tables(); + let unavailable_calls = unavailable.metadata_calls(); + + let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await; + + assert!(table_names.contains(&"current".to_string())); + assert_eq!(unavailable.metadata_calls(), unavailable_calls); +} + +#[tokio::test] +async fn test_show_tables_only_refreshes_current_database() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "default", + "unavailable", + ])); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.set_table_names(vec!["current"]); + catalog.fail_list_tables_for("unavailable"); + let unavailable_calls = catalog.list_tables_calls_for("unavailable"); + + let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await; + + assert!(table_names.contains(&"current".to_string())); + assert_eq!( + catalog.list_tables_calls_for("unavailable"), + unavailable_calls + ); +} + +#[tokio::test] +async fn test_show_columns_refreshes_qualified_database() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "default", + "analytics", + ])); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + let default_calls = catalog.list_tables_calls_for("default"); + let analytics_calls = catalog.list_tables_calls_for("analytics"); + assert!(sql_context + .sql("SHOW COLUMNS IN paimon.analytics.metadata_only") + .await + .is_err()); + + assert_eq!(catalog.list_tables_calls_for("default"), default_calls); + assert_eq!( + catalog.list_tables_calls_for("analytics"), + analytics_calls + 1 + ); +} + +#[tokio::test] +async fn test_show_functions_does_not_refresh_catalog_metadata() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.fail_list_tables(); + let list_calls = catalog.list_tables_calls_for("default"); + sql_context.sql("SHOW FUNCTIONS").await.unwrap(); + + assert_eq!(catalog.list_tables_calls_for("default"), list_calls); +} + +#[tokio::test] +async fn test_create_table_does_not_refresh_unrelated_catalogs() { + let healthy = Arc::new(MetadataListingCatalog::new()); + let unrelated = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("healthy", healthy) + .await + .unwrap(); + sql_context + .register_catalog("unrelated", unrelated.clone()) + .await + .unwrap(); + let unrelated_calls = unrelated.metadata_calls(); + + sql_context + .sql("CREATE TABLE created (id BIGINT)") + .await + .unwrap(); + + assert_eq!(unrelated.metadata_calls(), unrelated_calls); +} + +#[tokio::test] +async fn test_successful_ddl_does_not_wait_for_metadata_listing() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + catalog.block_next_list_tables(); + + tokio::time::timeout( + std::time::Duration::from_millis(100), + sql_context.sql("CREATE TABLE direct_snapshot_update (id BIGINT)"), + ) + .await + .expect("committed DDL must not wait for a metadata listing") + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert!(provider + .schema("default") + .unwrap() + .table_names() + .contains(&"direct_snapshot_update".to_string())); +} + +#[tokio::test] +async fn test_ddl_applies_exact_snapshot_deltas() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let default_list_calls = catalog.list_tables_calls_for("default"); + + sql_context.sql("CREATE DATABASE analytics").await.unwrap(); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert!(provider.schema("analytics").is_some()); + + sql_context + .sql("CREATE TABLE analytics.events (id BIGINT)") + .await + .unwrap(); + assert!(provider.schema("analytics").unwrap().table_exist("events")); + + sql_context + .sql("ALTER TABLE analytics.events RENAME TO renamed_events") + .await + .unwrap(); + let schema = provider.schema("analytics").unwrap(); + assert!(!schema.table_exist("events")); + assert!(schema.table_exist("renamed_events")); + + sql_context + .sql("DROP TABLE analytics.renamed_events") + .await + .unwrap(); + assert!(!provider + .schema("analytics") + .unwrap() + .table_exist("renamed_events")); + + sql_context.sql("DROP DATABASE analytics").await.unwrap(); + assert!(provider.schema("analytics").is_none()); + assert_eq!( + catalog.list_tables_calls_for("default"), + default_list_calls, + "DDL deltas must not trigger metadata listings" + ); +} + +#[tokio::test] +async fn test_alter_table_if_exists_does_not_create_phantom_rename_delta() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog) + .await + .unwrap(); + + sql_context + .sql("ALTER TABLE IF EXISTS missing RENAME TO phantom") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(!schema.table_exist("missing")); + assert!(!schema.table_exist("phantom")); +} + +#[tokio::test] +async fn test_alter_table_if_exists_reconciles_stale_positive_rename_source() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["source"]); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + catalog.set_table_names(vec![]); + catalog.mark_rename_source_missing(); + sql_context + .sql("ALTER TABLE IF EXISTS source RENAME TO phantom") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(!schema.table_exist("source")); + assert!(!schema.table_exist("phantom")); + assert_eq!(catalog.rename_ignore_flags(), vec![false]); +} + +#[tokio::test] +async fn test_successful_rename_refreshes_target_when_source_was_not_snapshotted() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec![]); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + catalog.set_table_names(vec!["source"]); + + sql_context + .sql("ALTER TABLE source RENAME TO destination") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(!schema.table_exist("source")); + assert!(schema.table_exist("destination")); + assert!(schema.table_exist("destination$snapshots")); +} + +#[tokio::test] +async fn test_successful_rename_records_unknown_target_before_bounded_refresh() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec![]); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_timeout(std::time::Duration::from_millis(20)) + .build(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + catalog.set_table_names(vec!["source"]); + catalog.block_next_list_tables(); + + tokio::time::timeout( + std::time::Duration::from_millis(200), + sql_context.sql("ALTER TABLE source RENAME TO destination"), + ) + .await + .expect("rename reconciliation must be time bounded") + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("default").unwrap(); + assert!(!schema.table_exist("source")); + assert!(schema.table_exist("destination")); + assert!(!schema.table_exist("destination$snapshots")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_concurrent_ddl_delta_follows_serialized_commit_order() { + let catalog = Arc::new(MetadataListingCatalog::new()); + catalog.set_table_names(vec!["raced"]); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let sql_context = Arc::new(sql_context); + catalog.block_next_drop_response(); + + let drop_context = Arc::clone(&sql_context); + let drop = tokio::spawn(async move { drop_context.sql("DROP TABLE raced").await }); + catalog.drop_committed.notified().await; + + let create_finished = Arc::new(Notify::new()); + let create_context = Arc::clone(&sql_context); + let create_finished_signal = Arc::clone(&create_finished); + let create = tokio::spawn(async move { + let result = create_context.sql("CREATE TABLE raced (id BIGINT)").await; + create_finished_signal.notify_one(); + result + }); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(200), + create_finished.notified(), + ) + .await + .is_err(), + "same-database DDL must wait for the earlier mutation" + ); + catalog.release_drop_response.notify_one(); + + drop.await.unwrap().unwrap(); + create.await.unwrap().unwrap(); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert!(provider.schema("default").unwrap().table_exist("raced")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_database_and_table_ddl_share_serialized_commit_order() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let sql_context = Arc::new(sql_context); + catalog.block_next_drop_response(); + + let drop_context = Arc::clone(&sql_context); + let drop = tokio::spawn(async move { drop_context.sql("DROP DATABASE default").await }); + catalog.drop_committed.notified().await; + + let create_finished = Arc::new(Notify::new()); + let create_context = Arc::clone(&sql_context); + let create_finished_signal = Arc::clone(&create_finished); + let create = tokio::spawn(async move { + let result = create_context + .sql("CREATE TABLE default.after_drop (id BIGINT)") + .await; + create_finished_signal.notify_one(); + result + }); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(200), + create_finished.notified(), + ) + .await + .is_err(), + "same-database table DDL must wait for database DDL" + ); + catalog.release_drop_response.notify_one(); + + drop.await.unwrap().unwrap(); + create.await.unwrap().unwrap(); + let provider = sql_context.ctx().catalog("paimon").unwrap(); + assert!(provider + .schema("default") + .unwrap() + .table_exist("after_drop")); +} + +#[tokio::test] +async fn test_repeated_missing_tables_share_negative_refresh_ttl() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let list_calls = catalog.list_tables_calls_for("default"); + + assert!(sql_context + .sql("SELECT * FROM first_missing") + .await + .is_err()); + assert!(sql_context + .sql("SELECT * FROM second_missing") + .await + .is_err()); + + assert_eq!(catalog.list_tables_calls_for("default"), list_calls + 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_concurrent_missing_tables_share_single_refresh() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let sql_context = Arc::new(sql_context); + let list_calls = catalog.list_tables_calls_for("default"); + catalog.block_next_list_tables(); + + let first_context = Arc::clone(&sql_context); + let first = tokio::spawn(async move { first_context.sql("SELECT * FROM missing_a").await }); + catalog.blocked_list_tables_started.notified().await; + + let second_context = Arc::clone(&sql_context); + let second = tokio::spawn(async move { second_context.sql("SELECT * FROM missing_b").await }); + tokio::task::yield_now().await; + catalog.release_blocked_list_tables.notify_one(); + + assert!(first.await.unwrap().is_err()); + assert!(second.await.unwrap().is_err()); + assert_eq!(catalog.list_tables_calls_for("default"), list_calls + 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_stalled_missing_table_refresh_does_not_block_metadata_free_query() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let sql_context = Arc::new(sql_context); + catalog.block_next_list_tables(); + + let missing_context = Arc::clone(&sql_context); + let missing = tokio::spawn(async move { missing_context.sql("SELECT * FROM missing").await }); + catalog.blocked_list_tables_started.notified().await; + + tokio::time::timeout( + std::time::Duration::from_millis(100), + sql_context.sql("SELECT 1"), + ) + .await + .expect("metadata-free query must not wait for an unrelated refresh") + .unwrap(); + + catalog.release_blocked_list_tables.notify_one(); + assert!(missing.await.unwrap().is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_stalled_refresh_does_not_block_another_database() { + let catalog = Arc::new(MetadataListingCatalog::with_databases(vec![ + "first", "second", + ])); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let sql_context = Arc::new(sql_context); + catalog.block_next_list_tables_for("first"); + + let first_context = Arc::clone(&sql_context); + let first = tokio::spawn(async move { + first_context + .sql("SELECT * FROM paimon.first.missing") + .await + }); + catalog.blocked_list_tables_started.notified().await; + + let second_result = tokio::time::timeout( + std::time::Duration::from_millis(100), + sql_context.sql("SELECT * FROM paimon.second.missing"), + ) + .await + .expect("a different database must not wait for the stalled refresh"); + assert!(second_result.is_err()); + + catalog.release_blocked_list_tables.notify_one(); + assert!(first.await.unwrap().is_err()); +} + +#[tokio::test] +async fn test_show_tables_preserves_catalog_view_type() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + + let table_types = collect_string_column( + &sql_context, + "SELECT table_type FROM information_schema.tables \ + WHERE table_schema = 'default' AND table_name = 'metadata_view'", + "table_type", + ) + .await; + + assert_eq!(table_types, vec!["VIEW"]); + assert_eq!(catalog.get_table_calls(), 0); +} + +#[tokio::test] +async fn test_information_schema_refreshes_all_paimon_catalogs() { + let first = Arc::new(MetadataListingCatalog::new()); + let second = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); + sql_context.register_catalog("first", first).await.unwrap(); + sql_context + .register_catalog("second", second.clone()) + .await + .unwrap(); + second.set_table_names(vec!["second_new"]); + + let table_names = collect_string_column( + &sql_context, + "SELECT table_name FROM first.information_schema.tables \ + WHERE table_catalog = 'second' AND table_schema = 'default' \ + AND table_name = 'second_new'", + "table_name", + ) + .await; + + assert_eq!(table_names, vec!["second_new"]); +} + +#[tokio::test] +async fn test_information_schema_reuses_recent_catalog_refresh() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::from_millis(200)) + .build(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + let calls_after_first_query = catalog.metadata_calls(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + + assert_eq!(catalog.metadata_calls(), calls_after_first_query); } -async fn collect_i64_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { - let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - let mut values = Vec::new(); - for batch in batches { - let array = batch - .column_by_name(column) - .and_then(|c| c.as_any().downcast_ref::()) - .expect(column); - for row in 0..batch.num_rows() { - values.push(array.value(row)); - } - } - values.sort_unstable(); - values +#[tokio::test] +async fn test_information_schema_reuses_registration_snapshot() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let calls_after_registration = catalog.metadata_calls(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + + assert_eq!(catalog.metadata_calls(), calls_after_registration); } -async fn collect_string_column(sql_context: &SQLContext, sql: &str, column: &str) -> Vec { - let batches = sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - let mut values = Vec::new(); - for batch in batches { - let array = batch - .column_by_name(column) - .and_then(|c| c.as_any().downcast_ref::()) - .expect(column); - for row in 0..batch.num_rows() { - if !array.is_null(row) { - values.push(array.value(row).to_string()); - } - } - } - values.sort_unstable(); - values +#[tokio::test] +async fn test_information_schema_keeps_last_good_snapshot_on_catalog_refresh_failure() { + let healthy = Arc::new(MetadataListingCatalog::new()); + let unavailable = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); + sql_context + .register_catalog("healthy", healthy.clone()) + .await + .unwrap(); + sql_context + .register_catalog("unavailable", unavailable.clone()) + .await + .unwrap(); + healthy.set_table_names(vec!["healthy_new"]); + unavailable.fail_list_tables(); + + let table_names = collect_string_column( + &sql_context, + "SELECT table_name FROM healthy.information_schema.tables \ + WHERE table_catalog = 'healthy' AND table_schema = 'default' \ + AND table_name = 'healthy_new'", + "table_name", + ) + .await; + + assert_eq!(table_names, vec!["healthy_new"]); } -async fn assert_sql_error_contains(sql_context: &SQLContext, sql: &str, expected: &str) { - let err = match sql_context.sql(sql).await { - Ok(df) => df - .collect() - .await - .expect_err("SQL should fail but succeeded") - .to_string(), - Err(err) => err.to_string(), - }; - assert!( - err.contains(expected), - "expected error containing '{expected}', got: {err}" - ); +#[tokio::test] +async fn test_information_schema_backs_off_after_catalog_refresh_failure() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + catalog.fail_list_tables(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + let calls_after_failure = catalog.metadata_calls(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + + assert_eq!(catalog.metadata_calls(), calls_after_failure); } #[tokio::test] -async fn test_show_tables_is_enabled() { - let (_tmp, catalog) = create_test_env(); - let sql_context = create_sql_context(catalog).await; +async fn test_information_schema_backs_off_after_table_classification_failure() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + let table_names: Vec<_> = (0..100).map(|index| format!("table_{index}")).collect(); + *catalog.table_names.lock().unwrap() = table_names; + catalog.fail_all_list_table_types(); sql_context - .sql("SHOW TABLES") + .sql("SELECT * FROM paimon.information_schema.tables") .await - .expect("SHOW TABLES should be planned when information_schema is enabled") - .collect() + .unwrap(); + let calls_after_failure = catalog.list_table_types_calls(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") .await - .expect("SHOW TABLES should execute"); + .unwrap(); + + assert_eq!(catalog.list_table_types_calls(), calls_after_failure); } #[tokio::test] -async fn test_show_tables_does_not_load_table_providers() { +async fn test_information_schema_times_out_pending_catalog_refresh() { let catalog = Arc::new(MetadataListingCatalog::new()); - let mut sql_context = SQLContext::new(); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .with_catalog_metadata_refresh_timeout(std::time::Duration::from_millis(20)) + .build(); sql_context .register_catalog("paimon", catalog.clone()) .await .unwrap(); + catalog.block_next_list_tables(); - let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await; + tokio::time::timeout( + std::time::Duration::from_millis(200), + sql_context.sql("SELECT * FROM paimon.information_schema.tables"), + ) + .await + .expect("automatic catalog refresh must be time bounded") + .unwrap(); +} - assert!(table_names.contains(&"metadata_only".to_string())); - assert_eq!(catalog.get_table_calls(), 0); +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_information_schema_bounds_cross_catalog_refresh_concurrency() { + let listing_concurrency = Arc::new(MetadataListingConcurrency::default()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .with_max_concurrent_catalog_metadata_refreshes(2) + .build(); + for catalog_name in ["first", "second", "third", "fourth"] { + sql_context + .register_catalog( + catalog_name, + Arc::new(MetadataListingCatalog::with_listing_concurrency( + Arc::clone(&listing_concurrency), + )), + ) + .await + .unwrap(); + } + listing_concurrency.reset(); - assert!(sql_context - .sql("SELECT * FROM metadata_only") + sql_context + .sql("SELECT * FROM first.information_schema.tables") .await - .is_err()); - assert!(catalog.get_table_calls() > 0); + .unwrap(); + + assert_eq!(listing_concurrency.maximum(), 2); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_information_schema_bounds_metadata_io_across_catalogs_and_databases() { + let listing_concurrency = Arc::new(MetadataListingConcurrency::default()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .with_max_concurrent_catalog_metadata_refreshes(4) + .with_max_concurrent_catalog_metadata_requests(3) + .build(); + let mut catalogs = Vec::new(); + for catalog_name in ["first", "second", "third", "fourth"] { + let catalog = Arc::new(MetadataListingCatalog::with_listing_concurrency( + Arc::clone(&listing_concurrency), + )); + catalog.set_databases(vec!["db1", "db2", "db3", "db4"]); + sql_context + .register_catalog_with_default_db(catalog_name, catalog.clone(), None) + .await + .unwrap(); + catalogs.push((catalog_name, catalog)); + } + for (catalog_name, catalog) in catalogs { + for database in ["db1", "db2", "db3", "db4"] { + let discovered = format!("{catalog_name}_{database}_new"); + catalog.set_table_names_for(database, vec![&discovered]); + } + } + listing_concurrency.reset(); + + sql_context + .sql("SELECT * FROM first.information_schema.tables") + .await + .unwrap(); + + assert_eq!(listing_concurrency.maximum(), 3); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_metadata_type_probes_use_global_bounded_parallelism() { + let listing_concurrency = Arc::new(MetadataListingConcurrency::default()); + let catalog = Arc::new(MetadataListingCatalog::with_listing_concurrency( + Arc::clone(&listing_concurrency), + )); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .with_max_concurrent_catalog_metadata_requests(3) + .build(); + sql_context + .register_catalog("paimon", catalog.clone()) + .await + .unwrap(); + catalog.set_table_names(vec!["new1", "new2", "new3", "new4", "new5", "new6"]); + listing_concurrency.reset(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + + assert_eq!(listing_concurrency.maximum(), 3); } #[tokio::test] -async fn test_show_tables_preserves_catalog_view_type() { +async fn test_information_schema_does_not_hide_strict_object_refresh_failure() { let catalog = Arc::new(MetadataListingCatalog::new()); - let mut sql_context = SQLContext::new(); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); sql_context .register_catalog("paimon", catalog.clone()) .await .unwrap(); + catalog.fail_list_tables(); - let table_types = collect_string_column( - &sql_context, - "SELECT table_type FROM information_schema.tables \ - WHERE table_schema = 'default' AND table_name = 'metadata_view'", - "table_type", - ) - .await; + let error = sql_context + .sql( + "SELECT missing.* FROM paimon.information_schema.tables info \ + JOIN paimon.default.missing AS missing ON true", + ) + .await + .unwrap_err(); - assert_eq!(table_types, vec!["VIEW"]); - assert_eq!(catalog.get_table_calls(), 0); + assert!( + error + .to_string() + .contains("simulated metadata refresh failure"), + "unexpected error: {error}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_information_schema_refreshes_catalogs_concurrently() { + let first = Arc::new(MetadataListingCatalog::new()); + let second = Arc::new(MetadataListingCatalog::new()); + let mut sql_context = SQLContext::builder() + .with_catalog_metadata_refresh_ttl(std::time::Duration::ZERO) + .build(); + sql_context + .register_catalog("first", first.clone()) + .await + .unwrap(); + sql_context + .register_catalog("second", second.clone()) + .await + .unwrap(); + first.block_next_list_tables(); + second.block_next_list_tables(); + let sql_context = Arc::new(sql_context); + + let query_context = Arc::clone(&sql_context); + let query = tokio::spawn(async move { + query_context + .sql("SELECT * FROM first.information_schema.tables") + .await + }); + + tokio::time::timeout(std::time::Duration::from_millis(100), async { + tokio::join!( + first.blocked_list_tables_started.notified(), + second.blocked_list_tables_started.notified() + ); + }) + .await + .expect("information schema catalog refreshes must start concurrently"); + first.release_blocked_list_tables.notify_one(); + second.release_blocked_list_tables.notify_one(); + query.await.unwrap().unwrap(); } #[tokio::test] @@ -960,14 +2583,6 @@ async fn test_drop_schema() { #[tokio::test] async fn test_schema_names_via_catalog_provider() { let (_tmp, catalog) = create_test_env(); - let provider = PaimonCatalogProvider::new( - None, - catalog.clone(), - Default::default(), - Default::default(), - None, - ); - catalog .create_database("db_a", false, Default::default()) .await @@ -977,6 +2592,16 @@ async fn test_schema_names_via_catalog_provider() { .await .unwrap(); + let provider = PaimonCatalogProvider::try_new( + None, + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + let names = provider.schema_names(); assert!(names.contains(&"db_a".to_string())); assert!(names.contains(&"db_b".to_string())); @@ -1149,6 +2774,104 @@ async fn test_create_table_if_not_exists() { .expect("Second CREATE with IF NOT EXISTS should succeed"); } +#[tokio::test] +async fn test_create_table_if_not_exists_classifies_new_paimon_table() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + catalog + .create_database("mydb", false, Default::default()) + .await + .unwrap(); + + sql_context + .sql("CREATE TABLE IF NOT EXISTS paimon.mydb.records (id BIGINT)") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("mydb").unwrap(); + assert!(schema.table_exist("records$snapshots")); +} + +#[tokio::test] +async fn test_create_object_table_delta_does_not_claim_system_table_support() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + catalog + .create_database("mydb", false, Default::default()) + .await + .unwrap(); + + sql_context + .sql( + "CREATE TABLE paimon.mydb.objects (id BIGINT) \ + WITH ('type' = 'object-table')", + ) + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("mydb").unwrap(); + assert!(schema.table_exist("objects")); + assert!(!schema.table_exist("objects$snapshots")); +} + +#[tokio::test] +async fn test_create_table_if_not_exists_noop_does_not_overwrite_object_capability() { + let (_tmp, catalog) = create_test_env(); + catalog + .create_database("mydb", false, Default::default()) + .await + .unwrap(); + let object_schema = paimon::spec::Schema::builder() + .column("id", paimon::spec::DataType::BigInt(Default::default())) + .option("type", "object-table") + .build() + .unwrap(); + catalog + .create_table(&Identifier::new("mydb", "objects"), object_schema, false) + .await + .unwrap(); + let sql_context = create_sql_context(catalog).await; + + sql_context + .sql("CREATE TABLE IF NOT EXISTS paimon.mydb.objects (id BIGINT)") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("mydb").unwrap(); + assert!(schema.table_exist("objects")); + assert!(!schema.table_exist("objects$snapshots")); +} + +#[tokio::test] +async fn test_create_table_if_not_exists_noop_preserves_paimon_capability() { + let (_tmp, catalog) = create_test_env(); + catalog + .create_database("mydb", false, Default::default()) + .await + .unwrap(); + let schema = paimon::spec::Schema::builder() + .column("id", paimon::spec::DataType::BigInt(Default::default())) + .build() + .unwrap(); + catalog + .create_table(&Identifier::new("mydb", "records"), schema, false) + .await + .unwrap(); + let sql_context = create_sql_context(catalog).await; + + sql_context + .sql("CREATE TABLE IF NOT EXISTS paimon.mydb.records (id BIGINT)") + .await + .unwrap(); + + let provider = sql_context.ctx().catalog("paimon").unwrap(); + let schema = provider.schema("mydb").unwrap(); + assert!(schema.table_exist("records$snapshots")); +} + #[tokio::test] async fn test_create_external_table_rejected() { let (_tmp, catalog) = create_test_env(); diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index 3e979637a..1c67d6626 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -22,13 +22,16 @@ use async_trait::async_trait; use datafusion::arrow::array::{Array, Int32Array, StringArray}; use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::CatalogProvider; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use paimon::catalog::{Catalog, Database, Identifier, LoadedTable}; use paimon::spec::{Schema as PaimonSchema, SchemaChange, TableType}; use paimon::table::Table; use paimon::{CatalogOptions, FileSystemCatalog, Options, Result as PaimonResult}; -use paimon_datafusion::{EngineTableRequest, SQLContext, TableEngineResolver}; +use paimon_datafusion::{ + EngineTableRequest, PaimonCatalogProvider, SQLContext, TableEngineResolver, +}; use tempfile::TempDir; const CATALOG: &str = "cat"; @@ -309,6 +312,10 @@ async fn setup() -> TestEnv { let mut options = Options::new(); options.set(CatalogOptions::WAREHOUSE, warehouse); let fs_catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + fs_catalog + .create_database(DB, false, HashMap::new()) + .await + .unwrap(); let typed_catalog = Arc::new(TypedTestCatalog { inner: fs_catalog, declared_types: HashMap::from([ @@ -319,9 +326,6 @@ async fn setup() -> TestEnv { }); let mut ctx = SQLContext::new(); ctx.register_catalog(CATALOG, typed_catalog).await.unwrap(); - ctx.sql(&format!("CREATE SCHEMA {CATALOG}.{DB}")) - .await - .unwrap(); ctx.sql(&format!( "CREATE TABLE {CATALOG}.{DB}.pt (id INT NOT NULL, name STRING)" )) @@ -571,13 +575,99 @@ async fn system_tables_on_routed_tables_error() { } #[tokio::test] -async fn table_exist_mirrors_the_resolver() { +async fn registered_engine_none_is_not_found() { let env = setup().await; let provider = env.ctx.ctx().catalog(CATALOG).unwrap(); let schema = provider.schema(DB).unwrap(); - assert!(schema.table_exist("it")); + assert!(schema.table("ghost").await.unwrap().is_none()); assert!(!schema.table_exist("ghost")); + assert!(!schema.table_exist("ghost$snapshots")); + assert!(schema.table_type("ghost").await.unwrap().is_none()); +} + +#[tokio::test] +async fn metadata_refresh_clears_registered_engine_miss() { + let env = setup().await; + let provider = env.ctx.ctx().catalog(CATALOG).unwrap(); + let schema = provider.schema(DB).unwrap(); + assert!(schema.table("ghost").await.unwrap().is_none()); + assert!(!schema.table_exist("ghost")); + + provider + .downcast_ref::() + .unwrap() + .refresh_metadata() + .await + .unwrap(); + + assert!(provider.schema(DB).unwrap().table_exist("ghost")); +} + +#[tokio::test] +async fn table_exist_derives_system_table_from_snapshotted_base_table() { + let env = setup().await; + let provider = env.ctx.ctx().catalog(CATALOG).unwrap(); + let schema = provider.schema(DB).unwrap(); + + assert!(schema.table_exist("pt$snapshots")); + assert!(schema.table("it").await.unwrap().is_some()); + assert!(schema.table_exist("it")); assert!(!schema.table_exist("it$snapshots")); + assert!(!schema.table_exist("missing$snapshots")); + assert!(!schema.table_exist("pt$not_a_system_table")); +} + +#[tokio::test] +async fn object_table_does_not_advertise_paimon_system_tables_before_table_load() { + let paimon_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", paimon_dir.path().display()); + let mut options = Options::new(); + options.set(CatalogOptions::WAREHOUSE, warehouse); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + catalog + .create_database(DB, false, HashMap::new()) + .await + .unwrap(); + + let schema = PaimonSchema::builder() + .column( + "id", + paimon::spec::DataType::Int(paimon::spec::IntType::new()), + ) + .build() + .unwrap(); + catalog + .create_table(&Identifier::new(DB, "pt"), schema.clone(), false) + .await + .unwrap(); + let object_schema = PaimonSchema::builder() + .column( + "ignored", + paimon::spec::DataType::Int(paimon::spec::IntType::new()), + ) + .option("type", "object-table") + .build() + .unwrap(); + catalog + .create_table(&Identifier::new(DB, "objects"), object_schema, false) + .await + .unwrap(); + + let provider = PaimonCatalogProvider::try_new( + Some(CATALOG.to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + let schema = provider.schema(DB).unwrap(); + + assert!(schema.table_exist("objects")); + assert!(!schema.table_exist("objects$snapshots")); + assert!(schema.table("objects$snapshots").await.is_err()); + assert!(schema.table_exist("pt$snapshots")); } #[tokio::test] @@ -676,13 +766,17 @@ async fn an_unsupported_time_travel_clause_on_a_paimon_table_is_rejected() { let ctx = SessionContext::new(); ctx.register_catalog( CATALOG, - Arc::new(PaimonCatalogProvider::new( - Some(CATALOG.to_string()), - fs_catalog, - Default::default(), - Default::default(), - None, - )), + Arc::new( + PaimonCatalogProvider::try_new( + Some(CATALOG.to_string()), + fs_catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(), + ), ); paimon_datafusion::register_catalog_table_engine( &ctx, @@ -949,13 +1043,17 @@ async fn registering_on_a_raw_session_installs_the_planner() { let ctx = SessionContext::new(); ctx.register_catalog( CATALOG, - Arc::new(PaimonCatalogProvider::new( - Some(CATALOG.to_string()), - typed_catalog, - Default::default(), - Default::default(), - None, - )), + Arc::new( + PaimonCatalogProvider::try_new( + Some(CATALOG.to_string()), + typed_catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(), + ), ); register_catalog_table_engine( &ctx, @@ -1016,15 +1114,16 @@ async fn unregistered_external_table_does_not_break_information_schema_columns() let mut options = Options::new(); options.set(CatalogOptions::WAREHOUSE, warehouse); let fs_catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + fs_catalog + .create_database(DB, false, HashMap::new()) + .await + .unwrap(); let typed_catalog = Arc::new(TypedTestCatalog { inner: fs_catalog, declared_types: HashMap::from([("external".to_string(), TableType::IcebergTable)]), }); let mut ctx = SQLContext::new(); ctx.register_catalog(CATALOG, typed_catalog).await.unwrap(); - ctx.sql(&format!("CREATE SCHEMA {CATALOG}.{DB}")) - .await - .unwrap(); ctx.sql(&format!( "CREATE TABLE {CATALOG}.{DB}.pt (id INT NOT NULL, name STRING)" )) @@ -1135,6 +1234,24 @@ async fn the_default_load_table_classifies_for_a_catalog_that_only_has_get_table ); } +#[tokio::test] +async fn default_table_type_listing_does_not_claim_system_table_support() { + let (_dir, catalog) = legacy_catalog_with_iceberg_table().await; + let provider = PaimonCatalogProvider::try_new( + Some(CATALOG.to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + let schema = provider.schema(DB).unwrap(); + + assert!(schema.table_exist("it")); + assert!(!schema.table_exist("it$snapshots")); +} + #[tokio::test] async fn a_legacy_catalog_cannot_serve_an_external_table_as_paimon() { let (_dir, catalog) = legacy_catalog_with_iceberg_table().await; diff --git a/crates/paimon/src/catalog/filesystem.rs b/crates/paimon/src/catalog/filesystem.rs index 743f00286..d28f0df17 100644 --- a/crates/paimon/src/catalog/filesystem.rs +++ b/crates/paimon/src/catalog/filesystem.rs @@ -426,6 +426,21 @@ impl Catalog for FileSystemCatalog { Ok(tables) } + async fn list_table_types( + &self, + database_name: &str, + table_names: &[String], + ) -> Result> { + let mut types = HashMap::with_capacity(table_names.len()); + for table_name in table_names { + let identifier = Identifier::new(database_name, table_name); + let (_, schema) = self.fetch_table_schema(&identifier).await?; + let table_type = CoreOptions::new(schema.options()).table_type()?; + types.insert(table_name.clone(), table_type); + } + Ok(types) + } + async fn create_table( &self, identifier: &Identifier, diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index 0c66f8bef..9d367c436 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -444,6 +444,19 @@ pub trait Catalog: Send + Sync { /// * [`crate::Error::DatabaseNotExist`] - database does not exist. async fn list_tables(&self, database_name: &str) -> Result>; + /// Return the declared type for the requested tables. + /// + /// Catalogs that know their declared table types should override this method with an + /// implementation backed by their metadata store. The default is deliberately empty so a + /// mixed or legacy catalog is not assumed to support Paimon system tables. + async fn list_table_types( + &self, + _database_name: &str, + _table_names: &[String], + ) -> Result> { + Ok(HashMap::new()) + } + /// Create a table. /// /// * `ignore_if_exists` - if true, do nothing when the table already exists; diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 724dbcdbb..0cad642d7 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -37,7 +37,7 @@ use crate::catalog::{ use crate::common::{CatalogOptions, Options}; use crate::error::Error; use crate::io::cache::{create_local_cache_with_namespace, LocalCache}; -use crate::spec::{Partition, PartitionStatistics, Schema, SchemaChange}; +use crate::spec::{CoreOptions, Partition, PartitionStatistics, Schema, SchemaChange, TableType}; use crate::table::{RESTEnv, Table}; use crate::Result; @@ -326,6 +326,30 @@ impl Catalog for RESTCatalog { .map_err(|e| map_rest_error_for_database(e, database_name)) } + async fn list_table_types( + &self, + database_name: &str, + table_names: &[String], + ) -> Result> { + let mut types = HashMap::with_capacity(table_names.len()); + for table_name in table_names { + let identifier = Identifier::new(database_name, table_name); + let response = self + .api + .get_table(&identifier) + .await + .map_err(|error| map_rest_error_for_table(error, &identifier))?; + let table_type = response + .schema + .as_ref() + .map(|schema| CoreOptions::new(schema.options()).table_type()) + .transpose()? + .unwrap_or_default(); + types.insert(table_name.clone(), table_type); + } + Ok(types) + } + async fn create_table( &self, identifier: &Identifier, diff --git a/crates/paimon/tests/rest_catalog_test.rs b/crates/paimon/tests/rest_catalog_test.rs index bb617d056..fb3b9528c 100644 --- a/crates/paimon/tests/rest_catalog_test.rs +++ b/crates/paimon/tests/rest_catalog_test.rs @@ -36,7 +36,7 @@ use paimon::catalog::{Catalog, Function, FunctionDefinition, Identifier, RESTCat use paimon::common::Options; use paimon::spec::{ BigIntType, BlobType, BlobViewStruct, DataField, DataType, Datum, IntType, PartitionStatistics, - PredicateBuilder, Schema, SchemaChange, VarCharType, + PredicateBuilder, Schema, SchemaChange, TableType, VarCharType, }; use paimon::{CatalogOptions, FileSystemCatalog, Table}; @@ -837,6 +837,38 @@ async fn test_catalog_list_tables_empty() { ); } +#[tokio::test] +async fn test_catalog_lists_declared_table_types() { + let ctx = setup_catalog(vec!["default"]).await; + let table_schema = test_schema(); + let object_schema = Schema::builder() + .column("ignored", DataType::Int(IntType::new())) + .option("type", "object-table") + .build() + .unwrap(); + ctx.server.add_table_with_schema( + "default", + "plain", + table_schema, + "file:///tmp/test_warehouse/default.db/plain", + ); + ctx.server.add_table_with_schema( + "default", + "objects", + object_schema, + "file:///tmp/test_warehouse/default.db/objects", + ); + + let table_types = ctx + .catalog + .list_table_types("default", &["plain".to_string(), "objects".to_string()]) + .await + .unwrap(); + + assert_eq!(table_types.get("plain"), Some(&TableType::Table)); + assert_eq!(table_types.get("objects"), Some(&TableType::ObjectTable)); +} + #[tokio::test] async fn test_catalog_get_table() { let ctx = setup_catalog(vec!["default"]).await;