From 1397a0b85cf87dbfbd0a90cbdd1f072f96f97650 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:21:31 -0700 Subject: [PATCH 1/7] fix(datafusion): snapshot catalog metadata for sync callbacks --- bindings/python/src/context.rs | 31 +- bindings/python/tests/test_datafusion.py | 31 ++ crates/integrations/datafusion/src/catalog.rs | 503 ++++++++++-------- .../datafusion/src/sql_context.rs | 173 +++++- .../datafusion/tests/sql_context_tests.rs | 146 ++++- .../datafusion/tests/table_type_routing.rs | 42 +- 6 files changed, 650 insertions(+), 276 deletions(-) diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 1e285fd19..3a3237551 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -121,19 +121,30 @@ impl PaimonCatalog { #[new] fn new(py: Python<'_>, catalog_options: HashMap) -> PyResult { let catalog = py.detach(|| build_paimon_catalog(catalog_options))?; - let provider = Arc::new( - PaimonCatalogProvider::new( - None, - Arc::clone(&catalog), - Default::default(), - Default::default(), - None, - ) - .with_schema_force_view_types(false), - ); + let provider = { + let catalog = Arc::clone(&catalog); + py.detach(|| { + runtime().block_on(PaimonCatalogProvider::try_new( + None, + catalog, + Default::default(), + Default::default(), + None, + )) + }) + .map_err(df_to_py_err)? + }; + let provider = Arc::new(provider.with_schema_force_view_types(false)); 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, 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/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 06f34c3c4..7d7c6eee3 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -49,6 +49,53 @@ 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 { + database_names: Vec, + databases: HashMap, +} + +#[derive(Clone, Debug, Default)] +struct DatabaseMetadata { + object_names: Vec, + object_types: HashMap, +} + +type SharedCatalogMetadata = Arc>>; + +async fn load_database_metadata( + catalog: &dyn Catalog, + database: &str, +) -> DFResult { + let table_names = catalog + .list_tables(database) + .await + .map_err(to_datafusion_error)?; + let view_names = match catalog.list_views(database).await { + Ok(names) => names, + Err(paimon::Error::Unsupported { .. }) => vec![], + Err(error) => return Err(to_datafusion_error(error)), + }; + + let mut object_names = Vec::with_capacity(table_names.len() + view_names.len()); + let mut object_types = HashMap::with_capacity(object_names.capacity()); + for name in table_names { + if object_types.insert(name.clone(), TableType::Base).is_none() { + object_names.push(name); + } + } + for name in view_names { + if !object_types.contains_key(&name) { + object_types.insert(name.clone(), TableType::View); + object_names.push(name); + } + } + Ok(DatabaseMetadata { + object_names, + object_types, + }) +} + /// 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. @@ -232,8 +279,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 +314,8 @@ 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, } impl Debug for PaimonCatalogProvider { @@ -275,6 +326,10 @@ impl Debug for PaimonCatalogProvider { impl PaimonCatalogProvider { /// Creates a new [`PaimonCatalogProvider`]. + /// + /// The provider starts with an empty metadata snapshot. Call + /// [`Self::refresh_metadata`] before exposing it to DataFusion, or prefer + /// [`Self::try_new`] when an async construction boundary is available. pub fn new( catalog_name: Option, catalog: Arc, @@ -291,7 +346,56 @@ impl PaimonCatalogProvider { session_state, schema_force_view_types: true, table_engines: Arc::new(RwLock::new(HashMap::new())), + metadata: Arc::new(RwLock::new(Arc::new(CatalogMetadataSnapshot::default()))), + } + } + + /// Creates a provider and initializes its metadata snapshot before returning it. + pub async fn try_new( + catalog_name: Option, + catalog: Arc, + dynamic_options: DynamicOptions, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + ) -> DFResult { + let provider = Self::new( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + ); + provider.refresh_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<()> { + let mut database_names = 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 mut databases = HashMap::with_capacity(database_names.len()); + for database in &database_names { + databases.insert( + database.clone(), + load_database_metadata(self.catalog.as_ref(), database).await?, + ); } + + let next = Arc::new(CatalogMetadataSnapshot { + database_names, + databases, + }); + *self.metadata.write().unwrap_or_else(|e| e.into_inner()) = next; + Ok(()) } /// Configure whether table schemas use Arrow view types when available. @@ -333,79 +437,73 @@ 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 { + 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.object_types.contains_key(&object_name)) + } + 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( + 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_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 = self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .database_names + .clone(); + 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 +521,8 @@ 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 name = name.to_string(); block_on_with_runtime( async move { @@ -430,6 +530,15 @@ impl CatalogProvider for PaimonCatalogProvider { .create_database(&name, false, HashMap::new()) .await .map_err(to_datafusion_error)?; + { + let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + if !next.database_names.contains(&name) { + next.database_names.push(name.clone()); + } + next.databases.entry(name.clone()).or_default(); + *current = Arc::new(next); + } Ok(Some(Arc::new( PaimonSchemaProvider::new( catalog_name, @@ -440,7 +549,9 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) - .with_schema_force_view_types(schema_force_view_types), + .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 +569,8 @@ 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 name = name.to_string(); block_on_with_runtime( async move { @@ -465,6 +578,13 @@ impl CatalogProvider for PaimonCatalogProvider { .drop_database(&name, false, cascade) .await .map_err(to_datafusion_error)?; + { + let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + next.database_names.retain(|database| database != &name); + next.databases.remove(&name); + *current = Arc::new(next); + } Ok(Some(Arc::new( PaimonSchemaProvider::new( catalog_name, @@ -475,7 +595,9 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) - .with_schema_force_view_types(schema_force_view_types), + .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 +617,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.object_types.contains_key(table_name)) + { log::warn!( "Temporary table '{database}.{table_name}' shadows an existing Paimon table" ); @@ -575,8 +691,8 @@ 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, @@ -595,6 +711,10 @@ impl Debug for PaimonSchemaProvider { impl PaimonSchemaProvider { /// Creates a new [`PaimonSchemaProvider`]. + /// + /// The provider starts with an empty metadata snapshot. Call + /// [`Self::refresh_metadata`] before exposing it to DataFusion, or prefer + /// [`Self::try_new`] when an async construction boundary is available. pub fn new( catalog_name: Option, catalog: Arc, @@ -610,7 +730,7 @@ impl PaimonSchemaProvider { database, dynamic_options, temp_provider, - catalog_table_types: RwLock::new(HashMap::new()), + metadata: Arc::new(RwLock::new(Arc::new(CatalogMetadataSnapshot::default()))), blob_reader_registry, session_state, schema_force_view_types: true, @@ -618,6 +738,42 @@ impl PaimonSchemaProvider { } } + /// Creates a schema provider and initializes its metadata snapshot. + 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( + catalog_name, + catalog, + database, + dynamic_options, + temp_provider, + blob_reader_registry, + session_state, + ); + provider.refresh_metadata().await?; + Ok(provider) + } + + /// Refresh this database in the snapshot used by synchronous callbacks. + pub async fn refresh_metadata(&self) -> DFResult<()> { + let database = load_database_metadata(self.catalog.as_ref(), &self.database).await?; + let mut current = self.metadata.write().unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + if !next.database_names.contains(&self.database) { + next.database_names.push(self.database.clone()); + } + next.databases.insert(self.database.clone(), database); + *current = Arc::new(next); + Ok(()) + } + fn with_schema_force_view_types(mut self, schema_force_view_types: bool) -> Self { self.schema_force_view_types = schema_force_view_types; self @@ -627,52 +783,24 @@ impl PaimonSchemaProvider { 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 = self + .metadata + .read() + .unwrap_or_else(|e| e.into_inner()) + .databases + .get(&self.database) + .map(|database| database.object_names.clone()) + .unwrap_or_default(); if let Some(temp) = &self.temp_provider { names.extend(temp.table_names()); @@ -897,10 +1025,12 @@ impl SchemaProvider for PaimonSchemaProvider { } 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.object_types.get(name)) { return Ok(Some(*table_type)); } @@ -924,101 +1054,22 @@ impl SchemaProvider for PaimonSchemaProvider { return false; } }; - if let Some(system_name) = object.system_table() { - if !system_tables::is_registered(system_name) { - return false; - } + if object + .system_table() + .is_some_and(|system_name| !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 - .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 + // This callback cannot await an external engine resolver. Treat a catalog + // declaration as existing and let async `table()` surface resolver absence + // or unsupported system-table access during planning. + 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) + .is_some_and(|database| database.object_types.contains_key(object.table())) } fn register_table( @@ -1033,7 +1084,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 +1101,15 @@ impl SchemaProvider for PaimonSchemaProvider { .drop_table(&identifier, false) .await .map_err(to_datafusion_error)?; + { + let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + if let Some(database) = next.databases.get_mut(&database) { + database.object_names.retain(|object| object != &name); + database.object_types.remove(&name); + } + *current = Arc::new(next); + } Ok(Some(Arc::new(provider) as Arc)) }, "paimon catalog access thread panicked", diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index aba77b159..a2b55d254 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -67,8 +67,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; @@ -252,16 +253,17 @@ 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( Some(catalog_name.clone()), catalog.clone(), self.dynamic_options.clone(), self.blob_reader_registry.clone(), Some(session_state), - )), + ) + .await?, ); + self.ctx.register_catalog(&catalog_name, provider); register_table_functions( &self.ctx, &catalog, @@ -473,7 +475,12 @@ impl SQLContext { )); } - match &statements[0] { + if self.statement_needs_catalog_metadata(&statements[0])? { + self.refresh_all_catalog_metadata().await?; + } + + let refresh_metadata_after = statement_changes_catalog_metadata(&statements[0]); + let result = match &statements[0] { Statement::ShowDatabases { terse, history, @@ -515,6 +522,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 +671,7 @@ impl SQLContext { } } Statement::Drop { - object_type: ObjectType::Database, + object_type: object_type @ (ObjectType::Database | ObjectType::Schema), if_exists, names, cascade, @@ -650,15 +680,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 +787,73 @@ impl SQLContext { self.ctx.sql(&expanded.to_string()).await } _ => self.ctx.sql(sql).await, + }; + + if refresh_metadata_after && result.is_ok() { + if let Err(error) = self.refresh_all_catalog_metadata().await { + log::warn!("catalog metadata refresh after DDL failed: {error}"); + } + } + result + } + + fn statement_needs_catalog_metadata(&self, statement: &Statement) -> DFResult { + if matches!( + statement, + Statement::ShowTables { .. } + | Statement::ShowColumns { .. } + | Statement::ShowFunctions { .. } + ) { + return Ok(true); + } + + 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); + if schema.eq_ignore_ascii_case("information_schema") { + return Ok(true); + } + + 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 !provider.metadata_contains_object(schema, reference.table()) { + return Ok(true); + } + } + Ok(false) + } + + async fn refresh_all_catalog_metadata(&self) -> DFResult<()> { + let catalog_names: Vec<_> = self.catalogs.keys().cloned().collect(); + for catalog_name in catalog_names { + 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_metadata().await?; } + Ok(()) } /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). @@ -2304,6 +2406,23 @@ impl SQLContext { } } +fn statement_changes_catalog_metadata(statement: &Statement) -> bool { + match statement { + Statement::CreateDatabase { .. } + | Statement::CreateSchema { .. } + | Statement::AlterTable(_) + | Statement::Drop { + object_type: + ObjectType::Database | ObjectType::Schema | ObjectType::Table | ObjectType::View, + temporary: false, + .. + } => true, + Statement::CreateTable(create) => !create.temporary, + Statement::CreateView(create) => !create.temporary, + _ => false, + } +} + fn validate_immutable_scalar_plan(plan: &LogicalPlan) -> DFResult<()> { let mut violation = None; plan.apply(|node| { @@ -3905,7 +4024,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 +6189,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(), diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 2e550631c..bd90f41b1 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -52,23 +52,56 @@ 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, + table_names: Mutex>, } 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), + table_names: Mutex::new(vec!["metadata_only".to_string()]), } } 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 fail_list_tables(&self) { + self.fail_list_tables.store(true, Ordering::SeqCst); + } + + 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> { + self.record_remote_call(); Ok(vec!["default".to_string()]) } @@ -82,6 +115,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(), @@ -99,6 +133,7 @@ impl Catalog for MetadataListingCatalog { } 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(), @@ -106,10 +141,17 @@ impl Catalog for MetadataListingCatalog { } async fn list_tables(&self, _database_name: &str) -> paimon::Result> { - Ok(vec!["metadata_only".to_string()]) + self.record_remote_call(); + if self.fail_list_tables.load(Ordering::SeqCst) { + return Err(paimon::Error::Unsupported { + message: "simulated metadata refresh failure".to_string(), + }); + } + Ok(self.table_names.lock().unwrap().clone()) } async fn list_views(&self, _database_name: &str) -> paimon::Result> { + self.record_remote_call(); Ok(vec!["metadata_view".to_string()]) } @@ -149,6 +191,74 @@ impl Catalog for MetadataListingCatalog { } } +#[tokio::test] +async fn test_refreshed_catalog_callbacks_do_not_access_remote_catalog() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new( + 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); +} + +#[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] +async fn test_temp_table_registration_does_not_access_remote_catalog() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new( + 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, @@ -403,6 +513,22 @@ async fn test_show_tables_does_not_load_table_providers() { 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_preserves_catalog_view_type() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -960,14 +1086,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 +1095,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())); diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index 3e979637a..b6df2da3b 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -571,13 +571,13 @@ async fn system_tables_on_routed_tables_error() { } #[tokio::test] -async fn table_exist_mirrors_the_resolver() { +async fn table_exist_uses_catalog_declarations_without_resolving_engines() { 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_exist("ghost")); - assert!(!schema.table_exist("it$snapshots")); + assert!(schema.table_exist("ghost")); + assert!(schema.table_exist("it$snapshots")); } #[tokio::test] @@ -676,13 +676,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 +953,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, From 70f155d375629ac5ff046d0c322615e9601b0c4c Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sat, 19 Sep 2026 21:41:43 -0700 Subject: [PATCH 2/7] fix(datafusion): harden catalog metadata refresh --- Cargo.lock | 1 + bindings/python/src/context.rs | 27 +- bindings/python/tests/test_catalog_gil.py | 21 + crates/integrations/datafusion/Cargo.toml | 3 +- crates/integrations/datafusion/src/catalog.rs | 419 +++++++++++------ .../datafusion/src/sql_context.rs | 220 +++++++-- .../datafusion/tests/read_tables.rs | 4 +- .../datafusion/tests/sql_context_tests.rs | 421 +++++++++++++++++- .../datafusion/tests/table_type_routing.rs | 5 +- 9 files changed, 922 insertions(+), 199 deletions(-) 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 3a3237551..40f7cfab1 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -121,20 +121,16 @@ impl PaimonCatalog { #[new] fn new(py: Python<'_>, catalog_options: HashMap) -> PyResult { let catalog = py.detach(|| build_paimon_catalog(catalog_options))?; - let provider = { - let catalog = Arc::clone(&catalog); - py.detach(|| { - runtime().block_on(PaimonCatalogProvider::try_new( - None, - catalog, - Default::default(), - Default::default(), - None, - )) - }) - .map_err(df_to_py_err)? - }; - let provider = Arc::new(provider.with_schema_force_view_types(false)); + let provider = Arc::new( + PaimonCatalogProvider::new_uninitialized( + None, + Arc::clone(&catalog), + Default::default(), + Default::default(), + None, + ) + .with_schema_force_view_types(false), + ); Ok(Self { catalog, provider }) } @@ -151,6 +147,9 @@ impl PaimonCatalog { 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/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 7d7c6eee3..13ae91916 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -19,6 +19,8 @@ use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt::Debug; +use std::ops::Deref; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::sync::RwLock; @@ -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; @@ -51,49 +55,114 @@ type TableEngines = Arc, - databases: HashMap, + databases: IndexMap>, } #[derive(Clone, Debug, Default)] struct DatabaseMetadata { - object_names: Vec, - object_types: HashMap, + objects: IndexMap, } -type SharedCatalogMetadata = Arc>>; +#[derive(Debug)] +struct CatalogMetadataState { + snapshot: RwLock>, + next_generation: AtomicU64, + published_generation: AtomicU64, +} + +impl Default for CatalogMetadataState { + fn default() -> Self { + Self { + snapshot: RwLock::new(Arc::new(CatalogMetadataSnapshot::default())), + next_generation: AtomicU64::new(0), + published_generation: AtomicU64::new(0), + } + } +} + +impl Deref for CatalogMetadataState { + type Target = RwLock>; + + fn deref(&self) -> &Self::Target { + &self.snapshot + } +} + +impl CatalogMetadataState { + fn begin_refresh(&self) -> u64 { + self.next_generation.fetch_add(1, Ordering::AcqRel) + 1 + } + + fn publish(&self, generation: u64, next: Arc) { + let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); + if generation >= self.published_generation.load(Ordering::Acquire) { + *current = next; + self.published_generation + .store(generation, Ordering::Release); + } + } + + fn publish_update(&self, generation: u64, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { + let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); + if generation < self.published_generation.load(Ordering::Acquire) { + return; + } + let mut next = (**current).clone(); + update(&mut next); + *current = Arc::new(next); + self.published_generation + .store(generation, Ordering::Release); + } + + fn mutate(&self, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { + let generation = self.begin_refresh(); + self.publish_update(generation, update); + } +} + +type SharedCatalogMetadata = Arc; + +const MAX_CONCURRENT_METADATA_LISTINGS: usize = 16; async fn load_database_metadata( catalog: &dyn Catalog, database: &str, + ignore_missing_views_endpoint: bool, ) -> DFResult { - let table_names = catalog - .list_tables(database) - .await - .map_err(to_datafusion_error)?; - let view_names = match catalog.list_views(database).await { - Ok(names) => names, - Err(paimon::Error::Unsupported { .. }) => vec![], - Err(error) => return Err(to_datafusion_error(error)), + let tables = async { + catalog + .list_tables(database) + .await + .map_err(to_datafusion_error) }; + let views = async { + 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 (table_names, view_names) = futures::try_join!(tables, views)?; - let mut object_names = Vec::with_capacity(table_names.len() + view_names.len()); - let mut object_types = HashMap::with_capacity(object_names.capacity()); + let mut objects = IndexMap::with_capacity(table_names.len() + view_names.len()); for name in table_names { - if object_types.insert(name.clone(), TableType::Base).is_none() { - object_names.push(name); - } + objects.entry(name).or_insert(TableType::Base); } for name in view_names { - if !object_types.contains_key(&name) { - object_types.insert(name.clone(), TableType::View); - object_names.push(name); - } + objects.entry(name).or_insert(TableType::View); } - Ok(DatabaseMetadata { - object_names, - object_types, - }) + Ok(DatabaseMetadata { objects }) } /// What an engine is asked to resolve. Non-exhaustive so later releases can @@ -209,17 +278,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) } } @@ -325,12 +389,29 @@ impl Debug for PaimonCatalogProvider { } impl PaimonCatalogProvider { - /// Creates a new [`PaimonCatalogProvider`]. + /// Creates a provider with an initialized metadata snapshot. + pub async fn new( + catalog_name: Option, + catalog: Arc, + dynamic_options: DynamicOptions, + blob_reader_registry: BlobReaderRegistry, + session_state: Option, + ) -> DFResult { + let provider = Self::new_uninitialized( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + ); + provider.initialize_metadata().await?; + Ok(provider) + } + + /// Creates a provider with an empty metadata snapshot. /// - /// The provider starts with an empty metadata snapshot. Call - /// [`Self::refresh_metadata`] before exposing it to DataFusion, or prefer - /// [`Self::try_new`] when an async construction boundary is available. - pub fn new( + /// Callers must refresh it before exposing synchronous discovery callbacks. + pub fn new_uninitialized( catalog_name: Option, catalog: Arc, dynamic_options: DynamicOptions, @@ -346,11 +427,11 @@ impl PaimonCatalogProvider { session_state, schema_force_view_types: true, table_engines: Arc::new(RwLock::new(HashMap::new())), - metadata: Arc::new(RwLock::new(Arc::new(CatalogMetadataSnapshot::default()))), + metadata: Arc::new(CatalogMetadataState::default()), } } - /// Creates a provider and initializes its metadata snapshot before returning it. + /// Backward-compatible alias for [`Self::new`]. pub async fn try_new( catalog_name: Option, catalog: Arc, @@ -358,15 +439,14 @@ impl PaimonCatalogProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, ) -> DFResult { - let provider = Self::new( + Self::new( catalog_name, catalog, dynamic_options, blob_reader_registry, session_state, - ); - provider.refresh_metadata().await?; - Ok(provider) + ) + .await } /// Refresh the metadata consumed by DataFusion's synchronous catalog callbacks. @@ -374,6 +454,16 @@ impl PaimonCatalogProvider { /// 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.refresh_metadata_inner(false).await + } + + /// Initialize metadata while tolerating a REST server without the optional views endpoint. + pub async fn initialize_metadata(&self) -> DFResult<()> { + self.refresh_metadata_inner(true).await + } + + async fn refresh_metadata_inner(&self, ignore_missing_views_endpoint: bool) -> DFResult<()> { + let generation = self.metadata.begin_refresh(); let mut database_names = self .catalog .list_databases() @@ -382,19 +472,43 @@ impl PaimonCatalogProvider { let mut seen_databases = HashSet::new(); database_names.retain(|name| seen_databases.insert(name.clone())); - let mut databases = HashMap::with_capacity(database_names.len()); - for database in &database_names { - databases.insert( - database.clone(), - load_database_metadata(self.catalog.as_ref(), database).await?, - ); - } + let entries: HashMap<_, _> = stream::iter(database_names.iter().cloned()) + .map(|database| async move { + let metadata = load_database_metadata( + self.catalog.as_ref(), + database.as_str(), + ignore_missing_views_endpoint, + ) + .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 next = Arc::new(CatalogMetadataSnapshot { - database_names, - databases, + let next = Arc::new(CatalogMetadataSnapshot { databases }); + self.metadata.publish(generation, next); + Ok(()) + } + + /// Refresh one database in the metadata snapshot. + pub(crate) async fn refresh_database_metadata(&self, database: &str) -> DFResult<()> { + let generation = self.metadata.begin_refresh(); + let database_metadata = + load_database_metadata(self.catalog.as_ref(), database, true).await?; + self.metadata.publish_update(generation, |next| { + next.databases + .insert(database.to_string(), Arc::new(database_metadata)); }); - *self.metadata.write().unwrap_or_else(|e| e.into_inner()) = next; Ok(()) } @@ -438,6 +552,15 @@ impl PaimonCatalogProvider { } 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()); @@ -446,7 +569,7 @@ impl PaimonCatalogProvider { .unwrap_or_else(|e| e.into_inner()) .databases .get(database) - .is_some_and(|metadata| metadata.object_types.contains_key(&object_name)) + .is_some_and(|metadata| metadata.objects.contains_key(&object_name)) } fn paimon_schema(&self, name: &str) -> Option> { @@ -465,7 +588,7 @@ impl PaimonCatalogProvider { } Some(Arc::new( - PaimonSchemaProvider::new( + PaimonSchemaProvider::new_uninitialized( self.catalog_name.clone(), Arc::clone(&self.catalog), name.to_string(), @@ -483,12 +606,14 @@ impl PaimonCatalogProvider { impl CatalogProvider for PaimonCatalogProvider { fn schema_names(&self) -> Vec { - let mut names = self + let mut names: Vec = self .metadata .read() .unwrap_or_else(|e| e.into_inner()) - .database_names - .clone(); + .databases + .keys() + .cloned() + .collect(); let mut temp_names: Vec<_> = self .temp_tables .read() @@ -530,17 +655,13 @@ impl CatalogProvider for PaimonCatalogProvider { .create_database(&name, false, HashMap::new()) .await .map_err(to_datafusion_error)?; - { - let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); - let mut next = (**current).clone(); - if !next.database_names.contains(&name) { - next.database_names.push(name.clone()); - } - next.databases.entry(name.clone()).or_default(); - *current = Arc::new(next); - } + metadata.mutate(|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, @@ -578,15 +699,11 @@ impl CatalogProvider for PaimonCatalogProvider { .drop_database(&name, false, cascade) .await .map_err(to_datafusion_error)?; - { - let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); - let mut next = (**current).clone(); - next.database_names.retain(|database| database != &name); - next.databases.remove(&name); - *current = Arc::new(next); - } + metadata.mutate(|next| { + next.databases.shift_remove(&name); + }); Ok(Some(Arc::new( - PaimonSchemaProvider::new( + PaimonSchemaProvider::new_uninitialized( catalog_name, Arc::clone(&catalog), name, @@ -624,7 +741,7 @@ impl PaimonCatalogProvider { .unwrap_or_else(|e| e.into_inner()) .databases .get(database) - .is_some_and(|metadata| metadata.object_types.contains_key(table_name)) + .is_some_and(|metadata| metadata.objects.contains_key(table_name)) { log::warn!( "Temporary table '{database}.{table_name}' shadows an existing Paimon table" @@ -710,12 +827,31 @@ impl Debug for PaimonSchemaProvider { } impl PaimonSchemaProvider { - /// Creates a new [`PaimonSchemaProvider`]. - /// - /// The provider starts with an empty metadata snapshot. Call - /// [`Self::refresh_metadata`] before exposing it to DataFusion, or prefer - /// [`Self::try_new`] when an async construction boundary is available. - pub fn new( + /// Creates a schema provider with initialized metadata. + pub async fn 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) + } + + /// Creates a schema provider with an empty metadata snapshot. + fn new_uninitialized( catalog_name: Option, catalog: Arc, database: String, @@ -730,7 +866,7 @@ impl PaimonSchemaProvider { database, dynamic_options, temp_provider, - metadata: Arc::new(RwLock::new(Arc::new(CatalogMetadataSnapshot::default()))), + metadata: Arc::new(CatalogMetadataState::default()), blob_reader_registry, session_state, schema_force_view_types: true, @@ -738,7 +874,7 @@ impl PaimonSchemaProvider { } } - /// Creates a schema provider and initializes its metadata snapshot. + /// Backward-compatible alias for [`Self::new`]. pub async fn try_new( catalog_name: Option, catalog: Arc, @@ -748,7 +884,7 @@ impl PaimonSchemaProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, ) -> DFResult { - let provider = Self::new( + Self::new( catalog_name, catalog, database, @@ -756,21 +892,32 @@ impl PaimonSchemaProvider { temp_provider, blob_reader_registry, session_state, - ); - provider.refresh_metadata().await?; - Ok(provider) + ) + .await } /// Refresh this database in the snapshot used by synchronous callbacks. pub async fn refresh_metadata(&self) -> DFResult<()> { - let database = load_database_metadata(self.catalog.as_ref(), &self.database).await?; - let mut current = self.metadata.write().unwrap_or_else(|e| e.into_inner()); - let mut next = (**current).clone(); - if !next.database_names.contains(&self.database) { - next.database_names.push(self.database.clone()); - } - next.databases.insert(self.database.clone(), database); - *current = Arc::new(next); + self.refresh_metadata_inner(false).await + } + + /// Initialize metadata while tolerating a REST server without the optional views endpoint. + pub async fn initialize_metadata(&self) -> DFResult<()> { + self.refresh_metadata_inner(true).await + } + + async fn refresh_metadata_inner(&self, ignore_missing_views_endpoint: bool) -> DFResult<()> { + let generation = self.metadata.begin_refresh(); + let database = load_database_metadata( + self.catalog.as_ref(), + &self.database, + ignore_missing_views_endpoint, + ) + .await?; + self.metadata.publish_update(generation, |next| { + next.databases + .insert(self.database.clone(), Arc::new(database)); + }); Ok(()) } @@ -793,13 +940,13 @@ impl PaimonSchemaProvider { #[async_trait] impl SchemaProvider for PaimonSchemaProvider { fn table_names(&self) -> Vec { - let mut names = self + let mut names: Vec = self .metadata .read() .unwrap_or_else(|e| e.into_inner()) .databases .get(&self.database) - .map(|database| database.object_names.clone()) + .map(|database| database.objects.keys().cloned().collect()) .unwrap_or_default(); if let Some(temp) = &self.temp_provider { @@ -844,6 +991,14 @@ impl SchemaProvider for PaimonSchemaProvider { let schema_force_view_types = self.schema_force_view_types; 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() @@ -879,18 +1034,20 @@ impl SchemaProvider for PaimonSchemaProvider { identifier.full_name() )); } + let metadata_schema = match external.fields() { + Some(fields) => crate::table::datafusion_arrow_schema( + fields, + schema_force_view_types, + )?, + None => Arc::new(datafusion::arrow::datatypes::Schema::empty()), + }; let Some(resolver) = table_engines.get(&declared) else { - let schema = match external.fields() { - Some(fields) => crate::table::datafusion_arrow_schema( - fields, - schema_force_view_types, - )?, - None => Arc::new(datafusion::arrow::datatypes::Schema::empty()), - }; return Ok(Some(Arc::new(UnavailableEngineTableProvider { - schema, - declared, - table_name: identifier.full_name(), + schema: metadata_schema, + 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 @@ -911,12 +1068,19 @@ impl SchemaProvider for PaimonSchemaProvider { declared, )) .await?; - Ok(resolved.map(|inner| { - Arc::new(ReadOnlyTableProvider { + Ok(Some(match resolved { + Some(inner) => Arc::new(ReadOnlyTableProvider { inner, declared, table_name: identifier.full_name(), - }) as Arc + }) as Arc, + None => Arc::new(UnavailableEngineTableProvider { + schema: metadata_schema, + error_message: format!( + "registered table engine did not resolve '{declared}' table '{}'", + identifier.full_name() + ), + }) as Arc, })) } Ok(paimon::catalog::LoadedTable::Paimon(table)) => { @@ -955,16 +1119,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 { .. }) @@ -1030,7 +1184,7 @@ impl SchemaProvider for PaimonSchemaProvider { .unwrap_or_else(|e| e.into_inner()) .databases .get(&self.database) - .and_then(|database| database.object_types.get(name)) + .and_then(|database| database.objects.get(name)) { return Ok(Some(*table_type)); } @@ -1060,6 +1214,9 @@ impl SchemaProvider for PaimonSchemaProvider { { return false; } + if object.system_table().is_some() { + return false; + } // This callback cannot await an external engine resolver. Treat a catalog // declaration as existing and let async `table()` surface resolver absence @@ -1069,7 +1226,7 @@ impl SchemaProvider for PaimonSchemaProvider { .unwrap_or_else(|e| e.into_inner()) .databases .get(&self.database) - .is_some_and(|database| database.object_types.contains_key(object.table())) + .is_some_and(|database| database.objects.contains_key(object.table())) } fn register_table( @@ -1101,15 +1258,11 @@ impl SchemaProvider for PaimonSchemaProvider { .drop_table(&identifier, false) .await .map_err(to_datafusion_error)?; - { - let mut current = metadata.write().unwrap_or_else(|e| e.into_inner()); - let mut next = (**current).clone(); + metadata.mutate(|next| { if let Some(database) = next.databases.get_mut(&database) { - database.object_names.retain(|object| object != &name); - database.object_types.remove(&name); + Arc::make_mut(database).objects.shift_remove(&name); } - *current = Arc::new(next); - } + }); Ok(Some(Arc::new(provider) as Arc)) }, "paimon catalog access thread panicked", diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index a2b55d254..eb4032939 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}; +use std::time::{Duration, Instant}; use datafusion::arrow::array::{ new_null_array, ArrayRef, BooleanArray, Date32Array, Float32Array, Float64Array, Int16Array, @@ -104,6 +105,17 @@ 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>, + metadata_refresh_gate: tokio::sync::Mutex<()>, +} + +const MISSING_OBJECT_REFRESH_TTL: Duration = Duration::from_secs(1); + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum MetadataRefreshTarget { + Catalog(String), + Database { catalog: String, database: String }, } /// Builder for [`SQLContext`]. @@ -171,6 +183,8 @@ impl SQLContextBuilder { catalogs: HashMap::new(), dynamic_options: Default::default(), blob_reader_registry: BlobReaderRegistry::default(), + missing_object_refreshes: Mutex::new(HashMap::new()), + metadata_refresh_gate: tokio::sync::Mutex::new(()), } } } @@ -475,11 +489,17 @@ impl SQLContext { )); } - if self.statement_needs_catalog_metadata(&statements[0])? { - self.refresh_all_catalog_metadata().await?; + let refresh_metadata_after = self.metadata_change_targets(&statements[0], false)?; + let metadata_mutation_targets = self.metadata_change_targets(&statements[0], true)?; + { + let _refresh_guard = self.metadata_refresh_gate.lock().await; + let mut metadata_to_refresh = self.metadata_refresh_targets(&statements[0])?; + metadata_to_refresh.retain(|target| !metadata_mutation_targets.contains(target)); + if !metadata_to_refresh.is_empty() { + self.refresh_metadata_targets(metadata_to_refresh).await?; + } } - let refresh_metadata_after = statement_changes_catalog_metadata(&statements[0]); let result = match &statements[0] { Statement::ShowDatabases { terse, @@ -789,22 +809,43 @@ impl SQLContext { _ => self.ctx.sql(sql).await, }; - if refresh_metadata_after && result.is_ok() { - if let Err(error) = self.refresh_all_catalog_metadata().await { + if result.is_ok() && !refresh_metadata_after.is_empty() { + if let Err(error) = self.refresh_metadata_targets(refresh_metadata_after).await { log::warn!("catalog metadata refresh after DDL failed: {error}"); } } result } - fn statement_needs_catalog_metadata(&self, statement: &Statement) -> DFResult { - if matches!( - statement, - Statement::ShowTables { .. } - | Statement::ShowColumns { .. } - | Statement::ShowFunctions { .. } - ) { - return Ok(true); + fn metadata_refresh_targets( + &self, + statement: &Statement, + ) -> DFResult> { + let mut 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); + } + 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); + } + if matches!(statement, Statement::ShowFunctions { .. }) { + return Ok(targets); } let statement = datafusion::sql::parser::Statement::Statement(Box::new(statement.clone())); @@ -813,10 +854,6 @@ impl SQLContext { 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); - if schema.eq_ignore_ascii_case("information_schema") { - return Ok(true); - } - let catalog_name = reference.catalog().unwrap_or(&default_catalog); if !self.catalogs.contains_key(catalog_name) { continue; @@ -831,17 +868,109 @@ impl SQLContext { "Catalog '{catalog_name}' is not a Paimon catalog" )) })?; - if !provider.metadata_contains_object(schema, reference.table()) { - return Ok(true); + if schema.eq_ignore_ascii_case("information_schema") { + targets.insert(MetadataRefreshTarget::Catalog(catalog_name.to_string())); + } 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(false) + Ok(targets) } - async fn refresh_all_catalog_metadata(&self) -> DFResult<()> { - let catalog_names: Vec<_> = self.catalogs.keys().cloned().collect(); - for catalog_name in catalog_names { - let provider = self.ctx.catalog(&catalog_name).ok_or_else(|| { + fn metadata_change_targets( + &self, + statement: &Statement, + include_temporary: bool, + ) -> 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, _) = self.resolve_catalog_and_database(db_name)?; + targets.insert(MetadataRefreshTarget::Catalog(catalog)); + } + Statement::CreateSchema { + schema_name: SchemaName::Simple(name), + .. + } => { + let (_, catalog, _) = self.resolve_catalog_and_database(name)?; + targets.insert(MetadataRefreshTarget::Catalog(catalog)); + } + Statement::CreateTable(create) if include_temporary || !create.temporary => { + add_database(&create.name)? + } + Statement::CreateView(create) if include_temporary || !create.temporary => { + 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, _) = self.resolve_catalog_and_database(name)?; + targets.insert(MetadataRefreshTarget::Catalog(catalog)); + } + } + Statement::Drop { + object_type: ObjectType::Table | ObjectType::View, + names, + temporary, + .. + } if include_temporary || !temporary => { + for name in names { + add_database(name)?; + } + } + _ => {} + } + Ok(targets) + } + + 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 @@ -851,11 +980,33 @@ impl SQLContext { "Catalog '{catalog_name}' is not a Paimon catalog" )) })?; - provider.refresh_metadata().await?; + 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.initialize_metadata().await?; + } } 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) + } + /// Handle SQL queries containing time-travel syntax (`VERSION AS OF` / `TIMESTAMP AS OF`). /// /// DataFusion's default SQL parser does not support these clauses, so we: @@ -2406,23 +2557,6 @@ impl SQLContext { } } -fn statement_changes_catalog_metadata(statement: &Statement) -> bool { - match statement { - Statement::CreateDatabase { .. } - | Statement::CreateSchema { .. } - | Statement::AlterTable(_) - | Statement::Drop { - object_type: - ObjectType::Database | ObjectType::Schema | ObjectType::Table | ObjectType::View, - temporary: false, - .. - } => true, - Statement::CreateTable(create) => !create.temporary, - Statement::CreateView(create) => !create.temporary, - _ => false, - } -} - fn validate_immutable_scalar_plan(plan: &LogicalPlan) -> DFResult<()> { let mut violation = None; plan.apply(|node| { diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index f867b07d8..07e5bbd10 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -763,7 +763,9 @@ async fn test_missing_database_returns_no_schema() { 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 bd90f41b1..e394f206d 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -34,6 +34,7 @@ use paimon::table::{BranchManager, SnapshotManager, TagManager}; use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options}; use paimon_datafusion::{PaimonCatalogProvider, 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"); @@ -55,6 +56,21 @@ struct MetadataListingCatalog { 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, + blocked_list_tables_started: Notify, + release_blocked_list_tables: Notify, table_names: Mutex>, } @@ -65,10 +81,32 @@ impl MetadataListingCatalog { 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), + blocked_list_tables_started: Notify::new(), + release_blocked_list_tables: Notify::new(), table_names: Mutex::new(vec!["metadata_only".to_string()]), } } + 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 get_table_calls(&self) -> usize { self.get_table_calls.load(Ordering::SeqCst) } @@ -89,6 +127,38 @@ impl MetadataListingCatalog { 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 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 record_remote_call(&self) { assert!( !self.reject_remote_calls.load(Ordering::SeqCst), @@ -102,7 +172,7 @@ impl MetadataListingCatalog { impl Catalog for MetadataListingCatalog { async fn list_databases(&self) -> paimon::Result> { self.record_remote_call(); - Ok(vec!["default".to_string()]) + Ok(self.database_names.lock().unwrap().clone()) } async fn create_database( @@ -140,18 +210,73 @@ impl Catalog for MetadataListingCatalog { }) } - async fn list_tables(&self, _database_name: &str) -> paimon::Result> { + async fn list_tables(&self, database_name: &str) -> paimon::Result> { self.record_remote_call(); - if self.fail_list_tables.load(Ordering::SeqCst) { + *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(); + } + } + if self.block_next_list_tables.swap(false, Ordering::SeqCst) { + self.blocked_list_tables_started.notify_one(); + self.release_blocked_list_tables.notified().await; + } Ok(self.table_names.lock().unwrap().clone()) } 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()]) } @@ -194,7 +319,7 @@ impl Catalog for MetadataListingCatalog { #[tokio::test] async fn test_refreshed_catalog_callbacks_do_not_access_remote_catalog() { let catalog = Arc::new(MetadataListingCatalog::new()); - let provider = PaimonCatalogProvider::new( + let provider = PaimonCatalogProvider::new_uninitialized( Some("paimon".to_string()), catalog.clone(), Default::default(), @@ -215,6 +340,22 @@ async fn test_refreshed_catalog_callbacks_do_not_access_remote_catalog() { assert_eq!(catalog.metadata_calls(), calls_after_refresh); } +#[tokio::test] +async fn test_public_catalog_constructor_returns_ready_provider() { + let catalog = Arc::new(MetadataListingCatalog::new()); + let provider = PaimonCatalogProvider::new( + Some("paimon".to_string()), + catalog, + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + + assert_eq!(provider.schema_names(), vec!["default"]); +} + #[tokio::test] async fn test_failed_catalog_refresh_preserves_last_good_snapshot() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -237,10 +378,118 @@ async fn test_failed_catalog_refresh_preserves_last_good_snapshot() { 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_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( + let provider = PaimonCatalogProvider::new_uninitialized( Some("paimon".to_string()), catalog.clone(), Default::default(), @@ -529,6 +778,168 @@ async fn test_show_tables_refreshes_catalog_metadata() { 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_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] async fn test_show_tables_preserves_catalog_view_type() { let catalog = Arc::new(MetadataListingCatalog::new()); diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index b6df2da3b..904ad6531 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -571,13 +571,14 @@ async fn system_tables_on_routed_tables_error() { } #[tokio::test] -async fn table_exist_uses_catalog_declarations_without_resolving_engines() { +async fn table_exist_agrees_with_routed_table_resolution() { 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_exist("ghost")); - assert!(schema.table_exist("it$snapshots")); + assert!(schema.table("ghost").await.unwrap().is_some()); + assert!(!schema.table_exist("it$snapshots")); } #[tokio::test] From 4f68a94b9b133fdd4e1782d5c2f80a2684ed7049 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sun, 20 Sep 2026 00:35:32 -0700 Subject: [PATCH 3/7] fix(datafusion): address metadata refresh review --- crates/integrations/datafusion/src/catalog.rs | 176 ++++++--- .../datafusion/src/sql_context.rs | 218 ++++++++-- .../datafusion/tests/sql_context_tests.rs | 372 +++++++++++++++++- .../datafusion/tests/table_type_routing.rs | 32 +- 4 files changed, 701 insertions(+), 97 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 13ae91916..fd3602661 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -67,7 +67,7 @@ struct DatabaseMetadata { struct CatalogMetadataState { snapshot: RwLock>, next_generation: AtomicU64, - published_generation: AtomicU64, + database_generations: RwLock>, } impl Default for CatalogMetadataState { @@ -75,7 +75,7 @@ impl Default for CatalogMetadataState { Self { snapshot: RwLock::new(Arc::new(CatalogMetadataSnapshot::default())), next_generation: AtomicU64::new(0), - published_generation: AtomicU64::new(0), + database_generations: RwLock::new(HashMap::new()), } } } @@ -93,30 +93,68 @@ impl CatalogMetadataState { self.next_generation.fetch_add(1, Ordering::AcqRel) + 1 } - fn publish(&self, generation: u64, next: Arc) { + fn publish(&self, generation: u64, refreshed: CatalogMetadataSnapshot) { let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); - if generation >= self.published_generation.load(Ordering::Acquire) { - *current = next; - self.published_generation - .store(generation, Ordering::Release); + let mut database_generations = self + .database_generations + .write() + .unwrap_or_else(|e| e.into_inner()); + let mut next = (**current).clone(); + 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, generation); + } + } + 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, generation); + } } + *current = Arc::new(next); } - fn publish_update(&self, generation: u64, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { + fn publish_database(&self, generation: u64, database: String, metadata: Arc) { let mut current = self.snapshot.write().unwrap_or_else(|e| e.into_inner()); - if generation < self.published_generation.load(Ordering::Acquire) { + 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; } let mut next = (**current).clone(); - update(&mut next); + next.databases.insert(database.clone(), metadata); *current = Arc::new(next); - self.published_generation - .store(generation, Ordering::Release); + database_generations.insert(database, generation); } - fn mutate(&self, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { + fn mutate_database(&self, database: &str, update: impl FnOnce(&mut CatalogMetadataSnapshot)) { let generation = self.begin_refresh(); - self.publish_update(generation, update); + 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); + *current = Arc::new(next); + database_generations.insert(database.to_string(), generation); } } @@ -495,8 +533,8 @@ impl PaimonCatalogProvider { }) .collect(); - let next = Arc::new(CatalogMetadataSnapshot { databases }); - self.metadata.publish(generation, next); + self.metadata + .publish(generation, CatalogMetadataSnapshot { databases }); Ok(()) } @@ -505,10 +543,11 @@ impl PaimonCatalogProvider { let generation = self.metadata.begin_refresh(); let database_metadata = load_database_metadata(self.catalog.as_ref(), database, true).await?; - self.metadata.publish_update(generation, |next| { - next.databases - .insert(database.to_string(), Arc::new(database_metadata)); - }); + self.metadata.publish_database( + generation, + database.to_string(), + Arc::new(database_metadata), + ); Ok(()) } @@ -572,6 +611,52 @@ impl PaimonCatalogProvider { .is_some_and(|metadata| metadata.objects.contains_key(&object_name)) } + pub(crate) fn record_object_created(&self, database: &str, name: &str, table_type: TableType) { + self.metadata.mutate_database(database, |next| { + let database = next + .databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + Arc::make_mut(database) + .objects + .insert(name.to_string(), table_type); + }); + } + + 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); + }); + } + + 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) { + Arc::make_mut(metadata).objects.shift_remove(name); + } + }); + } + + pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) { + self.metadata.mutate_database(database, |next| { + let metadata = next + .databases + .entry(database.to_string()) + .or_insert_with(|| Arc::new(DatabaseMetadata::default())); + let objects = &mut Arc::make_mut(metadata).objects; + let table_type = objects.shift_remove(from).unwrap_or(TableType::Base); + objects.insert(to.to_string(), table_type); + }); + } + fn paimon_schema(&self, name: &str) -> Option> { let temp_provider = { let databases = self.temp_tables.read().unwrap_or_else(|e| e.into_inner()); @@ -655,7 +740,7 @@ impl CatalogProvider for PaimonCatalogProvider { .create_database(&name, false, HashMap::new()) .await .map_err(to_datafusion_error)?; - metadata.mutate(|next| { + metadata.mutate_database(&name, |next| { next.databases .entry(name.clone()) .or_insert_with(|| Arc::new(DatabaseMetadata::default())); @@ -699,7 +784,7 @@ impl CatalogProvider for PaimonCatalogProvider { .drop_database(&name, false, cascade) .await .map_err(to_datafusion_error)?; - metadata.mutate(|next| { + metadata.mutate_database(&name, |next| { next.databases.shift_remove(&name); }); Ok(Some(Arc::new( @@ -914,10 +999,8 @@ impl PaimonSchemaProvider { ignore_missing_views_endpoint, ) .await?; - self.metadata.publish_update(generation, |next| { - next.databases - .insert(self.database.clone(), Arc::new(database)); - }); + self.metadata + .publish_database(generation, self.database.clone(), Arc::new(database)); Ok(()) } @@ -1034,16 +1117,16 @@ impl SchemaProvider for PaimonSchemaProvider { identifier.full_name() )); } - let metadata_schema = match external.fields() { - Some(fields) => crate::table::datafusion_arrow_schema( - fields, - schema_force_view_types, - )?, - None => Arc::new(datafusion::arrow::datatypes::Schema::empty()), - }; let Some(resolver) = table_engines.get(&declared) else { + let schema = match external.fields() { + Some(fields) => crate::table::datafusion_arrow_schema( + fields, + schema_force_view_types, + )?, + None => Arc::new(datafusion::arrow::datatypes::Schema::empty()), + }; return Ok(Some(Arc::new(UnavailableEngineTableProvider { - schema: metadata_schema, + schema, error_message: format!( "no table engine is registered for '{declared}' tables ('{}')", identifier.full_name() @@ -1068,19 +1151,12 @@ impl SchemaProvider for PaimonSchemaProvider { declared, )) .await?; - Ok(Some(match resolved { - Some(inner) => Arc::new(ReadOnlyTableProvider { + Ok(resolved.map(|inner| { + Arc::new(ReadOnlyTableProvider { inner, declared, table_name: identifier.full_name(), - }) as Arc, - None => Arc::new(UnavailableEngineTableProvider { - schema: metadata_schema, - error_message: format!( - "registered table engine did not resolve '{declared}' table '{}'", - identifier.full_name() - ), - }) as Arc, + }) as Arc })) } Ok(paimon::catalog::LoadedTable::Paimon(table)) => { @@ -1214,13 +1290,9 @@ impl SchemaProvider for PaimonSchemaProvider { { return false; } - if object.system_table().is_some() { - return false; - } - - // This callback cannot await an external engine resolver. Treat a catalog - // declaration as existing and let async `table()` surface resolver absence - // or unsupported system-table access during planning. + // System tables derive their existence from a snapshotted base table. + // This callback cannot await an external engine resolver, so async + // `table()` remains responsible for surfacing unsupported routed tables. self.metadata .read() .unwrap_or_else(|e| e.into_inner()) @@ -1258,7 +1330,7 @@ impl SchemaProvider for PaimonSchemaProvider { .drop_table(&identifier, false) .await .map_err(to_datafusion_error)?; - metadata.mutate(|next| { + metadata.mutate_database(&database, |next| { if let Some(database) = next.databases.get_mut(&database) { Arc::make_mut(database).objects.shift_remove(&name); } diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index eb4032939..2130e5b6e 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -45,7 +45,7 @@ //! - `TRUNCATE TABLE db.t PARTITION (col = val, ...)` use std::collections::{HashMap, HashSet}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, Weak}; use std::time::{Duration, Instant}; use datafusion::arrow::array::{ @@ -61,7 +61,7 @@ use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::SessionStateBuilder; -use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, Volatility}; +use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, TableType, Volatility}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::planner::IdentNormalizer; use datafusion::sql::sqlparser::ast::{ @@ -107,7 +107,7 @@ pub struct SQLContext { blob_reader_registry: BlobReaderRegistry, /// Last successful refresh used to resolve a missing object, keyed by database. missing_object_refreshes: Mutex>, - metadata_refresh_gate: tokio::sync::Mutex<()>, + metadata_refresh_gates: Mutex>>>, } const MISSING_OBJECT_REFRESH_TTL: Duration = Duration::from_secs(1); @@ -184,7 +184,7 @@ impl SQLContextBuilder { dynamic_options: Default::default(), blob_reader_registry: BlobReaderRegistry::default(), missing_object_refreshes: Mutex::new(HashMap::new()), - metadata_refresh_gate: tokio::sync::Mutex::new(()), + metadata_refresh_gates: Mutex::new(HashMap::new()), } } } @@ -489,16 +489,31 @@ impl SQLContext { )); } - let refresh_metadata_after = self.metadata_change_targets(&statements[0], false)?; - let metadata_mutation_targets = self.metadata_change_targets(&statements[0], true)?; - { - let _refresh_guard = self.metadata_refresh_gate.lock().await; - let mut metadata_to_refresh = self.metadata_refresh_targets(&statements[0])?; - metadata_to_refresh.retain(|target| !metadata_mutation_targets.contains(target)); - if !metadata_to_refresh.is_empty() { - self.refresh_metadata_targets(metadata_to_refresh).await?; + 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); + if let Err(error) = self.refresh_metadata_targets([target]).await { + if best_effort { + log::warn!("information schema metadata refresh failed: {error}"); + } else { + return Err(error); + } + } } - } + Ok::<(), DataFusionError>(()) + })) + .await?; let result = match &statements[0] { Statement::ShowDatabases { @@ -809,10 +824,8 @@ impl SQLContext { _ => self.ctx.sql(sql).await, }; - if result.is_ok() && !refresh_metadata_after.is_empty() { - if let Err(error) = self.refresh_metadata_targets(refresh_metadata_after).await { - log::warn!("catalog metadata refresh after DDL failed: {error}"); - } + if result.is_ok() { + self.apply_metadata_change(&statements[0])?; } result } @@ -820,15 +833,19 @@ impl SQLContext { fn metadata_refresh_targets( &self, statement: &Statement, - ) -> DFResult> { + ) -> 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); + return Ok((targets, best_effort_targets)); } if let Statement::ShowColumns { show_options, .. } = statement { if let Some(show_in) = &show_options.show_in { @@ -842,10 +859,10 @@ impl SQLContext { } } } - return Ok(targets); + return Ok((targets, best_effort_targets)); } if matches!(statement, Statement::ShowFunctions { .. }) { - return Ok(targets); + return Ok((targets, best_effort_targets)); } let statement = datafusion::sql::parser::Statement::Statement(Box::new(statement.clone())); @@ -869,7 +886,13 @@ impl SQLContext { )) })?; if schema.eq_ignore_ascii_case("information_schema") { - targets.insert(MetadataRefreshTarget::Catalog(catalog_name.to_string())); + let catalogs = self + .catalogs + .keys() + .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(), @@ -880,13 +903,12 @@ impl SQLContext { } } } - Ok(targets) + Ok((targets, best_effort_targets)) } fn metadata_change_targets( &self, statement: &Statement, - include_temporary: bool, ) -> DFResult> { let mut targets = HashSet::new(); let mut add_database = |name: &ObjectName| -> DFResult<()> { @@ -914,12 +936,8 @@ impl SQLContext { let (_, catalog, _) = self.resolve_catalog_and_database(name)?; targets.insert(MetadataRefreshTarget::Catalog(catalog)); } - Statement::CreateTable(create) if include_temporary || !create.temporary => { - add_database(&create.name)? - } - Statement::CreateView(create) if include_temporary || !create.temporary => { - add_database(&create.name)? - } + 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, @@ -935,9 +953,8 @@ impl SQLContext { Statement::Drop { object_type: ObjectType::Table | ObjectType::View, names, - temporary, .. - } if include_temporary || !temporary => { + } => { for name in names { add_database(name)?; } @@ -947,6 +964,129 @@ impl SQLContext { Ok(targets) } + 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)?; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_object_created( + identifier.database(), + identifier.object(), + TableType::Base, + ) + })?; + 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_object_created( + identifier.database(), + identifier.object(), + TableType::View, + ) + })?; + Ok(()) + } + Statement::AlterTable(alter) => { + let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&alter.name)?; + for operation in &alter.operations { + if let AlterTableOperation::RenameTable { table_name } = operation { + let name = match table_name { + RenameTableNameKind::To(name) | RenameTableNameKind::As(name) => { + object_name_to_string(name) + } + }; + self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_table_renamed( + identifier.database(), + identifier.object(), + &name, + ) + })?; + } + } + 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), + ) -> 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")) + })?; + update(provider); + Ok(()) + } + async fn refresh_metadata_targets( &self, targets: impl IntoIterator, @@ -1007,6 +1147,20 @@ impl SQLContext { .is_some_and(|refreshed| refreshed.elapsed() < MISSING_OBJECT_REFRESH_TTL) } + 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`). /// /// DataFusion's default SQL parser does not support these clauses, so we: diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index e394f206d..3a9f2c0ea 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -69,9 +69,11 @@ struct MetadataListingCatalog { 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>>, } impl MetadataListingCatalog { @@ -94,9 +96,11 @@ impl MetadataListingCatalog { 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()), } } @@ -107,6 +111,11 @@ impl MetadataListingCatalog { catalog } + 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) } @@ -123,6 +132,13 @@ impl MetadataListingCatalog { *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); } @@ -159,6 +175,10 @@ impl MetadataListingCatalog { 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 record_remote_call(&self) { assert!( !self.reject_remote_calls.load(Ordering::SeqCst), @@ -265,11 +285,29 @@ impl Catalog for MetadataListingCatalog { self.parallel_database_listing_started.notify_waiters(); } } - if self.block_next_list_tables.swap(false, Ordering::SeqCst) { + 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(self.table_names.lock().unwrap().clone()) + Ok(table_names) } async fn list_views(&self, _database_name: &str) -> paimon::Result> { @@ -411,6 +449,89 @@ async fn test_older_refresh_cannot_overwrite_newer_snapshot() { 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() + .refresh_metadata() + .await + }); + catalog.blocked_list_tables_started.notified().await; + + 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 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_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()); @@ -892,6 +1013,78 @@ async fn test_create_table_does_not_refresh_unrelated_catalogs() { 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_repeated_missing_tables_share_negative_refresh_ttl() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -940,6 +1133,66 @@ async fn test_concurrent_missing_tables_share_single_refresh() { 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()); @@ -961,6 +1214,121 @@ async fn test_show_tables_preserves_catalog_view_type() { 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::new(); + 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_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::new(); + 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"]); +} + +#[tokio::test] +async fn test_information_schema_does_not_hide_strict_object_refresh_failure() { + 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 error = sql_context + .sql( + "SELECT missing.* FROM paimon.information_schema.tables info \ + JOIN paimon.default.missing AS missing ON true", + ) + .await + .unwrap_err(); + + 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::new(); + 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] async fn test_select_branch_table_reads_branch_snapshot() { 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 904ad6531..344fe48ee 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -309,6 +309,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 +323,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,14 +572,22 @@ async fn system_tables_on_routed_tables_error() { } #[tokio::test] -async fn table_exist_agrees_with_routed_table_resolution() { +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_exist("ghost")); - assert!(schema.table("ghost").await.unwrap().is_some()); - assert!(!schema.table_exist("it$snapshots")); + assert!(schema.table("ghost").await.unwrap().is_none()); +} + +#[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_exist("missing$snapshots")); + assert!(!schema.table_exist("pt$not_a_system_table")); } #[tokio::test] @@ -1025,15 +1034,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)" )) From 1b705e02eab22edc04520e74a5b1ce6b6770d44c Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sun, 20 Sep 2026 02:47:24 -0700 Subject: [PATCH 4/7] fix(datafusion): coordinate metadata refreshes --- crates/integrations/datafusion/src/catalog.rs | 454 +++++++++++++++--- .../datafusion/src/sql_context.rs | 145 +++++- .../datafusion/tests/read_tables.rs | 2 +- .../datafusion/tests/sql_context_tests.rs | 357 +++++++++++++- .../datafusion/tests/table_type_routing.rs | 28 +- 5 files changed, 910 insertions(+), 76 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index fd3602661..f2efb6e1a 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -17,18 +17,17 @@ //! Paimon catalog integration for DataFusion. -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt::Debug; use std::ops::Deref; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::Arc; -use std::sync::RwLock; +use std::sync::{Arc, Mutex, RwLock}; 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; @@ -63,11 +62,20 @@ struct DatabaseMetadata { objects: IndexMap, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ObjectResolution { + Paimon, + Routed, + Unavailable, +} + #[derive(Debug)] struct CatalogMetadataState { snapshot: RwLock>, next_generation: AtomicU64, database_generations: RwLock>, + active_refreshes: Mutex>, + object_resolutions: RwLock>, } impl Default for CatalogMetadataState { @@ -76,10 +84,29 @@ impl Default for CatalogMetadataState { 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>; @@ -89,17 +116,50 @@ impl Deref for CatalogMetadataState { } impl CatalogMetadataState { - fn begin_refresh(&self) -> u64 { + fn next_generation(&self) -> u64 { self.next_generation.fetch_add(1, Ordering::AcqRel) + 1 } - fn publish(&self, generation: u64, refreshed: CatalogMetadataSnapshot) { + 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 @@ -112,19 +172,33 @@ impl CatalogMetadataState { for name in removed_names { if database_generations.get(&name).copied().unwrap_or_default() <= generation { next.databases.shift_remove(&name); - database_generations.insert(name, generation); + 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, generation); + 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) { + 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 @@ -136,16 +210,24 @@ impl CatalogMetadataState { .unwrap_or_default() > generation { - return; + 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.begin_refresh(); + 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 @@ -153,14 +235,114 @@ impl CatalogMetadataState { .unwrap_or_else(|e| e.into_inner()); let mut next = (**current).clone(); update(&mut next); - *current = Arc::new(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 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; +// 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; async fn load_database_metadata( catalog: &dyn Catalog, @@ -427,23 +609,24 @@ impl Debug for PaimonCatalogProvider { } impl PaimonCatalogProvider { - /// Creates a provider with an initialized metadata snapshot. - pub async fn new( + /// Creates a provider with an empty metadata snapshot. + /// + /// This preserves the original synchronous constructor signature. New callers should use + /// [`Self::try_new`] so discovery callbacks are not exposed before metadata is initialized. + pub fn new( catalog_name: Option, catalog: Arc, dynamic_options: DynamicOptions, blob_reader_registry: BlobReaderRegistry, session_state: Option, - ) -> DFResult { - let provider = Self::new_uninitialized( + ) -> Self { + Self::new_uninitialized( catalog_name, catalog, dynamic_options, blob_reader_registry, session_state, - ); - provider.initialize_metadata().await?; - Ok(provider) + ) } /// Creates a provider with an empty metadata snapshot. @@ -469,7 +652,7 @@ impl PaimonCatalogProvider { } } - /// Backward-compatible alias for [`Self::new`]. + /// Creates a provider with an initialized metadata snapshot. pub async fn try_new( catalog_name: Option, catalog: Arc, @@ -477,14 +660,15 @@ impl PaimonCatalogProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, ) -> DFResult { - Self::new( + let provider = Self::new_uninitialized( catalog_name, catalog, dynamic_options, blob_reader_registry, session_state, - ) - .await + ); + provider.initialize_metadata().await?; + Ok(provider) } /// Refresh the metadata consumed by DataFusion's synchronous catalog callbacks. @@ -533,22 +717,63 @@ impl PaimonCatalogProvider { }) .collect(); - self.metadata - .publish(generation, CatalogMetadataSnapshot { databases }); + let mut conflicts = self + .metadata + .publish(generation.get(), CatalogMetadataSnapshot { databases }); + if !conflicts.is_empty() { + let current_databases: HashSet<_> = 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| async move { + self.refresh_database_metadata_inner( + database.as_str(), + ignore_missing_views_endpoint, + ) + .await + }) + .buffer_unordered(MAX_CONCURRENT_METADATA_LISTINGS) + .try_collect::>() + .await?; Ok(()) } /// Refresh one database in the metadata snapshot. pub(crate) async fn refresh_database_metadata(&self, database: &str) -> DFResult<()> { - let generation = self.metadata.begin_refresh(); - let database_metadata = - load_database_metadata(self.catalog.as_ref(), database, true).await?; - self.metadata.publish_database( - generation, - database.to_string(), - Arc::new(database_metadata), - ); - Ok(()) + self.refresh_database_metadata_inner(database, true).await + } + + async fn refresh_database_metadata_inner( + &self, + database: &str, + ignore_missing_views_endpoint: bool, + ) -> 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, + ) + .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. @@ -621,6 +846,12 @@ impl PaimonCatalogProvider { .objects .insert(name.to_string(), table_type); }); + if table_type == TableType::Base { + self.metadata + .set_object_resolution(database, name, ObjectResolution::Paimon); + } else { + self.metadata.remove_object_resolution(database, name); + } } pub(crate) fn record_database_created(&self, database: &str) { @@ -635,6 +866,7 @@ impl PaimonCatalogProvider { 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) { @@ -643,18 +875,23 @@ impl PaimonCatalogProvider { Arc::make_mut(metadata).objects.shift_remove(name); } }); + self.metadata.remove_object_resolution(database, name); } pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) { + 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 objects = &mut Arc::make_mut(metadata).objects; - let table_type = objects.shift_remove(from).unwrap_or(TableType::Base); - objects.insert(to.to_string(), table_type); + if let Some(metadata) = next.databases.get_mut(database) { + let objects = &mut Arc::make_mut(metadata).objects; + if let Some(table_type) = objects.shift_remove(from) { + objects.insert(to.to_string(), table_type); + renamed = true; + } + } }); + if renamed { + self.metadata.rename_object_resolution(database, from, to); + } } fn paimon_schema(&self, name: &str) -> Option> { @@ -912,8 +1149,11 @@ impl Debug for PaimonSchemaProvider { } impl PaimonSchemaProvider { - /// Creates a schema provider with initialized metadata. - pub async fn new( + /// Creates a schema provider with an empty metadata snapshot. + /// + /// This preserves the original synchronous constructor signature. New callers should use + /// [`Self::try_new`] so discovery callbacks are not exposed before metadata is initialized. + pub fn new( catalog_name: Option, catalog: Arc, database: String, @@ -921,8 +1161,8 @@ impl PaimonSchemaProvider { temp_provider: Option>, blob_reader_registry: BlobReaderRegistry, session_state: Option, - ) -> DFResult { - let provider = Self::new_uninitialized( + ) -> Self { + Self::new_uninitialized( catalog_name, catalog, database, @@ -930,9 +1170,7 @@ impl PaimonSchemaProvider { temp_provider, blob_reader_registry, session_state, - ); - provider.initialize_metadata().await?; - Ok(provider) + ) } /// Creates a schema provider with an empty metadata snapshot. @@ -959,7 +1197,7 @@ impl PaimonSchemaProvider { } } - /// Backward-compatible alias for [`Self::new`]. + /// Creates a schema provider with initialized metadata. pub async fn try_new( catalog_name: Option, catalog: Arc, @@ -969,7 +1207,7 @@ impl PaimonSchemaProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, ) -> DFResult { - Self::new( + let provider = Self::new_uninitialized( catalog_name, catalog, database, @@ -977,8 +1215,9 @@ impl PaimonSchemaProvider { temp_provider, blob_reader_registry, session_state, - ) - .await + ); + provider.initialize_metadata().await?; + Ok(provider) } /// Refresh this database in the snapshot used by synchronous callbacks. @@ -999,8 +1238,16 @@ impl PaimonSchemaProvider { ignore_missing_views_endpoint, ) .await?; - self.metadata - .publish_database(generation, self.database.clone(), Arc::new(database)); + 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 + ))); + } Ok(()) } @@ -1072,6 +1319,7 @@ 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() @@ -1090,6 +1338,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::Paimon, + ); if branch.is_some() { return Err(plan_datafusion_err!( "branches are not supported for 'object-table' tables ('{}')", @@ -1110,6 +1363,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 ('{}')", @@ -1151,6 +1409,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, @@ -1160,6 +1425,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 @@ -1253,6 +1523,11 @@ 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 .metadata @@ -1290,15 +1565,25 @@ impl SchemaProvider for PaimonSchemaProvider { { return false; } - // System tables derive their existence from a snapshotted base table. - // This callback cannot await an external engine resolver, so async - // `table()` remains responsible for surfacing unsupported routed tables. - self.metadata + let table_type = self + .metadata .read() .unwrap_or_else(|e| e.into_inner()) .databases .get(&self.database) - .is_some_and(|database| database.objects.contains_key(object.table())) + .and_then(|database| database.objects.get(object.table())) + .copied(); + let resolution = self + .metadata + .object_resolution(&self.database, object.table()); + if resolution == Some(ObjectResolution::Unavailable) { + return false; + } + if object.system_table().is_some() { + return table_type == Some(TableType::Base) + && resolution != Some(ObjectResolution::Routed); + } + table_type.is_some() } fn register_table( @@ -1620,6 +1905,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 2130e5b6e..bfb575432 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -107,12 +107,23 @@ pub struct SQLContext { 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, } 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, PartialEq)] +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] enum MetadataRefreshTarget { Catalog(String), Database { catalog: String, database: String }, @@ -145,6 +156,9 @@ enum MetadataRefreshTarget { #[derive(Default)] pub struct SQLContextBuilder { runtime_env: Option>, + catalog_metadata_refresh_ttl: Option, + catalog_metadata_refresh_timeout: Option, + max_concurrent_catalog_metadata_refreshes: Option, } impl SQLContextBuilder { @@ -159,6 +173,28 @@ 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 + } + /// Builds a [`SQLContext`]. pub fn build(self) -> SQLContext { let mut state_builder = SessionStateBuilder::new() @@ -184,7 +220,19 @@ impl SQLContextBuilder { 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), + )), } } } @@ -285,6 +333,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 { @@ -503,8 +555,38 @@ impl SQLContext { current_targets.retain(|target| !metadata_mutation_targets.contains(target)); if current_targets.contains(&target) { let best_effort = best_effort_targets.contains(&target); - if let Err(error) = self.refresh_metadata_targets([target]).await { + 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); @@ -515,6 +597,17 @@ impl SQLContext { })) .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, @@ -889,6 +982,7 @@ impl SQLContext { let catalogs = self .catalogs .keys() + .filter(|catalog| !self.catalog_refresh_is_recent(catalog)) .cloned() .map(MetadataRefreshTarget::Catalog); targets.extend(catalogs.clone()); @@ -923,18 +1017,17 @@ impl SQLContext { }); Ok(()) }; - match statement { Statement::CreateDatabase { db_name, .. } => { - let (_, catalog, _) = self.resolve_catalog_and_database(db_name)?; - targets.insert(MetadataRefreshTarget::Catalog(catalog)); + 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, _) = self.resolve_catalog_and_database(name)?; - targets.insert(MetadataRefreshTarget::Catalog(catalog)); + 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)?, @@ -946,8 +1039,8 @@ impl SQLContext { .. } => { for name in names { - let (_, catalog, _) = self.resolve_catalog_and_database(name)?; - targets.insert(MetadataRefreshTarget::Catalog(catalog)); + let (_, catalog, database) = self.resolve_catalog_and_database(name)?; + targets.insert(MetadataRefreshTarget::Database { catalog, database }); } } Statement::Drop { @@ -1131,6 +1224,14 @@ impl SQLContext { ); } else { provider.initialize_metadata().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(()) @@ -1147,6 +1248,32 @@ impl SQLContext { .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 diff --git a/crates/integrations/datafusion/tests/read_tables.rs b/crates/integrations/datafusion/tests/read_tables.rs index 07e5bbd10..0b476b73f 100644 --- a/crates/integrations/datafusion/tests/read_tables.rs +++ b/crates/integrations/datafusion/tests/read_tables.rs @@ -757,7 +757,7 @@ 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(), diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 3a9f2c0ea..ae5a2aa32 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -32,7 +32,7 @@ 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; @@ -74,6 +74,34 @@ struct MetadataListingCatalog { release_blocked_list_tables: Notify, table_names: Mutex>, table_names_by_database: Mutex>>, + listing_concurrency: Option>, + block_drop_response: AtomicBool, + drop_committed: Notify, + release_drop_response: Notify, +} + +#[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 { @@ -101,6 +129,10 @@ impl MetadataListingCatalog { 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()), + listing_concurrency: None, + block_drop_response: AtomicBool::new(false), + drop_committed: Notify::new(), + release_drop_response: Notify::new(), } } @@ -111,6 +143,13 @@ impl MetadataListingCatalog { 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(); @@ -179,6 +218,10 @@ impl MetadataListingCatalog { *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 record_remote_call(&self) { assert!( !self.reject_remote_calls.load(Ordering::SeqCst), @@ -219,6 +262,10 @@ 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(()) } @@ -232,6 +279,9 @@ impl Catalog for MetadataListingCatalog { 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() @@ -332,6 +382,10 @@ 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(()) } @@ -381,7 +435,7 @@ async fn test_refreshed_catalog_callbacks_do_not_access_remote_catalog() { #[tokio::test] async fn test_public_catalog_constructor_returns_ready_provider() { let catalog = Arc::new(MetadataListingCatalog::new()); - let provider = PaimonCatalogProvider::new( + let provider = PaimonCatalogProvider::try_new( Some("paimon".to_string()), catalog, Default::default(), @@ -394,6 +448,27 @@ async fn test_public_catalog_constructor_returns_ready_provider() { assert_eq!(provider.schema_names(), vec!["default"]); } +#[test] +fn test_public_provider_constructors_remain_synchronous() { + let catalog: Arc = Arc::new(MetadataListingCatalog::new()); + let _: PaimonCatalogProvider = PaimonCatalogProvider::new( + Some("paimon".to_string()), + Arc::clone(&catalog), + Default::default(), + Default::default(), + None, + ); + let _: PaimonSchemaProvider = PaimonSchemaProvider::new( + Some("paimon".to_string()), + catalog, + "default".to_string(), + Default::default(), + None, + Default::default(), + None, + ); +} + #[tokio::test] async fn test_failed_catalog_refresh_preserves_last_good_snapshot() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -494,6 +569,42 @@ async fn test_full_refresh_merges_non_conflicting_database_updates() { ); } +#[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![ @@ -1085,6 +1196,112 @@ async fn test_ddl_applies_exact_snapshot_deltas() { ); } +#[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(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()); @@ -1218,7 +1435,9 @@ async fn test_show_tables_preserves_catalog_view_type() { 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::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()) @@ -1238,11 +1457,57 @@ async fn test_information_schema_refreshes_all_paimon_catalogs() { 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); +} + +#[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); +} + #[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::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 @@ -1266,10 +1531,88 @@ async fn test_information_schema_keeps_last_good_snapshot_on_catalog_refresh_fai assert_eq!(table_names, vec!["healthy_new"]); } +#[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_information_schema_times_out_pending_catalog_refresh() { + let catalog = Arc::new(MetadataListingCatalog::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(); + + 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(); +} + +#[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(); + + sql_context + .sql("SELECT * FROM first.information_schema.tables") + .await + .unwrap(); + + assert_eq!(listing_concurrency.maximum(), 2); +} + #[tokio::test] 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 @@ -1296,7 +1639,9 @@ async fn test_information_schema_does_not_hide_strict_object_refresh_failure() { 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::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 diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index 344fe48ee..36db8c746 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -28,7 +28,9 @@ 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"; @@ -577,6 +579,27 @@ async fn registered_engine_none_is_not_found() { 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")); + 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] @@ -586,6 +609,9 @@ async fn table_exist_derives_system_table_from_snapshotted_base_table() { 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")); } From 4019607c8a56832733f303fdf0a6c97ddd6e37eb Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sun, 20 Sep 2026 04:20:25 -0700 Subject: [PATCH 5/7] fix(datafusion): close metadata snapshot review gaps --- crates/integrations/datafusion/src/catalog.rs | 221 ++++++++++++------ .../datafusion/src/sql_context.rs | 44 ++-- .../datafusion/tests/sql_context_tests.rs | 120 +++++++++- .../datafusion/tests/table_type_routing.rs | 54 +++++ crates/paimon/src/catalog/filesystem.rs | 18 ++ crates/paimon/src/catalog/mod.rs | 17 ++ .../paimon/src/catalog/rest/rest_catalog.rs | 29 ++- crates/paimon/tests/rest_catalog_test.rs | 34 ++- 8 files changed, 438 insertions(+), 99 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index f2efb6e1a..6b29af268 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -60,11 +60,20 @@ struct CatalogMetadataSnapshot { #[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, } @@ -255,6 +264,16 @@ impl CatalogMetadataState { .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() @@ -348,6 +367,7 @@ async fn load_database_metadata( catalog: &dyn Catalog, database: &str, ignore_missing_views_endpoint: bool, + known_capabilities: HashMap, ) -> DFResult { let tables = async { catalog @@ -375,14 +395,51 @@ async fn load_database_metadata( }; let (table_names, view_names) = futures::try_join!(tables, views)?; + let unresolved_names: Vec<_> = table_names + .iter() + .filter(|name| !known_capabilities.contains_key(*name)) + .cloned() + .collect(); + let declared_types = if unresolved_names.is_empty() { + HashMap::new() + } else { + match catalog.list_table_types(database, &unresolved_names).await { + Ok(types) => types, + Err(error) => { + log::debug!( + "unable to classify table types while refreshing database '{database}': \ + {error}" + ); + HashMap::new() + } + } + }; + 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 = known_capabilities.get(&name).copied().unwrap_or_else(|| { + 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 }) + Ok(DatabaseMetadata { + objects, + system_table_capabilities, + }) } /// What an engine is asked to resolve. Non-exhaustive so later releases can @@ -609,26 +666,6 @@ impl Debug for PaimonCatalogProvider { } impl PaimonCatalogProvider { - /// Creates a provider with an empty metadata snapshot. - /// - /// This preserves the original synchronous constructor signature. New callers should use - /// [`Self::try_new`] so discovery callbacks are not exposed before metadata is initialized. - pub fn new( - catalog_name: Option, - catalog: Arc, - dynamic_options: DynamicOptions, - blob_reader_registry: BlobReaderRegistry, - session_state: Option, - ) -> Self { - Self::new_uninitialized( - catalog_name, - catalog, - dynamic_options, - blob_reader_registry, - session_state, - ) - } - /// Creates a provider with an empty metadata snapshot. /// /// Callers must refresh it before exposing synchronous discovery callbacks. @@ -653,6 +690,26 @@ impl PaimonCatalogProvider { } /// 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, @@ -695,14 +752,18 @@ impl PaimonCatalogProvider { database_names.retain(|name| seen_databases.insert(name.clone())); let entries: HashMap<_, _> = stream::iter(database_names.iter().cloned()) - .map(|database| async move { - let metadata = load_database_metadata( - self.catalog.as_ref(), - database.as_str(), - ignore_missing_views_endpoint, - ) - .await?; - Ok::<_, datafusion::error::DataFusionError>((database, metadata)) + .map(|database| { + let known_capabilities = self.metadata.system_table_capabilities(&database); + async move { + let metadata = load_database_metadata( + self.catalog.as_ref(), + database.as_str(), + ignore_missing_views_endpoint, + known_capabilities, + ) + .await?; + Ok::<_, datafusion::error::DataFusionError>((database, metadata)) + } }) .buffer_unordered(MAX_CONCURRENT_METADATA_LISTINGS) .try_collect() @@ -761,6 +822,7 @@ impl PaimonCatalogProvider { self.catalog.as_ref(), database, ignore_missing_views_endpoint, + self.metadata.system_table_capabilities(database), ) .await?; if self.metadata.publish_database( @@ -842,9 +904,15 @@ impl PaimonCatalogProvider { .databases .entry(database.to_string()) .or_insert_with(|| Arc::new(DatabaseMetadata::default())); - Arc::make_mut(database) - .objects - .insert(name.to_string(), table_type); + let database = Arc::make_mut(database); + database.objects.insert(name.to_string(), table_type); + if table_type == TableType::Base { + database + .system_table_capabilities + .insert(name.to_string(), SystemTableCapability::Paimon); + } else { + database.system_table_capabilities.remove(name); + } }); if table_type == TableType::Base { self.metadata @@ -872,7 +940,9 @@ impl PaimonCatalogProvider { 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) { - Arc::make_mut(metadata).objects.shift_remove(name); + let metadata = Arc::make_mut(metadata); + metadata.objects.shift_remove(name); + metadata.system_table_capabilities.remove(name); } }); self.metadata.remove_object_resolution(database, name); @@ -882,9 +952,15 @@ impl PaimonCatalogProvider { let mut renamed = false; self.metadata.mutate_database(database, |next| { if let Some(metadata) = next.databases.get_mut(database) { - let objects = &mut Arc::make_mut(metadata).objects; + 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); + if let Some(capability) = metadata.system_table_capabilities.remove(from) { + metadata + .system_table_capabilities + .insert(to.to_string(), capability); + } renamed = true; } } @@ -1150,31 +1226,7 @@ impl Debug for PaimonSchemaProvider { impl PaimonSchemaProvider { /// Creates a schema provider with an empty metadata snapshot. - /// - /// This preserves the original synchronous constructor signature. New callers should use - /// [`Self::try_new`] so discovery callbacks are not exposed before metadata is initialized. - pub fn new( - catalog_name: Option, - catalog: Arc, - database: String, - dynamic_options: DynamicOptions, - temp_provider: Option>, - blob_reader_registry: BlobReaderRegistry, - session_state: Option, - ) -> Self { - Self::new_uninitialized( - catalog_name, - catalog, - database, - dynamic_options, - temp_provider, - blob_reader_registry, - session_state, - ) - } - - /// Creates a schema provider with an empty metadata snapshot. - fn new_uninitialized( + pub fn new_uninitialized( catalog_name: Option, catalog: Arc, database: String, @@ -1198,6 +1250,28 @@ impl PaimonSchemaProvider { } /// 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, @@ -1236,6 +1310,7 @@ impl PaimonSchemaProvider { self.catalog.as_ref(), &self.database, ignore_missing_views_endpoint, + self.metadata.system_table_capabilities(&self.database), ) .await?; if !self.metadata.publish_database( @@ -1341,7 +1416,7 @@ impl SchemaProvider for PaimonSchemaProvider { metadata.set_object_resolution( identifier.database(), identifier.object(), - ObjectResolution::Paimon, + ObjectResolution::Object, ); if branch.is_some() { return Err(plan_datafusion_err!( @@ -1565,14 +1640,22 @@ impl SchemaProvider for PaimonSchemaProvider { { return false; } - let table_type = self + let (table_type, snapshot_capability) = self .metadata .read() .unwrap_or_else(|e| e.into_inner()) .databases .get(&self.database) - .and_then(|database| database.objects.get(object.table())) - .copied(); + .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()); @@ -1580,8 +1663,16 @@ impl SchemaProvider for PaimonSchemaProvider { return false; } if object.system_table().is_some() { - return table_type == Some(TableType::Base) - && resolution != Some(ObjectResolution::Routed); + 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() } diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index bfb575432..3cdadec6e 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -1106,26 +1106,6 @@ impl SQLContext { })?; Ok(()) } - Statement::AlterTable(alter) => { - let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&alter.name)?; - for operation in &alter.operations { - if let AlterTableOperation::RenameTable { table_name } = operation { - let name = match table_name { - RenameTableNameKind::To(name) | RenameTableNameKind::As(name) => { - object_name_to_string(name) - } - }; - self.update_catalog_metadata(&catalog_name, |provider| { - provider.record_table_renamed( - identifier.database(), - identifier.object(), - &name, - ) - })?; - } - } - Ok(()) - } Statement::Drop { object_type: ObjectType::Database | ObjectType::Schema, names, @@ -1838,7 +1818,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 @@ -1967,10 +1947,24 @@ 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(()) => self.update_catalog_metadata(&catalog_name, |provider| { + provider.record_table_renamed( + identifier.database(), + identifier.object(), + new_identifier.object(), + ) + })?, + 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() { @@ -7576,7 +7570,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/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index ae5a2aa32..89255abd6 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::{ @@ -74,10 +74,13 @@ struct MetadataListingCatalog { release_blocked_list_tables: Notify, table_names: Mutex>, table_names_by_database: Mutex>>, + list_table_types_calls: AtomicUsize, listing_concurrency: Option>, block_drop_response: AtomicBool, drop_committed: Notify, release_drop_response: Notify, + rename_source_missing: AtomicBool, + rename_ignore_flags: Mutex>, } #[derive(Default)] @@ -129,10 +132,13 @@ impl MetadataListingCatalog { 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), 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()), } } @@ -195,6 +201,10 @@ impl MetadataListingCatalog { .unwrap_or_default() } + fn list_table_types_calls(&self) -> usize { + self.list_table_types_calls.load(Ordering::SeqCst) + } + fn race_next_refreshes(&self) { self.racing_list_tables_calls.store(0, Ordering::SeqCst); self.race_refreshes.store(true, Ordering::SeqCst); @@ -222,6 +232,14 @@ impl MetadataListingCatalog { 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), @@ -360,6 +378,19 @@ impl Catalog for MetadataListingCatalog { 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); + 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) { @@ -391,10 +422,22 @@ impl Catalog for MetadataListingCatalog { async fn rename_table( &self, - _from: &Identifier, + from: &Identifier, _to: &Identifier, - _ignore_if_not_exists: bool, + 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(), + }); + } Ok(()) } @@ -446,19 +489,52 @@ async fn test_public_catalog_constructor_returns_ready_provider() { .unwrap(); assert_eq!(provider.schema_names(), vec!["default"]); + assert_eq!( + provider.schema("default").unwrap().table_names(), + vec!["metadata_only", "metadata_view"] + ); +} + +#[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); + + 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); } -#[test] -fn test_public_provider_constructors_remain_synchronous() { +#[tokio::test] +async fn test_explicit_uninitialized_providers_require_metadata_initialization() { let catalog: Arc = Arc::new(MetadataListingCatalog::new()); - let _: PaimonCatalogProvider = PaimonCatalogProvider::new( + let catalog_provider = PaimonCatalogProvider::new_uninitialized( Some("paimon".to_string()), Arc::clone(&catalog), Default::default(), Default::default(), None, ); - let _: PaimonSchemaProvider = PaimonSchemaProvider::new( + 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(), @@ -467,6 +543,12 @@ fn test_public_provider_constructors_remain_synchronous() { 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] @@ -1216,6 +1298,30 @@ async fn test_alter_table_if_exists_does_not_create_phantom_rename_delta() { 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(flavor = "multi_thread", worker_threads = 2)] async fn test_concurrent_ddl_delta_follows_serialized_commit_order() { let catalog = Arc::new(MetadataListingCatalog::new()); diff --git a/crates/integrations/datafusion/tests/table_type_routing.rs b/crates/integrations/datafusion/tests/table_type_routing.rs index 36db8c746..8a56b5b3a 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -22,6 +22,7 @@ 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}; @@ -616,6 +617,59 @@ async fn table_exist_derives_system_table_from_snapshotted_base_table() { 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] async fn paimon_managed_types_cannot_be_routed() { let env = setup().await; diff --git a/crates/paimon/src/catalog/filesystem.rs b/crates/paimon/src/catalog/filesystem.rs index 743f00286..6838fe83f 100644 --- a/crates/paimon/src/catalog/filesystem.rs +++ b/crates/paimon/src/catalog/filesystem.rs @@ -35,6 +35,7 @@ use crate::table::{ObjectTable, SchemaManager, Table}; use async_trait::async_trait; use bytes::Bytes; use chrono::TimeZone; +use futures::{stream, StreamExt, TryStreamExt}; use opendal::raw::get_basename; use snafu::OptionExt; @@ -426,6 +427,23 @@ impl Catalog for FileSystemCatalog { Ok(tables) } + async fn list_table_types( + &self, + database_name: &str, + table_names: &[String], + ) -> Result> { + stream::iter(table_names.iter().cloned()) + .map(|table_name| async move { + 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()?; + Ok((table_name, table_type)) + }) + .buffer_unordered(16) + .try_collect() + .await + } + 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..c5ff2d72c 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -444,6 +444,23 @@ 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 can contain non-Paimon table types should override this method with an + /// implementation backed by their metadata store. The default preserves compatibility for + /// catalogs that only implement the original Paimon-only [`Self::list_tables`] contract. + async fn list_table_types( + &self, + _database_name: &str, + table_names: &[String], + ) -> Result> { + Ok(table_names + .iter() + .cloned() + .map(|name| (name, TableType::Table)) + .collect()) + } + /// 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..bd54d2e80 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -24,6 +24,7 @@ use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; +use futures::{stream, StreamExt, TryStreamExt}; use crate::api::rest_api::RESTApi; use crate::api::rest_error::RestError; @@ -37,7 +38,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 +327,32 @@ 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> { + stream::iter(table_names.iter().cloned()) + .map(|table_name| async move { + 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(); + Ok((table_name, table_type)) + }) + .buffer_unordered(16) + .try_collect() + .await + } + 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; From 30711dcf24a844e0489d1d43844acbbfacb8a182 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sun, 20 Sep 2026 19:21:42 -0700 Subject: [PATCH 6/7] fix(datafusion): bound metadata discovery requests --- crates/integrations/datafusion/src/catalog.rs | 231 ++++++++++++++---- .../datafusion/src/sql_context.rs | 94 +++++-- .../datafusion/tests/sql_context_tests.rs | 193 ++++++++++++++- .../datafusion/tests/table_type_routing.rs | 18 ++ crates/paimon/src/catalog/filesystem.rs | 19 +- crates/paimon/src/catalog/mod.rs | 14 +- .../paimon/src/catalog/rest/rest_catalog.rs | 37 ++- 7 files changed, 494 insertions(+), 112 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 6b29af268..39439de06 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -359,6 +359,7 @@ impl CatalogMetadataState { type SharedCatalogMetadata = Arc; const MAX_CONCURRENT_METADATA_LISTINGS: usize = 16; +pub(crate) const DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS: usize = 16; // 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; @@ -368,14 +369,21 @@ async fn load_database_metadata( database: &str, ignore_missing_views_endpoint: bool, known_capabilities: HashMap, + metadata_io_semaphore: &tokio::sync::Semaphore, ) -> 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![]), @@ -397,29 +405,47 @@ async fn load_database_metadata( let unresolved_names: Vec<_> = table_names .iter() - .filter(|name| !known_capabilities.contains_key(*name)) + .filter(|name| { + !matches!( + known_capabilities.get(*name), + Some(SystemTableCapability::Paimon | SystemTableCapability::Unsupported) + ) + }) .cloned() .collect(); - let declared_types = if unresolved_names.is_empty() { - HashMap::new() - } else { - match catalog.list_table_types(database, &unresolved_names).await { - Ok(types) => types, - Err(error) => { - log::debug!( - "unable to classify table types while refreshing database '{database}': \ - {error}" - ); - HashMap::new() + let declared_types: HashMap<_, _> = stream::iter(unresolved_names) + .map(|name| async move { + 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) => Ok::<_, DataFusionError>( + types.remove(&name).map(|table_type| (name, table_type)), + ), + Err(error) => { + log::debug!("unable to classify table type for '{database}.{name}': {error}"); + Ok::<_, DataFusionError>(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 = known_capabilities.get(&name).copied().unwrap_or_else(|| { - declared_types + 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() { @@ -428,8 +454,8 @@ async fn load_database_metadata( SystemTableCapability::Paimon } }) - .unwrap_or(SystemTableCapability::Unknown) - }); + .unwrap_or(SystemTableCapability::Unknown), + }; system_table_capabilities.insert(name.clone(), capability); objects.entry(name).or_insert(TableType::Base); } @@ -657,6 +683,8 @@ pub struct PaimonCatalogProvider { 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, } impl Debug for PaimonCatalogProvider { @@ -675,6 +703,26 @@ impl PaimonCatalogProvider { 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, + )), + ) + } + + 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, ) -> Self { PaimonCatalogProvider { catalog_name, @@ -686,6 +734,7 @@ impl PaimonCatalogProvider { schema_force_view_types: true, table_engines: Arc::new(RwLock::new(HashMap::new())), metadata: Arc::new(CatalogMetadataState::default()), + metadata_io_semaphore, } } @@ -717,12 +766,34 @@ impl PaimonCatalogProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, ) -> DFResult { - let provider = Self::new_uninitialized( + 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, + )), + ) + .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, + ) -> DFResult { + let provider = Self::new_uninitialized_with_metadata_io_semaphore( + catalog_name, + catalog, + dynamic_options, + blob_reader_registry, + session_state, + metadata_io_semaphore, ); provider.initialize_metadata().await?; Ok(provider) @@ -743,11 +814,15 @@ impl PaimonCatalogProvider { async fn refresh_metadata_inner(&self, ignore_missing_views_endpoint: bool) -> DFResult<()> { let generation = self.metadata.begin_refresh(); - let mut database_names = self - .catalog - .list_databases() - .await - .map_err(to_datafusion_error)?; + 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())); @@ -760,6 +835,7 @@ impl PaimonCatalogProvider { database.as_str(), ignore_missing_views_endpoint, known_capabilities, + self.metadata_io_semaphore.as_ref(), ) .await?; Ok::<_, datafusion::error::DataFusionError>((database, metadata)) @@ -782,13 +858,17 @@ impl PaimonCatalogProvider { .metadata .publish(generation.get(), CatalogMetadataSnapshot { databases }); if !conflicts.is_empty() { - let current_databases: HashSet<_> = self - .catalog - .list_databases() - .await - .map_err(to_datafusion_error)? - .into_iter() - .collect(); + 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) @@ -823,6 +903,7 @@ impl PaimonCatalogProvider { database, ignore_missing_views_endpoint, self.metadata.system_table_capabilities(database), + self.metadata_io_semaphore.as_ref(), ) .await?; if self.metadata.publish_database( @@ -898,30 +979,56 @@ impl PaimonCatalogProvider { .is_some_and(|metadata| metadata.objects.contains_key(&object_name)) } - pub(crate) fn record_object_created(&self, database: &str, name: &str, table_type: TableType) { + 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(), table_type); - if table_type == TableType::Base { - database - .system_table_capabilities - .insert(name.to_string(), SystemTableCapability::Paimon); - } else { - database.system_table_capabilities.remove(name); - } + 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); }); - if table_type == TableType::Base { - self.metadata - .set_object_resolution(database, name, ObjectResolution::Paimon); - } else { - self.metadata.remove_object_resolution(database, name); + 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) 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 @@ -948,7 +1055,7 @@ impl PaimonCatalogProvider { self.metadata.remove_object_resolution(database, name); } - pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) { + pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) -> bool { let mut renamed = false; self.metadata.mutate_database(database, |next| { if let Some(metadata) = next.databases.get_mut(database) { @@ -956,11 +1063,13 @@ impl PaimonCatalogProvider { let objects = &mut metadata.objects; if let Some(table_type) = objects.shift_remove(from) { objects.insert(to.to_string(), table_type); - if let Some(capability) = metadata.system_table_capabilities.remove(from) { - metadata - .system_table_capabilities - .insert(to.to_string(), capability); - } + let capability = metadata + .system_table_capabilities + .remove(from) + .unwrap_or(SystemTableCapability::Unknown); + metadata + .system_table_capabilities + .insert(to.to_string(), capability); renamed = true; } } @@ -968,6 +1077,7 @@ impl PaimonCatalogProvider { if renamed { self.metadata.rename_object_resolution(database, from, to); } + renamed } fn paimon_schema(&self, name: &str) -> Option> { @@ -995,6 +1105,7 @@ impl PaimonCatalogProvider { self.blob_reader_registry.clone(), self.session_state.clone(), ) + .with_metadata_io_semaphore(Arc::clone(&self.metadata_io_semaphore)) .with_schema_force_view_types(self.schema_force_view_types) .with_table_engines(self.table_engines()) .with_metadata_snapshot(Arc::clone(&self.metadata)), @@ -1046,6 +1157,7 @@ impl CatalogProvider for PaimonCatalogProvider { 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 { @@ -1068,6 +1180,7 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) + .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), @@ -1090,6 +1203,7 @@ impl CatalogProvider for PaimonCatalogProvider { 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 { @@ -1110,6 +1224,7 @@ impl CatalogProvider for PaimonCatalogProvider { blob_reader_registry, session_state, ) + .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), @@ -1213,6 +1328,8 @@ pub struct PaimonSchemaProvider { 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, } impl Debug for PaimonSchemaProvider { @@ -1246,6 +1363,9 @@ impl PaimonSchemaProvider { 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, + )), } } @@ -1311,6 +1431,7 @@ impl PaimonSchemaProvider { &self.database, ignore_missing_views_endpoint, self.metadata.system_table_capabilities(&self.database), + self.metadata_io_semaphore.as_ref(), ) .await?; if !self.metadata.publish_database( @@ -1331,6 +1452,14 @@ impl PaimonSchemaProvider { self } + fn with_metadata_io_semaphore( + mut self, + metadata_io_semaphore: Arc, + ) -> Self { + self.metadata_io_semaphore = metadata_io_semaphore; + self + } + pub(crate) fn with_table_engines(mut self, table_engines: TableEngines) -> Self { self.table_engines = table_engines; self diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 3cdadec6e..78e6c2d8a 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -61,7 +61,7 @@ use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::SessionStateBuilder; -use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, TableType, Volatility}; +use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, Volatility}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion::sql::planner::IdentNormalizer; use datafusion::sql::sqlparser::ast::{ @@ -115,6 +115,7 @@ pub struct SQLContext { 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); @@ -159,6 +160,7 @@ pub struct SQLContextBuilder { catalog_metadata_refresh_ttl: Option, catalog_metadata_refresh_timeout: Option, max_concurrent_catalog_metadata_refreshes: Option, + max_concurrent_catalog_metadata_requests: Option, } impl SQLContextBuilder { @@ -195,6 +197,16 @@ impl SQLContextBuilder { 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() @@ -233,6 +245,10 @@ impl SQLContextBuilder { 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), + )), } } } @@ -316,12 +332,13 @@ impl SQLContext { let session_state: crate::catalog::SessionStateProvider = Arc::new(move || weak_state.upgrade().map(|state| state.read().clone())); let provider = Arc::new( - crate::catalog::PaimonCatalogProvider::try_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), ) .await?, ); @@ -1082,11 +1099,23 @@ impl SQLContext { return Ok(()); } let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&create.name)?; + let declared_type = if create.if_not_exists { + None + } else { + let options: HashMap<_, _> = extract_options(&create.table_options)? + .into_iter() + .collect(); + Some( + CoreOptions::new(&options) + .table_type() + .map_err(to_datafusion_error)?, + ) + }; self.update_catalog_metadata(&catalog_name, |provider| { - provider.record_object_created( + provider.record_table_created( identifier.database(), identifier.object(), - TableType::Base, + declared_type, ) })?; Ok(()) @@ -1098,11 +1127,7 @@ impl SQLContext { } let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&create.name)?; self.update_catalog_metadata(&catalog_name, |provider| { - provider.record_object_created( - identifier.database(), - identifier.object(), - TableType::View, - ) + provider.record_view_created(identifier.database(), identifier.object()) })?; Ok(()) } @@ -1142,10 +1167,27 @@ impl SQLContext { } } - fn update_catalog_metadata( + fn update_catalog_metadata( &self, catalog_name: &str, - update: impl FnOnce(&crate::catalog::PaimonCatalogProvider), + 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 @@ -1156,8 +1198,7 @@ impl SQLContext { .ok_or_else(|| { DataFusionError::Plan(format!("Catalog '{catalog_name}' is not a Paimon catalog")) })?; - update(provider); - Ok(()) + provider.refresh_database_metadata(database).await } async fn refresh_metadata_targets( @@ -1951,13 +1992,26 @@ impl SQLContext { .rename_table(&identifier, &new_identifier, false) .await { - Ok(()) => self.update_catalog_metadata(&catalog_name, |provider| { - provider.record_table_renamed( - identifier.database(), - identifier.object(), - new_identifier.object(), - ) - })?, + 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()) diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 89255abd6..713192a9d 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -75,6 +75,7 @@ struct MetadataListingCatalog { table_names: Mutex>, table_names_by_database: Mutex>>, list_table_types_calls: AtomicUsize, + fail_next_list_table_types: AtomicBool, listing_concurrency: Option>, block_drop_response: AtomicBool, drop_committed: Notify, @@ -133,6 +134,7 @@ impl MetadataListingCatalog { 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), listing_concurrency: None, block_drop_response: AtomicBool::new(false), drop_committed: Notify::new(), @@ -205,6 +207,11 @@ impl MetadataListingCatalog { 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 race_next_refreshes(&self) { self.racing_list_tables_calls.store(0, Ordering::SeqCst); self.race_refreshes.store(true, Ordering::SeqCst); @@ -384,6 +391,17 @@ impl Catalog for MetadataListingCatalog { 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_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() @@ -423,7 +441,7 @@ impl Catalog for MetadataListingCatalog { async fn rename_table( &self, from: &Identifier, - _to: &Identifier, + to: &Identifier, ignore_if_not_exists: bool, ) -> paimon::Result<()> { self.rename_ignore_flags @@ -438,6 +456,18 @@ impl Catalog for MetadataListingCatalog { 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(()) } @@ -520,6 +550,31 @@ async fn test_metadata_refresh_classifies_only_new_objects() { assert_eq!(catalog.list_table_types_calls(), 2); } +#[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); + + provider.refresh_metadata().await.unwrap(); + + assert_eq!(catalog.list_table_types_calls(), 2); + let schema = provider.schema("default").unwrap(); + assert!(schema.table_exist("existing$snapshots")); +} + #[tokio::test] async fn test_explicit_uninitialized_providers_require_metadata_initialization() { let catalog: Arc = Arc::new(MetadataListingCatalog::new()); @@ -1322,6 +1377,29 @@ async fn test_alter_table_if_exists_reconciles_stale_positive_rename_source() { 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(flavor = "multi_thread", worker_threads = 2)] async fn test_concurrent_ddl_delta_follows_serialized_commit_order() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -1713,6 +1791,67 @@ async fn test_information_schema_bounds_cross_catalog_refresh_concurrency() { 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_information_schema_does_not_hide_strict_object_refresh_failure() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -2507,6 +2646,58 @@ async fn test_create_table_if_not_exists() { .expect("Second CREATE with IF NOT EXISTS should succeed"); } +#[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_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 8a56b5b3a..1c67d6626 100644 --- a/crates/integrations/datafusion/tests/table_type_routing.rs +++ b/crates/integrations/datafusion/tests/table_type_routing.rs @@ -1234,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 6838fe83f..d28f0df17 100644 --- a/crates/paimon/src/catalog/filesystem.rs +++ b/crates/paimon/src/catalog/filesystem.rs @@ -35,7 +35,6 @@ use crate::table::{ObjectTable, SchemaManager, Table}; use async_trait::async_trait; use bytes::Bytes; use chrono::TimeZone; -use futures::{stream, StreamExt, TryStreamExt}; use opendal::raw::get_basename; use snafu::OptionExt; @@ -432,16 +431,14 @@ impl Catalog for FileSystemCatalog { database_name: &str, table_names: &[String], ) -> Result> { - stream::iter(table_names.iter().cloned()) - .map(|table_name| async move { - 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()?; - Ok((table_name, table_type)) - }) - .buffer_unordered(16) - .try_collect() - .await + 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( diff --git a/crates/paimon/src/catalog/mod.rs b/crates/paimon/src/catalog/mod.rs index c5ff2d72c..9d367c436 100644 --- a/crates/paimon/src/catalog/mod.rs +++ b/crates/paimon/src/catalog/mod.rs @@ -446,19 +446,15 @@ pub trait Catalog: Send + Sync { /// Return the declared type for the requested tables. /// - /// Catalogs that can contain non-Paimon table types should override this method with an - /// implementation backed by their metadata store. The default preserves compatibility for - /// catalogs that only implement the original Paimon-only [`Self::list_tables`] contract. + /// 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], + _table_names: &[String], ) -> Result> { - Ok(table_names - .iter() - .cloned() - .map(|name| (name, TableType::Table)) - .collect()) + Ok(HashMap::new()) } /// Create a table. diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index bd54d2e80..0cad642d7 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -24,7 +24,6 @@ use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; -use futures::{stream, StreamExt, TryStreamExt}; use crate::api::rest_api::RESTApi; use crate::api::rest_error::RestError; @@ -332,25 +331,23 @@ impl Catalog for RESTCatalog { database_name: &str, table_names: &[String], ) -> Result> { - stream::iter(table_names.iter().cloned()) - .map(|table_name| async move { - 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(); - Ok((table_name, table_type)) - }) - .buffer_unordered(16) - .try_collect() - .await + 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( From 1dccb9779d1f1dccf3a0957ba7a38bd877c6950e Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:30:59 -0700 Subject: [PATCH 7/7] fix(datafusion): harden metadata reconciliation --- crates/integrations/datafusion/src/catalog.rs | 313 +++++++++++++++--- .../datafusion/src/sql_context.rs | 43 ++- .../datafusion/tests/sql_context_tests.rs | 180 +++++++++- 3 files changed, 472 insertions(+), 64 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 39439de06..d88a2c305 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -20,8 +20,9 @@ use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt::Debug; use std::ops::Deref; -use std::sync::atomic::{AtomicU64, Ordering}; +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}; @@ -359,17 +360,89 @@ impl CatalogMetadataState { 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(|_| { @@ -401,7 +474,11 @@ async fn load_database_metadata( Err(error) => Err(to_datafusion_error(error)), } }; - let (table_names, view_names) = futures::try_join!(tables, views)?; + 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() @@ -414,20 +491,34 @@ async fn load_database_metadata( .cloned() .collect(); let declared_types: HashMap<_, _> = stream::iter(unresolved_names) - .map(|name| async move { - 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) => Ok::<_, DataFusionError>( - types.remove(&name).map(|table_type| (name, table_type)), - ), - Err(error) => { - log::debug!("unable to classify table type for '{database}.{name}': {error}"); - Ok::<_, DataFusionError>(None) + .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) + } } } }) @@ -685,6 +776,7 @@ pub struct PaimonCatalogProvider { 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 { @@ -713,6 +805,7 @@ impl PaimonCatalogProvider { Arc::new(tokio::sync::Semaphore::new( DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, )), + DEFAULT_METADATA_REFRESH_TIMEOUT, ) } @@ -723,6 +816,7 @@ impl PaimonCatalogProvider { blob_reader_registry: BlobReaderRegistry, session_state: Option, metadata_io_semaphore: Arc, + metadata_refresh_timeout: Duration, ) -> Self { PaimonCatalogProvider { catalog_name, @@ -735,6 +829,7 @@ impl PaimonCatalogProvider { table_engines: Arc::new(RwLock::new(HashMap::new())), metadata: Arc::new(CatalogMetadataState::default()), metadata_io_semaphore, + metadata_refresh_timeout, } } @@ -775,6 +870,7 @@ impl PaimonCatalogProvider { Arc::new(tokio::sync::Semaphore::new( DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, )), + DEFAULT_METADATA_REFRESH_TIMEOUT, ) .await } @@ -786,6 +882,7 @@ impl PaimonCatalogProvider { 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, @@ -794,6 +891,7 @@ impl PaimonCatalogProvider { blob_reader_registry, session_state, metadata_io_semaphore, + metadata_refresh_timeout, ); provider.initialize_metadata().await?; Ok(provider) @@ -804,16 +902,49 @@ impl PaimonCatalogProvider { /// 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.refresh_metadata_inner(false).await + 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.refresh_metadata_inner(true).await + 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 refresh_metadata_inner(&self, ignore_missing_views_endpoint: bool) -> DFResult<()> { + 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()) @@ -829,6 +960,7 @@ impl PaimonCatalogProvider { 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(), @@ -836,6 +968,7 @@ impl PaimonCatalogProvider { ignore_missing_views_endpoint, known_capabilities, self.metadata_io_semaphore.as_ref(), + Arc::clone(&classification_budget), ) .await?; Ok::<_, datafusion::error::DataFusionError>((database, metadata)) @@ -872,28 +1005,43 @@ impl PaimonCatalogProvider { conflicts.retain(|database| current_databases.contains(database)); } stream::iter(conflicts) - .map(|database| async move { - self.refresh_database_metadata_inner( - database.as_str(), - ignore_missing_views_endpoint, - ) - .await + .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?; - Ok(()) + 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<()> { - self.refresh_database_metadata_inner(database, true).await + 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 { @@ -904,6 +1052,7 @@ impl PaimonCatalogProvider { 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( @@ -1016,6 +1165,37 @@ impl PaimonCatalogProvider { } } + 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 @@ -1058,24 +1238,34 @@ impl PaimonCatalogProvider { pub(crate) fn record_table_renamed(&self, database: &str, from: &str, to: &str) -> bool { let mut renamed = false; self.metadata.mutate_database(database, |next| { - if let Some(metadata) = next.databases.get_mut(database) { - 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; - } + 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 } @@ -1106,6 +1296,7 @@ impl PaimonCatalogProvider { 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)), @@ -1330,6 +1521,7 @@ pub struct PaimonSchemaProvider { 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 { @@ -1366,6 +1558,7 @@ impl PaimonSchemaProvider { metadata_io_semaphore: Arc::new(tokio::sync::Semaphore::new( DEFAULT_MAX_CONCURRENT_METADATA_REQUESTS, )), + metadata_refresh_timeout: DEFAULT_METADATA_REFRESH_TIMEOUT, } } @@ -1416,22 +1609,44 @@ impl PaimonSchemaProvider { /// Refresh this database in the snapshot used by synchronous callbacks. pub async fn refresh_metadata(&self) -> DFResult<()> { - self.refresh_metadata_inner(false).await + 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.refresh_metadata_inner(true).await + self.run_metadata_refresh(self.refresh_metadata_inner(true, false)) + .await } - async fn refresh_metadata_inner(&self, ignore_missing_views_endpoint: bool) -> DFResult<()> { + 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( @@ -1444,7 +1659,10 @@ impl PaimonSchemaProvider { self.database ))); } - Ok(()) + 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 { @@ -1460,6 +1678,11 @@ impl PaimonSchemaProvider { 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 diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 78e6c2d8a..67f5c2118 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -339,6 +339,7 @@ impl SQLContext { self.blob_reader_registry.clone(), Some(session_state), Arc::clone(&self.catalog_metadata_request_semaphore), + self.catalog_metadata_refresh_timeout, ) .await?, ); @@ -935,7 +936,7 @@ impl SQLContext { }; if result.is_ok() { - self.apply_metadata_change(&statements[0])?; + self.apply_metadata_change(&statements[0]).await?; } result } @@ -1074,7 +1075,7 @@ impl SQLContext { Ok(targets) } - fn apply_metadata_change(&self, statement: &Statement) -> DFResult<()> { + 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)?; @@ -1099,23 +1100,33 @@ impl SQLContext { return Ok(()); } let (_, catalog_name, identifier) = self.resolve_catalog_and_table(&create.name)?; - let declared_type = if create.if_not_exists { - None - } else { - let options: HashMap<_, _> = extract_options(&create.table_options)? - .into_iter() - .collect(); - Some( - CoreOptions::new(&options) - .table_type() - .map_err(to_datafusion_error)?, - ) - }; + 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(), - declared_type, + Some(declared_type), ) })?; Ok(()) @@ -1244,7 +1255,7 @@ impl SQLContext { Instant::now(), ); } else { - provider.initialize_metadata().await?; + provider.refresh_metadata_for_sql().await?; self.catalog_refreshes .lock() .unwrap_or_else(|e| e.into_inner()) diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 713192a9d..aa20c5ab6 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -76,6 +76,7 @@ struct MetadataListingCatalog { 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, @@ -135,6 +136,7 @@ impl MetadataListingCatalog { 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(), @@ -212,6 +214,10 @@ impl MetadataListingCatalog { .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); @@ -394,9 +400,10 @@ impl Catalog for MetadataListingCatalog { if let Some(listing_concurrency) = &self.listing_concurrency { listing_concurrency.observe().await; } - if self - .fail_next_list_table_types - .swap(false, Ordering::SeqCst) + 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(), @@ -550,6 +557,26 @@ async fn test_metadata_refresh_classifies_only_new_objects() { assert_eq!(catalog.list_table_types_calls(), 2); } +#[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"]); + + 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); + let names = provider.schema("default").unwrap().table_names(); + assert_eq!(names.iter().filter(|name| *name == "duplicate").count(), 1); +} + #[tokio::test] async fn test_metadata_refresh_retries_unknown_table_types() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -575,6 +602,32 @@ async fn test_metadata_refresh_retries_unknown_table_types() { assert!(schema.table_exist("existing$snapshots")); } +#[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(); + + let provider = PaimonCatalogProvider::try_new( + Some("paimon".to_string()), + catalog.clone(), + Default::default(), + Default::default(), + None, + ) + .await + .unwrap(); + + assert_eq!(provider.schema("default").unwrap().table_names().len(), 101); + assert!(catalog.list_table_types_calls() <= 16); + + 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); +} + #[tokio::test] async fn test_explicit_uninitialized_providers_require_metadata_initialization() { let catalog: Arc = Arc::new(MetadataListingCatalog::new()); @@ -606,6 +659,24 @@ async fn test_explicit_uninitialized_providers_require_metadata_initialization() ); } +#[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()); @@ -1400,6 +1471,35 @@ async fn test_successful_rename_refreshes_target_when_source_was_not_snapshotted 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()); @@ -1741,6 +1841,34 @@ async fn test_information_schema_backs_off_after_catalog_refresh_failure() { assert_eq!(catalog.metadata_calls(), calls_after_failure); } +#[tokio::test] +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("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + let calls_after_failure = catalog.list_table_types_calls(); + + sql_context + .sql("SELECT * FROM paimon.information_schema.tables") + .await + .unwrap(); + + assert_eq!(catalog.list_table_types_calls(), calls_after_failure); +} + #[tokio::test] async fn test_information_schema_times_out_pending_catalog_refresh() { let catalog = Arc::new(MetadataListingCatalog::new()); @@ -2646,6 +2774,25 @@ 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(); @@ -2698,6 +2845,33 @@ async fn test_create_table_if_not_exists_noop_does_not_overwrite_object_capabili 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();