1use std::collections::{HashMap, HashSet};
40use std::sync::Arc;
41
42use datafusion::arrow::array::{
43 new_null_array, ArrayRef, BooleanArray, Date32Array, Float32Array, Float64Array, Int16Array,
44 Int32Array, Int64Array, Int8Array, StringArray,
45};
46use datafusion::arrow::compute::cast;
47use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
48use datafusion::arrow::record_batch::RecordBatch;
49use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
50use datafusion::common::TableReference;
51use datafusion::datasource::{MemTable, TableProvider};
52use datafusion::error::{DataFusionError, Result as DFResult};
53use datafusion::execution::SessionStateBuilder;
54use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan, Volatility};
55use datafusion::prelude::{DataFrame, SessionContext};
56use datafusion::sql::planner::IdentNormalizer;
57use datafusion::sql::sqlparser::ast::{
58 AlterTableOperation, BinaryLength, CharacterLength, ColumnDef, ColumnOption, CreateFunction,
59 CreateFunctionBody, CreateTable, CreateTableOptions, CreateView, Delete, Expr as SqlExpr,
60 FromTable, FunctionBehavior, FunctionReturnType, Insert, Merge, ObjectName, ObjectType,
61 RenameTableNameKind, Reset, ResetStatement, Set, ShowCreateObject, SqlOption, Statement,
62 TableFactor, TableObject, Truncate, Update, Value as SqlValue,
63};
64use datafusion::sql::sqlparser::dialect::GenericDialect;
65use datafusion::sql::sqlparser::keywords::Keyword;
66use datafusion::sql::sqlparser::parser::Parser;
67use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
68use futures::StreamExt;
69use paimon::catalog::{parse_object_name, Catalog, Identifier};
70use paimon::spec::{
71 ArrayType as PaimonArrayType, BigIntType, BinaryType, BlobType, BooleanType, CharType,
72 DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType,
73 DoubleType, FloatType, IntType, LocalZonedTimestampType, MapType as PaimonMapType,
74 RowType as PaimonRowType, SchemaChange, SmallIntType, TimestampType, TinyIntType,
75 VarBinaryType, VarCharType, VariantType,
76};
77
78use crate::error::to_datafusion_error;
79use crate::table_loader::load_table_for_read;
80use crate::{BlobReaderRegistry, DynamicOptions};
81
82pub struct SQLContext {
92 ctx: SessionContext,
93 catalogs: HashMap<String, Arc<dyn Catalog>>,
94 dynamic_options: DynamicOptions,
96 blob_reader_registry: BlobReaderRegistry,
97}
98
99impl Default for SQLContext {
100 fn default() -> Self {
101 Self::new()
102 }
103}
104
105impl SQLContext {
106 pub fn new() -> Self {
108 let state = SessionStateBuilder::new()
109 .with_config(crate::lateral_vector_search::session_config())
110 .with_default_features()
111 .with_relation_planners(vec![Arc::new(
112 crate::relation_planner::PaimonRelationPlanner::new(),
113 )])
114 .with_optimizer_rules(crate::lateral_vector_search::optimizer_rules())
115 .with_query_planner(Arc::new(
116 crate::lateral_vector_search::PaimonQueryPlanner::new(),
117 ))
118 .build();
119 let ctx = SessionContext::new_with_state(state);
120 crate::blob_descriptor_functions::register_blob_descriptor_functions(&ctx);
121 crate::variant_functions::register_variant_functions(&ctx);
122 Self {
123 ctx,
124 catalogs: HashMap::new(),
125 dynamic_options: Default::default(),
126 blob_reader_registry: BlobReaderRegistry::default(),
127 }
128 }
129
130 pub fn blob_reader_registry(&self) -> BlobReaderRegistry {
131 self.blob_reader_registry.clone()
132 }
133
134 pub async fn register_catalog(
141 &mut self,
142 catalog_name: impl Into<String>,
143 catalog: Arc<dyn Catalog>,
144 ) -> DFResult<()> {
145 self.register_catalog_with_default_db(catalog_name, catalog, Some("default"))
146 .await
147 }
148
149 pub async fn register_catalog_with_default_db(
164 &mut self,
165 catalog_name: impl Into<String>,
166 catalog: Arc<dyn Catalog>,
167 default_db: Option<&str>,
168 ) -> DFResult<()> {
169 if matches!(default_db, Some("")) {
170 return Err(DataFusionError::Plan(
171 "default_db must not be empty; pass None to skip default-database init".to_string(),
172 ));
173 }
174 let catalog_name = catalog_name.into();
175 let is_first = self.catalogs.is_empty();
176 if let Some(default_db) = default_db {
177 match catalog.get_database(default_db).await {
178 Ok(_) => {}
179 Err(paimon::Error::DatabaseNotExist { .. }) => {
180 catalog
181 .create_database(default_db, true, Default::default())
182 .await
183 .map_err(|e| DataFusionError::External(Box::new(e)))?;
184 }
185 Err(e) => return Err(DataFusionError::External(Box::new(e))),
186 }
187 }
188 let weak_state = self.ctx.state_weak_ref();
189 let session_state: crate::catalog::SessionStateProvider =
190 Arc::new(move || weak_state.upgrade().map(|state| state.read().clone()));
191 self.ctx.register_catalog(
192 &catalog_name,
193 Arc::new(crate::catalog::PaimonCatalogProvider::new(
194 Some(catalog_name.clone()),
195 catalog.clone(),
196 self.dynamic_options.clone(),
197 self.blob_reader_registry.clone(),
198 Some(session_state),
199 )),
200 );
201 register_table_functions(&self.ctx, &catalog, default_db.unwrap_or("default"));
202 self.catalogs.insert(catalog_name.clone(), catalog);
203 if is_first {
204 self.set_current_catalog(catalog_name).await?;
205 if let Some(default_db) = default_db {
206 self.set_current_database(default_db).await?;
207 }
208 }
209 Ok(())
210 }
211
212 pub async fn set_current_catalog(&mut self, catalog_name: impl Into<String>) -> DFResult<()> {
214 let catalog_name = catalog_name.into();
215 if !self.catalogs.contains_key(&catalog_name) {
216 return Err(DataFusionError::Plan(format!(
217 "Unknown catalog '{catalog_name}'"
218 )));
219 }
220 if catalog_name.contains('\'') {
221 return Err(DataFusionError::Plan(
222 "Catalog name must not contain single quotes".to_string(),
223 ));
224 }
225 self.ctx
226 .sql(&format!(
227 "SET datafusion.catalog.default_catalog = '{catalog_name}'"
228 ))
229 .await?;
230 Ok(())
231 }
232
233 pub async fn set_current_database(&self, database_name: &str) -> DFResult<()> {
235 if database_name.contains('\'') {
236 return Err(DataFusionError::Plan(
237 "Database name must not contain single quotes".to_string(),
238 ));
239 }
240 self.ctx
241 .sql(&format!(
242 "SET datafusion.catalog.default_schema = '{database_name}'"
243 ))
244 .await?;
245 Ok(())
246 }
247
248 pub fn ctx(&self) -> &SessionContext {
250 &self.ctx
251 }
252
253 pub fn register_temp_table(
262 &self,
263 name: impl Into<TableReference>,
264 table: Arc<dyn TableProvider>,
265 ) -> DFResult<()> {
266 let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?;
267 let catalog_provider = self
268 .ctx
269 .catalog(&catalog)
270 .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?;
271
272 let paimon_provider = catalog_provider
273 .downcast_ref::<crate::catalog::PaimonCatalogProvider>()
274 .ok_or_else(|| {
275 DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog"))
276 })?;
277
278 paimon_provider.register_temp_table(&database, &table_name, table)
279 }
280
281 pub fn deregister_temp_table(
285 &self,
286 name: impl Into<TableReference>,
287 ) -> DFResult<Option<Arc<dyn TableProvider>>> {
288 let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?;
289 let catalog_provider = self
290 .ctx
291 .catalog(&catalog)
292 .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?;
293
294 let paimon_provider = catalog_provider
295 .downcast_ref::<crate::catalog::PaimonCatalogProvider>()
296 .ok_or_else(|| {
297 DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog"))
298 })?;
299
300 paimon_provider.deregister_temp_table(&database, &table_name)
301 }
302
303 pub fn temp_table_exist(&self, name: impl Into<TableReference>) -> DFResult<bool> {
307 let (catalog, database, table_name) = self.resolve_temp_table_name(name.into())?;
308 let catalog_provider = self
309 .ctx
310 .catalog(&catalog)
311 .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?;
312
313 let paimon_provider = catalog_provider
314 .downcast_ref::<crate::catalog::PaimonCatalogProvider>()
315 .ok_or_else(|| {
316 DataFusionError::Plan(format!("Catalog '{catalog}' is not a Paimon catalog"))
317 })?;
318
319 Ok(paimon_provider.temp_table_exist(&database, &table_name))
320 }
321
322 fn resolve_temp_table_name(&self, name: TableReference) -> DFResult<(String, String, String)> {
324 match name {
325 TableReference::Bare { table } => {
326 let catalog = self.current_catalog_name();
327 let database = self
328 .ctx
329 .state()
330 .config_options()
331 .catalog
332 .default_schema
333 .clone();
334 Ok((catalog, database, table.to_string()))
335 }
336 TableReference::Partial { schema, table } => {
337 let catalog = self.current_catalog_name();
338 Ok((catalog, schema.to_string(), table.to_string()))
339 }
340 TableReference::Full {
341 catalog,
342 schema,
343 table,
344 } => Ok((catalog.to_string(), schema.to_string(), table.to_string())),
345 }
346 }
347
348 #[cfg(test)]
349 pub(crate) fn dynamic_options(&self) -> &DynamicOptions {
350 &self.dynamic_options
351 }
352
353 pub async fn sql(&self, sql: &str) -> DFResult<DataFrame> {
356 let is_create_table = looks_like_create_table(sql);
357 let (rewritten_sql, partition_keys) = if is_create_table {
358 extract_partition_by(sql)?
359 } else {
360 (sql.to_string(), vec![])
361 };
362 if contains_time_travel_keyword(&rewritten_sql) {
363 return self.handle_time_travel_query(&rewritten_sql).await;
365 }
366
367 let statements = parse_sql_statements(&rewritten_sql)?;
368
369 if statements.len() != 1 {
370 return Err(DataFusionError::Plan(
371 "Expected exactly one SQL statement".to_string(),
372 ));
373 }
374
375 match &statements[0] {
376 Statement::CreateTable(create_table) => {
377 if create_table.temporary {
378 self.handle_create_temp_table(create_table).await
379 } else {
380 let (catalog, _catalog_name, _) =
381 self.resolve_catalog_and_table(&create_table.name)?;
382 self.handle_create_table(&catalog, create_table, partition_keys)
383 .await
384 }
385 }
386 Statement::ShowCreate {
387 obj_type: ShowCreateObject::Table,
388 obj_name,
389 } => self.handle_show_create_table(sql, obj_name).await,
390 Statement::AlterTable(alter_table) => {
391 let (catalog, _catalog_name, _) =
392 self.resolve_catalog_and_table(&alter_table.name)?;
393 self.handle_alter_table(
394 &catalog,
395 &alter_table.name,
396 &alter_table.operations,
397 alter_table.if_exists,
398 )
399 .await
400 }
401 Statement::Merge(merge) => self.handle_merge_into(merge).await,
402 Statement::Update(update) => self.handle_update(update).await,
403 Statement::Delete(delete) => self.handle_delete(delete).await,
404 Statement::Insert(insert)
405 if insert.overwrite
406 && insert.partitioned.as_ref().is_some_and(|p| !p.is_empty()) =>
407 {
408 self.handle_insert_overwrite_partition(insert).await
409 }
410 Statement::Set(Set::SingleAssignment {
411 variable, values, ..
412 }) => {
413 let key = variable.to_string();
414 let key = key.trim_matches('\'').trim_matches('"');
415 if let Some(paimon_key) = key.strip_prefix("paimon.") {
416 let value = values
417 .first()
418 .ok_or_else(|| DataFusionError::Plan("SET requires a value".to_string()))?
419 .to_string();
420 let value = value
421 .strip_prefix('\'')
422 .and_then(|s| s.strip_suffix('\''))
423 .unwrap_or(&value)
424 .to_string();
425 self.dynamic_options
426 .write()
427 .unwrap()
428 .insert(paimon_key.to_string(), value);
429 return ok_result(&self.ctx);
430 }
431 self.ctx.sql(sql).await
432 }
433 Statement::Reset(ResetStatement {
434 reset: Reset::ConfigurationParameter(name),
435 }) => {
436 let key = name.to_string();
437 let key = key.trim_matches('\'').trim_matches('"');
438 if let Some(paimon_key) = key.strip_prefix("paimon.") {
439 self.dynamic_options.write().unwrap().remove(paimon_key);
440 return ok_result(&self.ctx);
441 }
442 self.ctx.sql(sql).await
443 }
444 Statement::Truncate(truncate) => self.handle_truncate_table(truncate).await,
445 Statement::CreateView(create_view) => {
446 if create_view.temporary {
447 self.handle_create_view(create_view).await
449 } else {
450 let view_name = create_view.name.to_string();
452 let table_ref: TableReference = view_name.as_str().into();
453 if self.is_paimon_catalog_ref(&table_ref) {
454 self.handle_create_view(create_view).await
455 } else {
456 self.ctx.sql(sql).await
457 }
458 }
459 }
460 Statement::CreateFunction(create_function) => {
461 if self.is_paimon_function_name(&create_function.name) {
462 self.handle_create_function(create_function).await
463 } else {
464 self.ctx.sql(sql).await
465 }
466 }
467 Statement::Drop {
468 object_type,
469 if_exists,
470 names,
471 cascade,
472 restrict,
473 purge,
474 temporary,
475 table,
476 } if matches!(*object_type, ObjectType::Table | ObjectType::View) => {
477 if *temporary {
478 self.handle_drop_temp_table(names, *if_exists)
479 } else if *object_type == ObjectType::Table {
480 let table_ref: TableReference = names[0].to_string().as_str().into();
482 if self.is_paimon_catalog_ref(&table_ref) {
483 let (catalog, _catalog_name, _) =
484 self.resolve_catalog_and_table(&names[0])?;
485 self.handle_drop_table(&catalog, names, *if_exists).await
486 } else {
487 self.ctx.sql(sql).await
488 }
489 } else {
490 let targets_paimon_catalog = names.iter().any(|name| {
491 let table_ref: TableReference = name.to_string().as_str().into();
492 self.is_paimon_catalog_ref(&table_ref)
493 });
494 if !targets_paimon_catalog {
495 return self.ctx.sql(sql).await;
496 }
497 let [name] = names.as_slice() else {
498 return Err(DataFusionError::Plan(
499 "Persistent DROP VIEW does not support multiple views".to_string(),
500 ));
501 };
502 if *cascade {
503 return Err(DataFusionError::Plan(
504 "DROP VIEW CASCADE is not supported".to_string(),
505 ));
506 }
507 if *restrict {
508 return Err(DataFusionError::Plan(
509 "DROP VIEW RESTRICT is not supported".to_string(),
510 ));
511 }
512 if *purge {
513 return Err(DataFusionError::Plan(
514 "DROP VIEW PURGE is not supported".to_string(),
515 ));
516 }
517 if table.is_some() {
518 return Err(DataFusionError::Plan(
519 "DROP VIEW ON clauses are not supported".to_string(),
520 ));
521 }
522 let (catalog, _catalog_name, identifier) =
523 self.resolve_catalog_and_table(name)?;
524 self.handle_drop_view(&catalog, &identifier, *if_exists)
525 .await
526 }
527 }
528 Statement::Call(func) => {
529 crate::procedures::execute_call(
530 &self.ctx,
531 &self.catalogs,
532 &self.current_catalog_name(),
533 func,
534 )
535 .await
536 }
537 Statement::Query(_) | Statement::Explain { .. } => {
538 let current_catalog = self.current_catalog_name();
539 let current_database = self
540 .ctx
541 .state()
542 .config_options()
543 .catalog
544 .default_schema
545 .clone();
546 let expanded = crate::sql_function::expand_statement(
547 statements[0].clone(),
548 &self.catalogs,
549 ¤t_catalog,
550 ¤t_database,
551 )
552 .await?;
553 self.ctx.sql(&expanded.to_string()).await
554 }
555 _ => self.ctx.sql(sql).await,
556 }
557 }
558
559 async fn handle_time_travel_query(&self, sql: &str) -> DFResult<DataFrame> {
568 use crate::table::PaimonTableProvider;
569 use paimon::spec::{SCAN_TIMESTAMP_MILLIS_OPTION, SCAN_VERSION_OPTION};
570
571 let mut tracker = crate::merge_into::TempTableTracker::new(self);
572
573 let version_clauses = extract_all_version_as_of(sql);
574 let timestamp_clauses = extract_all_timestamp_as_of(sql);
575
576 if version_clauses.is_empty() && timestamp_clauses.is_empty() {
577 return Err(DataFusionError::Plan(
578 "Failed to parse time-travel clause in SQL".to_string(),
579 ));
580 }
581
582 let mut replacements: Vec<((usize, usize), String)> = Vec::new();
584
585 for info in &version_clauses {
587 let table_ref: datafusion::common::TableReference = info.table_name.as_str().into();
588 let (catalog, _catalog_name, identifier) =
589 self.resolve_table_name_from_ref(&table_ref)?;
590
591 let (paimon_table, base_identifier, system_name) =
592 load_table_for_read(&catalog, &identifier).await?;
593
594 let mut options = self.dynamic_options.read().unwrap().clone();
596 options.insert(SCAN_VERSION_OPTION.to_string(), info.version.clone());
597
598 let table_with_options = paimon_table
599 .copy_with_time_travel(options)
600 .await
601 .map_err(|e| DataFusionError::External(Box::new(e)))?;
602 let provider: Arc<dyn TableProvider> = if let Some(system_name) = system_name {
603 crate::system_tables::provider_for_table(
604 Arc::clone(&catalog),
605 base_identifier,
606 table_with_options,
607 &system_name,
608 )?
609 .ok_or_else(|| {
610 DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}"))
611 })?
612 } else {
613 Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry(
614 table_with_options,
615 self.blob_reader_registry.clone(),
616 )?)
617 };
618
619 let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple());
620 self.register_temp_table(uuid_name.as_str(), provider)?;
621 tracker.register(&uuid_name);
622 replacements.push((info.clause_range, uuid_name));
623 }
624
625 for info in ×tamp_clauses {
627 let table_ref: datafusion::common::TableReference = info.table_name.as_str().into();
628 let (catalog, _catalog_name, identifier) =
629 self.resolve_table_name_from_ref(&table_ref)?;
630
631 let (paimon_table, base_identifier, system_name) =
632 load_table_for_read(&catalog, &identifier).await?;
633
634 let millis = Self::parse_timestamp_to_millis(&info.timestamp)?;
635
636 let mut options = self.dynamic_options.read().unwrap().clone();
638 options.insert(SCAN_TIMESTAMP_MILLIS_OPTION.to_string(), millis.to_string());
639
640 let table_with_options = paimon_table
641 .copy_with_time_travel(options)
642 .await
643 .map_err(|e| DataFusionError::External(Box::new(e)))?;
644 let provider: Arc<dyn TableProvider> = if let Some(system_name) = system_name {
645 crate::system_tables::provider_for_table(
646 Arc::clone(&catalog),
647 base_identifier,
648 table_with_options,
649 &system_name,
650 )?
651 .ok_or_else(|| {
652 DataFusionError::Plan(format!("Unknown Paimon system table: {system_name}"))
653 })?
654 } else {
655 Arc::new(PaimonTableProvider::try_new_with_blob_reader_registry(
656 table_with_options,
657 self.blob_reader_registry.clone(),
658 )?)
659 };
660
661 let uuid_name = format!("__paimon_tt_{}", uuid::Uuid::new_v4().as_simple());
662 self.register_temp_table(uuid_name.as_str(), provider)?;
663 tracker.register(&uuid_name);
664 replacements.push((info.clause_range, uuid_name));
665 }
666
667 replacements.sort_by_key(|r| std::cmp::Reverse(r.0 .0));
670
671 let mut rewritten_sql = sql.to_string();
673 for ((start, end), uuid_name) in &replacements {
674 rewritten_sql = format!(
675 "{}{}{}",
676 &rewritten_sql[..*start],
677 uuid_name,
678 &rewritten_sql[*end..]
679 );
680 }
681
682 let current_catalog = self.current_catalog_name();
684 let current_database = self
685 .ctx
686 .state()
687 .config_options()
688 .catalog
689 .default_schema
690 .clone();
691 let expanded = crate::sql_function::expand_sql(
692 &rewritten_sql,
693 &self.catalogs,
694 ¤t_catalog,
695 ¤t_database,
696 )
697 .await?;
698 self.ctx.sql(&expanded).await
699 }
700
701 fn parse_timestamp_to_millis(ts: &str) -> DFResult<i64> {
703 use chrono::{Local, NaiveDateTime, TimeZone};
704
705 let naive = NaiveDateTime::parse_from_str(ts, "%Y-%m-%d %H:%M:%S").map_err(|e| {
706 DataFusionError::Plan(format!(
707 "Cannot parse time travel timestamp '{ts}': {e}. Expected format: YYYY-MM-DD HH:MM:SS"
708 ))
709 })?;
710 let local = Local.from_local_datetime(&naive).single().ok_or_else(|| {
711 DataFusionError::Plan(format!("Ambiguous or invalid local time: '{ts}'"))
712 })?;
713 Ok(local.timestamp_millis())
714 }
715
716 fn resolve_table_name_from_ref(
718 &self,
719 table_ref: &datafusion::common::TableReference,
720 ) -> DFResult<(Arc<dyn Catalog>, String, Identifier)> {
721 match table_ref {
722 datafusion::common::TableReference::Full {
723 catalog,
724 schema,
725 table,
726 } => {
727 let catalog_arc = self
728 .catalogs
729 .get(catalog.as_ref())
730 .ok_or_else(|| DataFusionError::Plan(format!("Unknown catalog '{catalog}'")))?;
731 Ok((
732 catalog_arc.clone(),
733 catalog.to_string(),
734 Identifier::new(schema.as_ref(), table.as_ref()),
735 ))
736 }
737 datafusion::common::TableReference::Partial { schema, table } => {
738 let catalog = self.current_catalog()?;
739 let catalog_name = self.current_catalog_name();
740 Ok((
741 catalog,
742 catalog_name,
743 Identifier::new(schema.as_ref(), table.as_ref()),
744 ))
745 }
746 datafusion::common::TableReference::Bare { table } => {
747 let catalog = self.current_catalog()?;
748 let catalog_name = self.current_catalog_name();
749 let default_schema = self
750 .ctx
751 .state()
752 .config_options()
753 .catalog
754 .default_schema
755 .clone();
756 Ok((
757 catalog,
758 catalog_name,
759 Identifier::new(default_schema, table.as_ref()),
760 ))
761 }
762 }
763 }
764
765 async fn handle_create_table(
766 &self,
767 catalog: &Arc<dyn Catalog>,
768 ct: &CreateTable,
769 partition_keys: Vec<String>,
770 ) -> DFResult<DataFrame> {
771 if ct.external {
772 return Err(DataFusionError::Plan(
773 "CREATE EXTERNAL TABLE is not supported. Use CREATE TABLE instead.".to_string(),
774 ));
775 }
776 if ct.location.is_some() {
777 return Err(DataFusionError::Plan(
778 "LOCATION is not supported for Paimon tables. Table path is determined by the catalog warehouse.".to_string(),
779 ));
780 }
781 if ct.query.is_some() {
782 return Err(DataFusionError::Plan(
783 "CREATE TABLE AS SELECT is not yet supported for Paimon tables.".to_string(),
784 ));
785 }
786
787 let identifier = self.resolve_table_name(&ct.name)?;
788
789 let mut builder = paimon::spec::Schema::builder();
790 let table_options = extract_options(&ct.table_options)?;
791
792 for col in &ct.columns {
794 let paimon_type = column_def_to_paimon_type(col)?;
795 let comment = column_def_comment(col);
796 builder = builder.column_with_description(col.name.value.clone(), paimon_type, comment);
797 }
798
799 for constraint in &ct.constraints {
801 if let datafusion::sql::sqlparser::ast::TableConstraint::PrimaryKey(pk) = constraint {
802 let pk_cols: Vec<String> = pk
803 .columns
804 .iter()
805 .map(|c| primary_key_column_name(&c.column.expr))
806 .collect();
807 builder = builder.primary_key(pk_cols);
808 }
809 }
810
811 if !partition_keys.is_empty() {
813 let col_names: Vec<&str> = ct.columns.iter().map(|c| c.name.value.as_str()).collect();
814 for pk in &partition_keys {
815 if !col_names.contains(&pk.as_str()) {
816 return Err(DataFusionError::Plan(format!(
817 "PARTITIONED BY column '{pk}' is not defined in the table"
818 )));
819 }
820 }
821 builder = builder.partition_keys(partition_keys);
822 }
823
824 for (k, v) in table_options {
826 builder = builder.option(k, v);
827 }
828
829 let schema = builder.build().map_err(to_datafusion_error)?;
830
831 catalog
832 .create_table(&identifier, schema, ct.if_not_exists)
833 .await
834 .map_err(to_datafusion_error)?;
835
836 ok_result(&self.ctx)
837 }
838
839 async fn handle_create_temp_table(&self, ct: &CreateTable) -> DFResult<DataFrame> {
840 let table_ref: TableReference = ct.name.to_string().as_str().into();
841
842 if ct.if_not_exists && self.temp_table_exist(table_ref.clone())? {
843 return ok_result(&self.ctx);
844 }
845
846 let declared_schema = if !ct.columns.is_empty() {
848 let fields: Vec<Field> = ct
849 .columns
850 .iter()
851 .map(|col| {
852 let paimon_type =
853 sql_data_type_to_paimon_type(&col.data_type, column_def_nullable(col))?;
854 let arrow_type = paimon::arrow::paimon_type_to_arrow(&paimon_type)
855 .map_err(to_datafusion_error)?;
856 Ok(Field::new(
857 &col.name.value,
858 arrow_type,
859 column_def_nullable(col),
860 ))
861 })
862 .collect::<DFResult<Vec<_>>>()?;
863 Some(Arc::new(Schema::new(fields)))
864 } else {
865 None
866 };
867
868 if let Some(query) = &ct.query {
869 let query_sql = query.to_string();
871 let df = self.ctx.sql(&query_sql).await?;
872 let schema = df.schema().inner().clone();
873 let batches = df.collect().await?;
874
875 let batches = if ct.columns.is_empty() {
877 batches
878 } else {
879 let target_fields: Vec<(String, ArrowDataType)> = ct
880 .columns
881 .iter()
882 .map(|col| {
883 let paimon_type =
884 sql_data_type_to_paimon_type(&col.data_type, column_def_nullable(col))?;
885 let arrow_type = paimon::arrow::paimon_type_to_arrow(&paimon_type)
886 .map_err(to_datafusion_error)?;
887 Ok((col.name.value.clone(), arrow_type))
888 })
889 .collect::<DFResult<Vec<_>>>()?;
890
891 let select_col_count = schema.fields().len();
892 let declared_col_count = target_fields.len();
893 if select_col_count < declared_col_count {
894 return Err(DataFusionError::Plan(format!(
895 "CREATE TEMPORARY TABLE AS SELECT: declared {declared_col_count} column(s) \
896 but SELECT query returns only {select_col_count} column(s)"
897 )));
898 }
899
900 batches
901 .into_iter()
902 .map(|batch| {
903 let columns = batch
904 .columns()
905 .iter()
906 .enumerate()
907 .map(|(i, col)| {
908 if i < target_fields.len() {
909 let target_dt = &target_fields[i].1;
910 if *col.data_type() != *target_dt {
911 cast(col, target_dt)
912 .map_err(|e| DataFusionError::External(e.into()))
913 } else {
914 Ok(col.clone())
915 }
916 } else {
917 Ok(col.clone())
918 }
919 })
920 .collect::<DFResult<Vec<_>>>()?;
921 let new_fields = target_fields
922 .iter()
923 .zip(schema.fields().iter())
924 .map(|((name, dt), _)| Field::new(name, dt.clone(), true))
925 .chain(
926 schema
927 .fields()
928 .iter()
929 .skip(target_fields.len())
930 .map(|f| f.as_ref().clone()),
931 )
932 .collect::<Vec<_>>();
933 let new_schema = Schema::new(new_fields);
934 RecordBatch::try_new(Arc::new(new_schema), columns)
935 .map_err(|e| DataFusionError::External(e.into()))
936 })
937 .collect::<DFResult<Vec<_>>>()?
938 };
939
940 let schema = batches.first().map(|b| b.schema()).unwrap_or(schema);
941 let mem_table = MemTable::try_new(schema, vec![batches])?;
942 self.register_temp_table(table_ref, Arc::new(mem_table))?;
943 } else if let Some(schema) = declared_schema {
944 let mem_table = MemTable::try_new(schema, vec![vec![]])?;
946 self.register_temp_table(table_ref, Arc::new(mem_table))?;
947 } else {
948 return Err(DataFusionError::Plan(
949 "CREATE TEMPORARY TABLE requires column definitions or AS SELECT".to_string(),
950 ));
951 }
952
953 ok_result(&self.ctx)
954 }
955
956 fn handle_drop_temp_table(&self, names: &[ObjectName], if_exists: bool) -> DFResult<DataFrame> {
957 for name in names {
958 let table_ref: TableReference = name.to_string().as_str().into();
959 if if_exists && !self.temp_table_exist(table_ref.clone())? {
960 continue;
961 }
962 self.deregister_temp_table(table_ref)?;
963 }
964 ok_result(&self.ctx)
965 }
966
967 async fn handle_drop_table(
968 &self,
969 catalog: &Arc<dyn Catalog>,
970 names: &[ObjectName],
971 if_exists: bool,
972 ) -> DFResult<DataFrame> {
973 for name in names {
974 let identifier = self.resolve_table_name(name)?;
975 catalog
976 .drop_table(&identifier, if_exists)
977 .await
978 .map_err(|e| DataFusionError::External(Box::new(e)))?;
979 }
980 ok_result(&self.ctx)
981 }
982
983 async fn handle_drop_view(
984 &self,
985 catalog: &Arc<dyn Catalog>,
986 identifier: &Identifier,
987 if_exists: bool,
988 ) -> DFResult<DataFrame> {
989 catalog
990 .drop_view(identifier, if_exists)
991 .await
992 .map_err(to_datafusion_error)?;
993 ok_result(&self.ctx)
994 }
995
996 async fn handle_show_create_table(&self, sql: &str, name: &ObjectName) -> DFResult<DataFrame> {
997 let (catalog, catalog_name, identifier) = self.resolve_catalog_and_table(name)?;
998 let table = match catalog.get_table(&identifier).await {
999 Ok(table) => table,
1000 Err(paimon::Error::TableNotExist { .. }) => return self.ctx.sql(sql).await,
1001 Err(e) => return Err(to_datafusion_error(e)),
1002 };
1003 let definition = crate::table::build_table_definition(&table)?;
1004
1005 let schema = Arc::new(Schema::new(vec![
1006 Field::new("table_catalog", ArrowDataType::Utf8, false),
1007 Field::new("table_schema", ArrowDataType::Utf8, false),
1008 Field::new("table_name", ArrowDataType::Utf8, false),
1009 Field::new("definition", ArrowDataType::Utf8, false),
1010 ]));
1011 let batch = RecordBatch::try_new(
1012 schema,
1013 vec![
1014 Arc::new(StringArray::from(vec![catalog_name])),
1015 Arc::new(StringArray::from(vec![identifier.database().to_string()])),
1016 Arc::new(StringArray::from(vec![identifier.object().to_string()])),
1017 Arc::new(StringArray::from(vec![definition])),
1018 ],
1019 )?;
1020 self.ctx.read_batch(batch)
1021 }
1022
1023 async fn handle_alter_table(
1024 &self,
1025 catalog: &Arc<dyn Catalog>,
1026 name: &ObjectName,
1027 operations: &[AlterTableOperation],
1028 if_exists: bool,
1029 ) -> DFResult<DataFrame> {
1030 Self::ensure_main_branch_write_target(name, "ALTER TABLE")?;
1031 let identifier = self.resolve_table_name(name)?;
1032
1033 let mut changes = Vec::new();
1034 let mut rename_to: Option<Identifier> = None;
1035
1036 for op in operations {
1037 match op {
1038 AlterTableOperation::AddColumn { column_def, .. } => {
1039 changes.push(column_def_to_add_column(column_def)?);
1040 }
1041 AlterTableOperation::DropColumn { column_names, .. } => {
1042 for col in column_names {
1043 changes.push(SchemaChange::drop_column(col.value.clone()));
1044 }
1045 }
1046 AlterTableOperation::RenameColumn {
1047 old_column_name,
1048 new_column_name,
1049 } => {
1050 changes.push(SchemaChange::rename_column(
1051 old_column_name.value.clone(),
1052 new_column_name.value.clone(),
1053 ));
1054 }
1055 AlterTableOperation::RenameTable { table_name } => {
1056 let new_name = match table_name {
1057 RenameTableNameKind::To(name) | RenameTableNameKind::As(name) => {
1058 object_name_to_string(name)
1059 }
1060 };
1061 rename_to = Some(Identifier::new(identifier.database().to_string(), new_name));
1062 }
1063 AlterTableOperation::SetTblProperties { table_properties } => {
1064 for opt in table_properties {
1065 if let SqlOption::KeyValue { key, value } = opt {
1066 let v = value.to_string();
1067 let v = v
1068 .strip_prefix('\'')
1069 .and_then(|s| s.strip_suffix('\''))
1070 .unwrap_or(&v)
1071 .to_string();
1072 changes.push(SchemaChange::set_option(key.value.clone(), v));
1073 }
1074 }
1075 }
1076 AlterTableOperation::DropPartitions {
1077 partitions,
1078 if_exists: partition_if_exists,
1079 } => {
1080 return self
1081 .handle_drop_partitions(
1082 catalog,
1083 &identifier,
1084 partitions,
1085 if_exists || *partition_if_exists,
1086 )
1087 .await;
1088 }
1089 other => {
1090 return Err(DataFusionError::Plan(format!(
1091 "Unsupported ALTER TABLE operation: {other}"
1092 )));
1093 }
1094 }
1095 }
1096
1097 if let Some(new_identifier) = rename_to {
1098 catalog
1099 .rename_table(&identifier, &new_identifier, if_exists)
1100 .await
1101 .map_err(to_datafusion_error)?;
1102 }
1103
1104 if !changes.is_empty() {
1105 catalog
1106 .alter_table(&identifier, changes, if_exists)
1107 .await
1108 .map_err(to_datafusion_error)?;
1109 }
1110
1111 ok_result(&self.ctx)
1112 }
1113
1114 fn ensure_no_time_travel_for_write(&self, operation: &str) -> DFResult<()> {
1124 use paimon::spec::{
1125 SCAN_SNAPSHOT_ID_OPTION, SCAN_TAG_NAME_OPTION, SCAN_TIMESTAMP_MILLIS_OPTION,
1126 SCAN_VERSION_OPTION,
1127 };
1128
1129 let options = self.dynamic_options.read().unwrap();
1130 for key in [
1131 SCAN_VERSION_OPTION,
1132 SCAN_TIMESTAMP_MILLIS_OPTION,
1133 SCAN_SNAPSHOT_ID_OPTION,
1134 SCAN_TAG_NAME_OPTION,
1135 ] {
1136 if options.contains_key(key) {
1137 return Err(DataFusionError::Plan(format!(
1138 "Cannot execute {operation} while time-travel option '{key}' is set; \
1139 RESET 'paimon.{key}' first"
1140 )));
1141 }
1142 }
1143 Ok(())
1144 }
1145
1146 async fn handle_merge_into(&self, merge: &Merge) -> DFResult<DataFrame> {
1147 self.ensure_no_time_travel_for_write("MERGE INTO")?;
1148 let table_name = match &merge.table {
1149 TableFactor::Table { name, .. } => name.clone(),
1150 other => {
1151 return Err(DataFusionError::Plan(format!(
1152 "Unsupported target table in MERGE INTO: {other}"
1153 )))
1154 }
1155 };
1156 Self::ensure_main_branch_write_target(&table_name, "MERGE INTO")?;
1157 let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
1158
1159 let table = catalog
1160 .get_table(&identifier)
1161 .await
1162 .map_err(to_datafusion_error)?;
1163
1164 crate::merge_into::execute_merge_into(self, merge, table).await
1165 }
1166
1167 async fn handle_update(&self, update: &Update) -> DFResult<DataFrame> {
1168 self.ensure_no_time_travel_for_write("UPDATE")?;
1169 let table_name = match &update.table.relation {
1170 TableFactor::Table { name, .. } => name.clone(),
1171 other => {
1172 return Err(DataFusionError::Plan(format!(
1173 "Unsupported target table in UPDATE: {other}"
1174 )))
1175 }
1176 };
1177 Self::ensure_main_branch_write_target(&table_name, "UPDATE")?;
1178 let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
1179
1180 let table = catalog
1181 .get_table(&identifier)
1182 .await
1183 .map_err(to_datafusion_error)?;
1184
1185 crate::update::execute_update(self, update, table).await
1186 }
1187
1188 async fn handle_delete(&self, delete: &Delete) -> DFResult<DataFrame> {
1189 self.ensure_no_time_travel_for_write("DELETE")?;
1190 let tables = match &delete.from {
1191 FromTable::WithFromKeyword(t) | FromTable::WithoutKeyword(t) => t,
1192 };
1193 let table_factor = tables
1194 .first()
1195 .map(|t| &t.relation)
1196 .ok_or_else(|| DataFusionError::Plan("DELETE requires a target table".to_string()))?;
1197 let table_name = match table_factor {
1198 TableFactor::Table { name, .. } => name.clone(),
1199 other => {
1200 return Err(DataFusionError::Plan(format!(
1201 "Unsupported target table in DELETE: {other}"
1202 )))
1203 }
1204 };
1205 Self::ensure_main_branch_write_target(&table_name, "DELETE")?;
1206 let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
1207
1208 let table = catalog
1209 .get_table(&identifier)
1210 .await
1211 .map_err(to_datafusion_error)?;
1212
1213 let table_ref = table_name.to_string();
1214 crate::delete::execute_delete(self, delete, table, &table_ref).await
1215 }
1216
1217 async fn handle_insert_overwrite_partition(&self, insert: &Insert) -> DFResult<DataFrame> {
1218 self.ensure_no_time_travel_for_write("INSERT OVERWRITE")?;
1219 let table_name = match &insert.table {
1220 TableObject::TableName(name) => name.clone(),
1221 other => {
1222 return Err(DataFusionError::Plan(format!(
1223 "Unsupported target table in INSERT OVERWRITE: {other}"
1224 )))
1225 }
1226 };
1227 Self::ensure_main_branch_write_target(&table_name, "INSERT OVERWRITE")?;
1228 let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?;
1229 let table = catalog
1230 .get_table(&identifier)
1231 .await
1232 .map_err(to_datafusion_error)?;
1233
1234 let partition_exprs = insert.partitioned.as_ref().ok_or_else(|| {
1235 DataFusionError::Plan("INSERT OVERWRITE PARTITION requires a PARTITION clause".into())
1236 })?;
1237 let partition_fields = table.schema().partition_fields();
1238 let static_partitions =
1239 parse_static_partitions(partition_exprs, &partition_fields, table.schema().fields())?;
1240
1241 let source = insert.source.as_ref().ok_or_else(|| {
1242 DataFusionError::Plan("INSERT OVERWRITE requires a source query".into())
1243 })?;
1244 let df = self.ctx.sql(&source.to_string()).await?;
1245
1246 let all_fields = table.schema().fields();
1247 let non_static_fields: Vec<&PaimonDataField> = all_fields
1248 .iter()
1249 .filter(|f| !static_partitions.contains_key(f.name()))
1250 .collect();
1251 let expected_source_cols = non_static_fields.len();
1252
1253 let target_columns: Option<Vec<String>> = if !insert.columns.is_empty() {
1256 Some(
1257 insert
1258 .columns
1259 .iter()
1260 .map(object_name_to_single_identifier)
1261 .collect::<DFResult<_>>()?,
1262 )
1263 } else if !insert.after_columns.is_empty() {
1264 Some(
1265 insert
1266 .after_columns
1267 .iter()
1268 .map(|ident| ident.value.clone())
1269 .collect(),
1270 )
1271 } else {
1272 None
1273 };
1274 let column_reorder: Option<Vec<usize>> = if let Some(cols) = target_columns.as_ref() {
1275 if cols.len() != expected_source_cols {
1276 return Err(DataFusionError::Plan(format!(
1277 "Column list has {} columns, but expected {} non-partition columns",
1278 cols.len(),
1279 expected_source_cols
1280 )));
1281 }
1282 let col_names: Vec<&str> = cols.iter().map(String::as_str).collect();
1283 let mut reorder = Vec::with_capacity(expected_source_cols);
1284 for field in &non_static_fields {
1285 let pos = col_names
1286 .iter()
1287 .position(|c| c == &field.name())
1288 .ok_or_else(|| {
1289 DataFusionError::Plan(format!(
1290 "Column '{}' not found in target column list",
1291 field.name()
1292 ))
1293 })?;
1294 reorder.push(pos);
1295 }
1296 Some(reorder)
1297 } else {
1298 None
1299 };
1300
1301 let source_col_count = df.schema().fields().len();
1303 if source_col_count != expected_source_cols {
1304 return Err(DataFusionError::Plan(format!(
1305 "Source query has {} columns, but expected {} non-partition columns",
1306 source_col_count, expected_source_cols
1307 )));
1308 }
1309
1310 let mut stream = df.execute_stream().await?;
1311
1312 let wb = table.new_write_builder().with_overwrite();
1313 let mut tw = wb.new_write().map_err(to_datafusion_error)?;
1314 let mut row_count = 0u64;
1315
1316 while let Some(batch_result) = stream.next().await {
1317 let batch = batch_result?;
1318 if batch.num_rows() == 0 {
1319 continue;
1320 }
1321 let batch = if let Some(ref reorder) = column_reorder {
1322 let reordered_cols: Vec<ArrayRef> =
1323 reorder.iter().map(|&i| batch.column(i).clone()).collect();
1324 let reordered_fields: Vec<Field> = reorder
1325 .iter()
1326 .map(|&i| batch.schema().field(i).clone())
1327 .collect();
1328 let reordered_schema = Arc::new(Schema::new(reordered_fields));
1329 RecordBatch::try_new(reordered_schema, reordered_cols)
1330 .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?
1331 } else {
1332 batch
1333 };
1334 let augmented = append_partition_columns(
1335 &batch,
1336 &static_partitions,
1337 expected_source_cols,
1338 all_fields,
1339 )?;
1340 row_count += augmented.num_rows() as u64;
1341 tw.write_arrow_batch(&augmented)
1342 .await
1343 .map_err(to_datafusion_error)?;
1344 }
1345
1346 let messages = tw.prepare_commit().await.map_err(to_datafusion_error)?;
1347 let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
1348
1349 let overwrite_partitions = if static_partitions.is_empty() {
1350 None
1351 } else {
1352 Some(static_partitions)
1353 };
1354 commit
1355 .overwrite(messages, overwrite_partitions)
1356 .await
1357 .map_err(to_datafusion_error)?;
1358
1359 crate::merge_into::ok_result(&self.ctx, row_count)
1360 }
1361
1362 async fn handle_truncate_table(&self, truncate: &Truncate) -> DFResult<DataFrame> {
1363 self.ensure_no_time_travel_for_write("TRUNCATE TABLE")?;
1364 if truncate.table_names.len() > 1 {
1365 return Err(DataFusionError::Plan(
1366 "TRUNCATE TABLE does not support multiple tables".to_string(),
1367 ));
1368 }
1369 let target = truncate.table_names.first().ok_or_else(|| {
1370 DataFusionError::Plan("TRUNCATE TABLE requires a table name".to_string())
1371 })?;
1372 Self::ensure_main_branch_write_target(&target.name, "TRUNCATE TABLE")?;
1373 let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&target.name)?;
1374 let table = match catalog.get_table(&identifier).await {
1375 Ok(t) => t,
1376 Err(e) if truncate.if_exists && is_table_not_exist(&e) => {
1377 return ok_result(&self.ctx);
1378 }
1379 Err(e) => return Err(to_datafusion_error(e)),
1380 };
1381
1382 let wb = table.new_write_builder();
1383 let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
1384
1385 if let Some(partitions) = &truncate.partitions {
1386 if partitions.is_empty() {
1387 return Err(DataFusionError::Plan(
1388 "PARTITION clause requires at least one column = value".to_string(),
1389 ));
1390 }
1391 let partition_values = parse_partition_values(
1392 partitions,
1393 table.schema().fields(),
1394 table.schema().partition_keys(),
1395 )?;
1396 commit
1397 .truncate_partitions(partition_values)
1398 .await
1399 .map_err(to_datafusion_error)?;
1400 return ok_result(&self.ctx);
1401 }
1402
1403 commit.truncate_table().await.map_err(to_datafusion_error)?;
1404 ok_result(&self.ctx)
1405 }
1406
1407 async fn handle_create_view(&self, create_view: &CreateView) -> DFResult<DataFrame> {
1408 if create_view.materialized {
1409 return Err(DataFusionError::Plan(
1410 "CREATE MATERIALIZED VIEW is not supported".to_string(),
1411 ));
1412 }
1413
1414 let query_sql = create_view.query.to_string();
1415
1416 if create_view.temporary {
1417 let view_name = create_view.name.to_string();
1418 let table_ref: TableReference = view_name.as_str().into();
1419 let (catalog, database, name) = self.resolve_temp_table_name(table_ref)?;
1420 let df = self.ctx.sql(&query_sql).await?;
1421 let logical_plan = df.logical_plan().clone();
1422 if create_view.if_not_exists
1423 && self.temp_table_exist(format!("{catalog}.{database}.{name}"))?
1424 {
1425 return ok_result(&self.ctx);
1426 }
1427 let view_table = datafusion::datasource::ViewTable::new(logical_plan, Some(query_sql));
1429 self.register_temp_table(format!("{catalog}.{database}.{name}"), Arc::new(view_table))?;
1430 ok_result(&self.ctx)
1431 } else {
1432 validate_persistent_create_view(create_view)?;
1433 let (catalog, catalog_name, identifier) =
1434 self.resolve_catalog_and_table(&create_view.name)?;
1435 let mut state = self.ctx.state();
1436 state.config_mut().options_mut().catalog.default_catalog = catalog_name.clone();
1437 state.config_mut().options_mut().catalog.default_schema =
1438 identifier.database().to_string();
1439 let expanded_query = crate::sql_function::expand_sql(
1440 &query_sql,
1441 &self.catalogs,
1442 &catalog_name,
1443 identifier.database(),
1444 )
1445 .await?;
1446 let logical_plan = state.create_logical_plan(&expanded_query).await?;
1447 let mut arrow_fields = logical_plan
1448 .schema()
1449 .as_arrow()
1450 .fields()
1451 .iter()
1452 .map(|field| field.as_ref().clone())
1453 .collect::<Vec<_>>();
1454 if !create_view.columns.is_empty() && create_view.columns.len() != arrow_fields.len() {
1455 return Err(DataFusionError::Plan(format!(
1456 "view column list has {} columns but query produces {} columns",
1457 create_view.columns.len(),
1458 arrow_fields.len()
1459 )));
1460 }
1461 let column_names = create_view
1462 .columns
1463 .iter()
1464 .map(|column| IdentNormalizer::default().normalize(column.name.clone()))
1465 .collect::<Vec<_>>();
1466 let mut unique_names = HashSet::with_capacity(column_names.len());
1467 for name in &column_names {
1468 if !unique_names.insert(name.clone()) {
1469 return Err(DataFusionError::Plan(format!(
1470 "duplicate view column name '{name}'"
1471 )));
1472 }
1473 }
1474 for (field, name) in arrow_fields.iter_mut().zip(column_names) {
1475 *field = field.clone().with_name(name);
1476 }
1477 let fields = paimon::arrow::arrow_fields_to_paimon(&arrow_fields)
1478 .map_err(to_datafusion_error)?;
1479 let schema = paimon::catalog::ViewSchema::new(
1480 fields,
1481 query_sql.clone(),
1482 HashMap::from([("datafusion".to_string(), query_sql)]),
1483 None,
1484 HashMap::new(),
1485 );
1486 catalog
1487 .create_view(&identifier, schema, create_view.if_not_exists)
1488 .await
1489 .map_err(to_datafusion_error)?;
1490 ok_result(&self.ctx)
1491 }
1492 }
1493
1494 async fn handle_create_function(
1495 &self,
1496 create_function: &CreateFunction,
1497 ) -> DFResult<DataFrame> {
1498 validate_persistent_create_function(create_function)?;
1499 if create_function
1500 .language
1501 .as_ref()
1502 .is_some_and(|language| !language.value.eq_ignore_ascii_case("sql"))
1503 {
1504 return Err(DataFusionError::Plan(
1505 "CREATE FUNCTION only supports LANGUAGE SQL".to_string(),
1506 ));
1507 }
1508 if matches!(
1509 create_function.behavior,
1510 Some(FunctionBehavior::Stable | FunctionBehavior::Volatile)
1511 ) {
1512 return Err(DataFusionError::Plan(
1513 "CREATE FUNCTION only supports deterministic SQL functions".to_string(),
1514 ));
1515 }
1516 let FunctionReturnType::DataType(return_type) = create_function
1517 .return_type
1518 .as_ref()
1519 .ok_or_else(|| DataFusionError::Plan("CREATE FUNCTION requires RETURNS".to_string()))?
1520 else {
1521 return Err(DataFusionError::Plan(
1522 "CREATE FUNCTION SETOF return types are not supported".to_string(),
1523 ));
1524 };
1525 let CreateFunctionBody::Return(body) = create_function
1526 .function_body
1527 .as_ref()
1528 .ok_or_else(|| DataFusionError::Plan("CREATE FUNCTION requires RETURN".to_string()))?
1529 else {
1530 return Err(DataFusionError::Plan(
1531 "CREATE FUNCTION only supports a RETURN expression".to_string(),
1532 ));
1533 };
1534
1535 let mut parameter_names = HashSet::new();
1536 let input_params = create_function
1537 .args
1538 .as_deref()
1539 .unwrap_or_default()
1540 .iter()
1541 .enumerate()
1542 .map(|(id, argument)| {
1543 if argument.mode.is_some() || argument.default_expr.is_some() {
1544 return Err(DataFusionError::Plan(
1545 "CREATE FUNCTION argument modes and defaults are not supported".to_string(),
1546 ));
1547 }
1548 let name = argument
1549 .name
1550 .clone()
1551 .map(normalize_create_function_argument_name)
1552 .ok_or_else(|| {
1553 DataFusionError::Plan(
1554 "CREATE FUNCTION arguments must have names".to_string(),
1555 )
1556 })?;
1557 if !parameter_names.insert(name.clone()) {
1558 return Err(DataFusionError::Plan(format!(
1559 "duplicate function argument name '{name}'"
1560 )));
1561 }
1562 Ok(PaimonDataField::new(
1563 id as i32,
1564 name,
1565 sql_data_type_to_paimon_type(&argument.data_type, true)?,
1566 ))
1567 })
1568 .collect::<DFResult<Vec<_>>>()?;
1569 let return_params = vec![PaimonDataField::new(
1570 0,
1571 "result".to_string(),
1572 sql_data_type_to_paimon_type(return_type, true)?,
1573 )];
1574 let (catalog, catalog_name, identifier) =
1575 self.resolve_catalog_and_function(&create_function.name)?;
1576 let function = paimon::catalog::Function::new(
1577 identifier,
1578 Some(input_params),
1579 Some(return_params),
1580 true,
1581 HashMap::from([(
1582 "datafusion".to_string(),
1583 paimon::catalog::FunctionDefinition::Sql {
1584 definition: body.to_string(),
1585 },
1586 )]),
1587 None,
1588 HashMap::new(),
1589 );
1590 self.validate_create_function(&function, &catalog_name)
1591 .await?;
1592 catalog
1593 .create_function(&function, create_function.if_not_exists)
1594 .await
1595 .map_err(to_datafusion_error)?;
1596 ok_result(&self.ctx)
1597 }
1598
1599 async fn validate_create_function(
1600 &self,
1601 function: &paimon::catalog::Function,
1602 catalog_name: &str,
1603 ) -> DFResult<()> {
1604 let arguments = function
1605 .input_params()
1606 .unwrap_or_default()
1607 .iter()
1608 .map(|field| {
1609 let sql_type =
1610 crate::table::data_type_to_sql(field.data_type()).map_err(|error| {
1611 DataFusionError::Plan(format!(
1612 "Invalid CREATE FUNCTION argument type '{:?}': {error}",
1613 field.data_type()
1614 ))
1615 })?;
1616 Ok(format!("CAST(NULL AS {sql_type})"))
1617 })
1618 .collect::<DFResult<Vec<_>>>()?
1619 .join(", ");
1620 let quote = |identifier: &str| format!("\"{}\"", identifier.replace('"', "\"\""));
1621 let validation_sql = format!(
1622 "SELECT {}.{}.{}({arguments})",
1623 quote(catalog_name),
1624 quote(function.identifier().database()),
1625 quote(function.name())
1626 );
1627 let expanded = crate::sql_function::expand_sql_with_candidate(
1628 &validation_sql,
1629 &self.catalogs,
1630 catalog_name,
1631 function.identifier().database(),
1632 function,
1633 )
1634 .await?;
1635 let mut state = self.ctx.state();
1636 state.config_mut().options_mut().catalog.default_catalog = catalog_name.to_string();
1637 state.config_mut().options_mut().catalog.default_schema =
1638 function.identifier().database().to_string();
1639 let logical_plan = state.create_logical_plan(&expanded).await?;
1640 validate_immutable_scalar_plan(&logical_plan)?;
1641 state.create_physical_plan(&logical_plan).await?;
1642 Ok(())
1643 }
1644
1645 async fn handle_drop_partitions(
1646 &self,
1647 catalog: &Arc<dyn Catalog>,
1648 identifier: &Identifier,
1649 partitions: &[SqlExpr],
1650 if_exists: bool,
1651 ) -> DFResult<DataFrame> {
1652 if partitions.is_empty() {
1653 return Err(DataFusionError::Plan(
1654 "DROP PARTITIONS requires at least one partition specification".to_string(),
1655 ));
1656 }
1657 let table = match catalog.get_table(identifier).await {
1658 Ok(t) => t,
1659 Err(e) if if_exists && is_table_not_exist(&e) => {
1660 return ok_result(&self.ctx);
1661 }
1662 Err(e) => return Err(to_datafusion_error(e)),
1663 };
1664
1665 let partition_values = parse_partition_values(
1666 partitions,
1667 table.schema().fields(),
1668 table.schema().partition_keys(),
1669 )?;
1670
1671 let wb = table.new_write_builder();
1672 let commit = wb.try_new_commit().map_err(to_datafusion_error)?;
1673 commit
1674 .truncate_partitions(partition_values)
1675 .await
1676 .map_err(to_datafusion_error)?;
1677
1678 ok_result(&self.ctx)
1679 }
1680
1681 pub(crate) fn current_catalog_name(&self) -> String {
1683 self.ctx
1684 .state()
1685 .config_options()
1686 .catalog
1687 .default_catalog
1688 .clone()
1689 }
1690
1691 fn current_catalog(&self) -> DFResult<Arc<dyn Catalog>> {
1692 let name = self.current_catalog_name();
1693 self.catalogs.get(&name).cloned().ok_or_else(|| {
1694 DataFusionError::Plan(
1695 "No catalog registered. Call register_catalog() first.".to_string(),
1696 )
1697 })
1698 }
1699
1700 fn is_paimon_catalog_ref(&self, table_ref: &TableReference) -> bool {
1702 let catalog_name = match table_ref {
1703 TableReference::Full { catalog, .. } => catalog.to_string(),
1704 TableReference::Partial { .. } | TableReference::Bare { .. } => {
1705 self.current_catalog_name()
1706 }
1707 };
1708 self.catalogs.contains_key(&catalog_name)
1709 }
1710
1711 fn is_paimon_function_name(&self, name: &ObjectName) -> bool {
1712 let Some(parts) = name
1713 .0
1714 .iter()
1715 .map(|part| {
1716 part.as_ident()
1717 .map(|identifier| IdentNormalizer::default().normalize(identifier.clone()))
1718 })
1719 .collect::<Option<Vec<_>>>()
1720 else {
1721 return false;
1722 };
1723 let catalog_name = match parts.as_slice() {
1724 [catalog, _, _] => catalog.clone(),
1725 [_] | [_, _] => self.current_catalog_name(),
1726 _ => return false,
1727 };
1728 self.catalogs.contains_key(&catalog_name)
1729 }
1730
1731 fn resolve_catalog_and_function(
1732 &self,
1733 name: &ObjectName,
1734 ) -> DFResult<(Arc<dyn Catalog>, String, Identifier)> {
1735 let parts = name
1736 .0
1737 .iter()
1738 .map(|part| {
1739 part.as_ident()
1740 .cloned()
1741 .map(|identifier| IdentNormalizer::default().normalize(identifier))
1742 .ok_or_else(|| {
1743 DataFusionError::Plan(format!("Invalid function reference: {name}"))
1744 })
1745 })
1746 .collect::<DFResult<Vec<_>>>()?;
1747 match parts.as_slice() {
1748 [catalog_name, database, function] => {
1749 let catalog = self.catalogs.get(catalog_name).ok_or_else(|| {
1750 DataFusionError::Plan(format!("Unknown catalog '{catalog_name}'"))
1751 })?;
1752 Ok((
1753 Arc::clone(catalog),
1754 catalog_name.clone(),
1755 Identifier::new(database, function),
1756 ))
1757 }
1758 [database, function] => Ok((
1759 self.current_catalog()?,
1760 self.current_catalog_name(),
1761 Identifier::new(database, function),
1762 )),
1763 [function] => Ok((
1764 self.current_catalog()?,
1765 self.current_catalog_name(),
1766 Identifier::new(
1767 self.ctx
1768 .state()
1769 .config_options()
1770 .catalog
1771 .default_schema
1772 .clone(),
1773 function,
1774 ),
1775 )),
1776 _ => Err(DataFusionError::Plan(format!(
1777 "Invalid function reference: {name}"
1778 ))),
1779 }
1780 }
1781
1782 fn resolve_catalog_and_table(
1784 &self,
1785 name: &ObjectName,
1786 ) -> DFResult<(Arc<dyn Catalog>, String, Identifier)> {
1787 let parts: Vec<String> = name
1788 .0
1789 .iter()
1790 .filter_map(|p| p.as_ident().map(|id| id.value.clone()))
1791 .collect();
1792 match parts.len() {
1793 3 => {
1794 let catalog = self.catalogs.get(&parts[0]).ok_or_else(|| {
1795 DataFusionError::Plan(format!("Unknown catalog '{}'", parts[0]))
1796 })?;
1797 Ok((
1798 catalog.clone(),
1799 parts[0].clone(),
1800 Identifier::new(parts[1].clone(), parts[2].clone()),
1801 ))
1802 }
1803 2 => {
1804 let catalog = self.current_catalog()?;
1805 Ok((
1806 catalog,
1807 self.current_catalog_name(),
1808 Identifier::new(parts[0].clone(), parts[1].clone()),
1809 ))
1810 }
1811 1 => {
1812 let catalog = self.current_catalog()?;
1813 let default_schema = self
1814 .ctx
1815 .state()
1816 .config_options()
1817 .catalog
1818 .default_schema
1819 .clone();
1820 Ok((
1821 catalog,
1822 self.current_catalog_name(),
1823 Identifier::new(default_schema, parts[0].clone()),
1824 ))
1825 }
1826 _ => Err(DataFusionError::Plan(format!(
1827 "Invalid table reference: {name}"
1828 ))),
1829 }
1830 }
1831
1832 fn ensure_main_branch_write_target(name: &ObjectName, operation: &str) -> DFResult<()> {
1833 let object = name
1834 .0
1835 .last()
1836 .and_then(|part| part.as_ident())
1837 .map(|ident| ident.value.as_str())
1838 .ok_or_else(|| DataFusionError::Plan(format!("Invalid table reference: {name}")))?;
1839 let parsed = parse_object_name(object).map_err(to_datafusion_error)?;
1840 if let Some(branch) = parsed.branch() {
1841 return Err(DataFusionError::NotImplemented(format!(
1842 "{operation} on Paimon branch '{branch}' is not supported"
1843 )));
1844 }
1845 Ok(())
1846 }
1847
1848 fn resolve_table_name(&self, name: &ObjectName) -> DFResult<Identifier> {
1850 let (_catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(name)?;
1851 Ok(identifier)
1852 }
1853}
1854
1855fn validate_immutable_scalar_plan(plan: &LogicalPlan) -> DFResult<()> {
1856 let mut violation = None;
1857 plan.apply(|node| {
1858 node.apply_expressions(|expression| {
1859 expression.apply(|expression| {
1860 match expression {
1861 LogicalExpr::ScalarFunction(function)
1862 if function.func.signature().volatility != Volatility::Immutable =>
1863 {
1864 violation = Some(format!(
1865 "CREATE FUNCTION body uses non-immutable function '{}'",
1866 function.func.name()
1867 ));
1868 }
1869 LogicalExpr::HigherOrderFunction(function)
1870 if function.func.signature().volatility != Volatility::Immutable =>
1871 {
1872 violation = Some(format!(
1873 "CREATE FUNCTION body uses non-immutable function '{}'",
1874 function.func.name()
1875 ));
1876 }
1877 LogicalExpr::AggregateFunction(_) | LogicalExpr::WindowFunction(_) => {
1878 violation = Some(
1879 "CREATE FUNCTION body must be a scalar expression; aggregate and window functions are not supported"
1880 .to_string(),
1881 );
1882 }
1883 LogicalExpr::Exists(_)
1884 | LogicalExpr::InSubquery(_)
1885 | LogicalExpr::SetComparison(_)
1886 | LogicalExpr::ScalarSubquery(_) => {
1887 violation = Some(
1888 "CREATE FUNCTION body must be a scalar expression; subqueries are not supported"
1889 .to_string(),
1890 );
1891 }
1892 LogicalExpr::Unnest(_) => {
1893 violation = Some(
1894 "CREATE FUNCTION body must be a scalar expression; UNNEST is not supported"
1895 .to_string(),
1896 );
1897 }
1898 _ => {}
1899 }
1900 Ok(TreeNodeRecursion::Continue)
1901 })?;
1902 Ok(TreeNodeRecursion::Continue)
1903 })?;
1904 Ok(TreeNodeRecursion::Continue)
1905 })?;
1906 if let Some(message) = violation {
1907 return Err(DataFusionError::Plan(message));
1908 }
1909 Ok(())
1910}
1911
1912fn normalize_create_function_argument_name(
1913 identifier: datafusion::sql::sqlparser::ast::Ident,
1914) -> String {
1915 if identifier.quote_style.is_none() {
1916 if let Some(value) = identifier
1917 .value
1918 .strip_prefix('"')
1919 .and_then(|value| value.strip_suffix('"'))
1920 {
1921 return value.replace("\"\"", "\"");
1922 }
1923 }
1924 IdentNormalizer::default().normalize(identifier)
1925}
1926
1927fn validate_persistent_create_function(create_function: &CreateFunction) -> DFResult<()> {
1928 let unsupported = if create_function.or_alter {
1929 Some("CREATE OR ALTER FUNCTION is not supported")
1930 } else if create_function.or_replace {
1931 Some("CREATE OR REPLACE FUNCTION is not supported")
1932 } else if create_function.temporary {
1933 Some("CREATE TEMPORARY FUNCTION is not supported")
1934 } else if create_function.args.is_none() {
1935 Some("CREATE FUNCTION requires a parenthesized argument list")
1936 } else if create_function.called_on_null.is_some() {
1937 Some("CREATE FUNCTION NULL INPUT clauses are not supported")
1938 } else if create_function.parallel.is_some() {
1939 Some("CREATE FUNCTION PARALLEL clauses are not supported")
1940 } else if create_function.security.is_some() {
1941 Some("CREATE FUNCTION SECURITY clauses are not supported")
1942 } else if !create_function.set_params.is_empty() {
1943 Some("CREATE FUNCTION SET clauses are not supported")
1944 } else if create_function.using.is_some() {
1945 Some("CREATE FUNCTION USING clauses are not supported")
1946 } else if create_function.determinism_specifier.is_some() {
1947 Some("CREATE FUNCTION determinism specifiers are not supported")
1948 } else if create_function.options.is_some() {
1949 Some("CREATE FUNCTION OPTIONS clauses are not supported")
1950 } else if create_function.remote_connection.is_some() {
1951 Some("CREATE FUNCTION REMOTE clauses are not supported")
1952 } else {
1953 None
1954 };
1955 if let Some(message) = unsupported {
1956 return Err(DataFusionError::Plan(message.to_string()));
1957 }
1958 Ok(())
1959}
1960
1961fn parse_sql_statements(sql: &str) -> DFResult<Vec<Statement>> {
1962 let dialect = GenericDialect {};
1963 let mut tokens = Tokenizer::new(&dialect, sql)
1964 .tokenize_with_location()
1965 .map_err(|error| DataFusionError::Plan(format!("SQL parse error: {error}")))?;
1966 let significant = tokens
1967 .iter()
1968 .enumerate()
1969 .filter_map(|(index, token)| match &token.token {
1970 Token::Whitespace(_) => None,
1971 Token::Word(word) => Some((index, word.keyword)),
1972 _ => Some((index, Keyword::NoKeyword)),
1973 })
1974 .take(5)
1975 .collect::<Vec<_>>();
1976 let create_function_if_not_exists = matches!(
1977 significant.as_slice(),
1978 [
1979 (_, Keyword::CREATE),
1980 (_, Keyword::FUNCTION),
1981 (_, Keyword::IF),
1982 (_, Keyword::NOT),
1983 (_, Keyword::EXISTS)
1984 ]
1985 );
1986 let create_function_or_alter = matches!(
1987 significant.as_slice(),
1988 [
1989 (_, Keyword::CREATE),
1990 (_, Keyword::OR),
1991 (_, Keyword::ALTER),
1992 (_, Keyword::FUNCTION),
1993 ..
1994 ]
1995 );
1996 if create_function_if_not_exists {
1997 let removed = significant[2..=4]
1998 .iter()
1999 .map(|(index, _)| *index)
2000 .collect::<HashSet<_>>();
2001 tokens = tokens
2002 .into_iter()
2003 .enumerate()
2004 .filter(|(index, _)| !removed.contains(index))
2005 .map(|(_, token)| token)
2006 .collect();
2007 }
2008 let mut statements = Parser::new(&dialect)
2009 .with_tokens_with_locations(tokens)
2010 .parse_statements()
2011 .map_err(|error| DataFusionError::Plan(format!("SQL parse error: {error}")))?;
2012 if create_function_if_not_exists {
2013 let Some(Statement::CreateFunction(create_function)) = statements.first_mut() else {
2014 return Err(DataFusionError::Plan(
2015 "SQL parse error: invalid CREATE FUNCTION IF NOT EXISTS statement".to_string(),
2016 ));
2017 };
2018 create_function.if_not_exists = true;
2019 }
2020 if create_function_or_alter {
2021 let Some(Statement::CreateFunction(create_function)) = statements.first_mut() else {
2022 return Err(DataFusionError::Plan(
2023 "SQL parse error: invalid CREATE OR ALTER FUNCTION statement".to_string(),
2024 ));
2025 };
2026 create_function.or_alter = true;
2027 }
2028 Ok(statements)
2029}
2030
2031fn validate_persistent_create_view(create_view: &CreateView) -> DFResult<()> {
2032 let unsupported = if create_view.or_alter {
2033 Some("CREATE OR ALTER VIEW is not supported")
2034 } else if create_view.or_replace {
2035 Some("CREATE OR REPLACE VIEW is not supported")
2036 } else if create_view.secure {
2037 Some("CREATE SECURE VIEW is not supported")
2038 } else if create_view.copy_grants {
2039 Some("CREATE VIEW COPY GRANTS is not supported")
2040 } else if create_view.name_before_not_exists {
2041 Some("CREATE VIEW with the name before IF NOT EXISTS is not supported")
2042 } else {
2043 match &create_view.options {
2044 CreateTableOptions::None => None,
2045 CreateTableOptions::With(_) => Some("CREATE VIEW WITH options are not supported"),
2046 CreateTableOptions::Options(_) => Some("CREATE VIEW OPTIONS are not supported"),
2047 _ => Some("CREATE VIEW options are not supported"),
2048 }
2049 };
2050 if let Some(message) = unsupported {
2051 return Err(DataFusionError::Plan(message.to_string()));
2052 }
2053 if create_view.comment.is_some() {
2054 return Err(DataFusionError::Plan(
2055 "CREATE VIEW COMMENT is not supported".to_string(),
2056 ));
2057 }
2058 if !create_view.cluster_by.is_empty() {
2059 return Err(DataFusionError::Plan(
2060 "CREATE VIEW CLUSTER BY is not supported".to_string(),
2061 ));
2062 }
2063 if create_view.to.is_some() {
2064 return Err(DataFusionError::Plan(
2065 "CREATE VIEW TO is not supported".to_string(),
2066 ));
2067 }
2068 if create_view.with_no_schema_binding {
2069 return Err(DataFusionError::Plan(
2070 "CREATE VIEW WITH NO SCHEMA BINDING is not supported".to_string(),
2071 ));
2072 }
2073 if create_view.params.is_some() {
2074 return Err(DataFusionError::Plan(
2075 "CREATE VIEW view parameters are not supported".to_string(),
2076 ));
2077 }
2078 if create_view
2079 .columns
2080 .iter()
2081 .any(|column| column.data_type.is_some())
2082 {
2083 return Err(DataFusionError::Plan(
2084 "CREATE VIEW column data types are not supported".to_string(),
2085 ));
2086 }
2087 if create_view
2088 .columns
2089 .iter()
2090 .any(|column| column.options.is_some())
2091 {
2092 return Err(DataFusionError::Plan(
2093 "CREATE VIEW column options are not supported".to_string(),
2094 ));
2095 }
2096 Ok(())
2097}
2098
2099fn looks_like_create_table(sql: &str) -> bool {
2102 let bytes = sql.as_bytes();
2103 let len = bytes.len();
2104 let mut i = 0;
2105 loop {
2107 while i < len && bytes[i].is_ascii_whitespace() {
2108 i += 1;
2109 }
2110 if i + 1 < len && bytes[i] == b'-' && bytes[i + 1] == b'-' {
2111 i += 2;
2112 while i < len && bytes[i] != b'\n' {
2113 i += 1;
2114 }
2115 continue;
2116 }
2117 if i + 1 < len && bytes[i] == b'/' && bytes[i + 1] == b'*' {
2118 i += 2;
2119 while i + 1 < len {
2120 if bytes[i] == b'*' && bytes[i + 1] == b'/' {
2121 i += 2;
2122 break;
2123 }
2124 i += 1;
2125 }
2126 continue;
2127 }
2128 break;
2129 }
2130 if i + 6 > len || !bytes[i..i + 6].eq_ignore_ascii_case(b"CREATE") {
2132 return false;
2133 }
2134 i += 6;
2135 if i >= len || !bytes[i].is_ascii_whitespace() {
2136 return false;
2137 }
2138 while i < len && bytes[i].is_ascii_whitespace() {
2139 i += 1;
2140 }
2141 if i + 9 <= len && bytes[i..i + 9].eq_ignore_ascii_case(b"TEMPORARY") {
2143 i += 9;
2144 while i < len && bytes[i].is_ascii_whitespace() {
2145 i += 1;
2146 }
2147 } else if i + 4 <= len && bytes[i..i + 4].eq_ignore_ascii_case(b"TEMP") {
2148 i += 4;
2149 while i < len && bytes[i].is_ascii_whitespace() {
2150 i += 1;
2151 }
2152 }
2153 if i + 4 <= len && bytes[i..i + 4].eq_ignore_ascii_case(b"VIEW") {
2155 return false;
2156 }
2157 i + 5 <= len && bytes[i..i + 5].eq_ignore_ascii_case(b"TABLE")
2158}
2159
2160fn find_partitioned_by(sql: &str) -> Option<(usize, usize)> {
2162 let bytes = sql.as_bytes();
2163 let len = bytes.len();
2164 let mut i = 0;
2165 while i < len {
2166 match bytes[i] {
2167 b'\'' => {
2168 i += 1;
2169 while i < len {
2170 if bytes[i] == b'\'' {
2171 i += 1;
2172 if i < len && bytes[i] == b'\'' {
2173 i += 1;
2174 } else {
2175 break;
2176 }
2177 } else {
2178 i += 1;
2179 }
2180 }
2181 }
2182 b'-' if i + 1 < len && bytes[i + 1] == b'-' => {
2183 i += 2;
2184 while i < len && bytes[i] != b'\n' {
2185 i += 1;
2186 }
2187 }
2188 b'/' if i + 1 < len && bytes[i + 1] == b'*' => {
2189 i += 2;
2190 while i + 1 < len {
2191 if bytes[i] == b'*' && bytes[i + 1] == b'/' {
2192 i += 2;
2193 break;
2194 }
2195 i += 1;
2196 }
2197 }
2198 b if b.is_ascii_alphabetic() && i + 11 <= len => {
2199 if bytes[i..i + 11].eq_ignore_ascii_case(b"PARTITIONED") {
2200 let rest = &bytes[i + 11..];
2201 let ws = rest.iter().take_while(|b| b.is_ascii_whitespace()).count();
2202 if ws > 0
2203 && i + 11 + ws + 2 <= len
2204 && rest[ws..ws + 2].eq_ignore_ascii_case(b"BY")
2205 {
2206 let by_end = i + 11 + ws + 2;
2207 return Some((i, by_end));
2208 }
2209 }
2210 i += 1;
2211 }
2212 _ => {
2213 i += 1;
2214 }
2215 }
2216 }
2217 None
2218}
2219
2220fn parse_partition_column(token: &str) -> DFResult<String> {
2222 let trimmed = token.trim();
2223 if trimmed.is_empty() {
2224 return Err(DataFusionError::Plan(
2225 "Empty column name in PARTITIONED BY".to_string(),
2226 ));
2227 }
2228
2229 let first = trimmed.as_bytes()[0];
2230 if first == b'"' || first == b'`' {
2231 let mut value = String::new();
2232 let mut end = None;
2233 let mut chars = trimmed[1..].char_indices().peekable();
2234 while let Some((idx, ch)) = chars.next() {
2235 if ch == first as char {
2236 if chars.peek().is_some_and(|(_, next)| *next == first as char) {
2237 value.push(ch);
2238 chars.next();
2239 } else {
2240 end = Some(1 + idx + ch.len_utf8());
2241 break;
2242 }
2243 } else {
2244 value.push(ch);
2245 }
2246 }
2247 if let Some(end) = end {
2248 if trimmed[end..].trim().is_empty() {
2249 return Ok(value);
2250 }
2251 }
2252 return Err(DataFusionError::Plan(format!(
2253 "Invalid quoted identifier in PARTITIONED BY: {trimmed}"
2254 )));
2255 }
2256
2257 let parts: Vec<&str> = trimmed.split_whitespace().collect();
2258 match parts.len() {
2259 1 => Ok(parts[0].to_string()),
2260 _ => Err(DataFusionError::Plan(format!(
2261 "PARTITIONED BY column '{}' should not specify a type. \
2262 Use column references only, e.g. PARTITIONED BY ({})",
2263 parts[0], parts[0]
2264 ))),
2265 }
2266}
2267
2268fn split_partition_columns(inner: &str) -> DFResult<Vec<&str>> {
2269 let mut columns = Vec::new();
2270 let mut start = 0;
2271 let mut quote = None;
2272 let mut chars = inner.char_indices().peekable();
2273 while let Some((idx, ch)) = chars.next() {
2274 match quote {
2275 Some(q) if ch == q => {
2276 if chars.peek().is_some_and(|(_, next)| *next == q) {
2277 chars.next();
2278 } else {
2279 quote = None;
2280 }
2281 }
2282 Some(_) => {}
2283 None if ch == '"' || ch == '`' => quote = Some(ch),
2284 None if ch == ',' => {
2285 columns.push(&inner[start..idx]);
2286 start = idx + ch.len_utf8();
2287 }
2288 None => {}
2289 }
2290 }
2291 if quote.is_some() {
2292 return Err(DataFusionError::Plan(
2293 "Unterminated quoted identifier in PARTITIONED BY".to_string(),
2294 ));
2295 }
2296 columns.push(&inner[start..]);
2297 Ok(columns)
2298}
2299
2300fn extract_partition_by(sql: &str) -> DFResult<(String, Vec<String>)> {
2307 let Some((kw_start, by_end)) = find_partitioned_by(sql) else {
2308 return Ok((sql.to_string(), vec![]));
2309 };
2310
2311 let after_by = sql[by_end..].trim_start();
2312 let paren_start = by_end + (sql[by_end..].len() - after_by.len());
2313
2314 if !after_by.starts_with('(') {
2315 return Err(DataFusionError::Plan(
2316 "Expected '(' after PARTITIONED BY".to_string(),
2317 ));
2318 }
2319
2320 let inner_start = paren_start + 1;
2321 let mut depth = 1;
2322 let mut paren_end = None;
2323 let mut quote = None;
2324 let mut chars = sql[inner_start..].char_indices().peekable();
2325 while let Some((i, ch)) = chars.next() {
2326 match quote {
2327 Some(q) if ch == q => {
2328 if chars.peek().is_some_and(|(_, next)| *next == q) {
2329 chars.next();
2330 } else {
2331 quote = None;
2332 }
2333 }
2334 Some(_) => {}
2335 None if ch == '"' || ch == '`' => quote = Some(ch),
2336 None if ch == '(' => depth += 1,
2337 None if ch == ')' => {
2338 depth -= 1;
2339 if depth == 0 {
2340 paren_end = Some(inner_start + i);
2341 break;
2342 }
2343 }
2344 None => {}
2345 }
2346 }
2347 let paren_end = paren_end.ok_or_else(|| {
2348 DataFusionError::Plan("Unmatched '(' in PARTITIONED BY clause".to_string())
2349 })?;
2350
2351 let inner = sql[inner_start..paren_end].trim();
2352 if inner.is_empty() {
2353 return Err(DataFusionError::Plan(
2354 "PARTITIONED BY must specify at least one column".to_string(),
2355 ));
2356 }
2357
2358 let mut partition_keys = Vec::new();
2359 for token in split_partition_columns(inner)? {
2360 partition_keys.push(parse_partition_column(token)?);
2361 }
2362
2363 let clause_end = paren_end + 1;
2364 let mut rewritten = String::with_capacity(sql.len());
2365 rewritten.push_str(&sql[..kw_start]);
2366 rewritten.push_str(&sql[clause_end..]);
2367 Ok((rewritten, partition_keys))
2368}
2369
2370fn column_def_to_add_column(col: &ColumnDef) -> DFResult<SchemaChange> {
2372 let paimon_type = column_def_to_paimon_type(col)?;
2373 let comment = column_def_comment(col);
2374
2375 Ok(SchemaChange::AddColumn {
2376 field_names: vec![col.name.value.clone()],
2377 data_type: paimon_type,
2378 comment,
2379 column_move: None,
2380 })
2381}
2382
2383fn column_def_to_paimon_type(col: &ColumnDef) -> DFResult<PaimonDataType> {
2384 sql_data_type_to_paimon_type(&col.data_type, column_def_nullable(col))
2385}
2386
2387fn column_def_comment(col: &ColumnDef) -> Option<String> {
2388 col.options.iter().find_map(|opt| match &opt.option {
2389 ColumnOption::Comment(comment) => Some(comment.clone()),
2390 _ => None,
2391 })
2392}
2393
2394fn primary_key_column_name(expr: &SqlExpr) -> String {
2395 match expr {
2396 SqlExpr::Identifier(ident) => ident.value.clone(),
2397 _ => expr.to_string(),
2398 }
2399}
2400
2401fn character_length_or_default(
2402 length: &Option<CharacterLength>,
2403 default_length: u32,
2404) -> DFResult<u32> {
2405 match length {
2406 Some(CharacterLength::IntegerLength { length, .. }) => (*length).try_into().map_err(|_| {
2407 DataFusionError::Plan(format!("Character length {length} exceeds supported range"))
2408 }),
2409 Some(CharacterLength::Max) => Ok(VarCharType::MAX_LENGTH),
2410 None => Ok(default_length),
2411 }
2412}
2413
2414fn u64_length_or_default(length: Option<u64>, default_length: usize) -> DFResult<usize> {
2415 match length {
2416 Some(length) => length.try_into().map_err(|_| {
2417 DataFusionError::Plan(format!("Binary length {length} exceeds supported range"))
2418 }),
2419 None => Ok(default_length),
2420 }
2421}
2422
2423fn binary_length_or_default(length: &Option<BinaryLength>, default_length: u32) -> DFResult<u32> {
2424 match length {
2425 Some(BinaryLength::IntegerLength { length }) => (*length).try_into().map_err(|_| {
2426 DataFusionError::Plan(format!("Binary length {length} exceeds supported range"))
2427 }),
2428 Some(BinaryLength::Max) => Ok(VarBinaryType::MAX_LENGTH),
2429 None => Ok(default_length),
2430 }
2431}
2432
2433fn column_def_nullable(col: &ColumnDef) -> bool {
2434 !col.options.iter().any(|opt| {
2435 matches!(
2436 opt.option,
2437 datafusion::sql::sqlparser::ast::ColumnOption::NotNull
2438 )
2439 })
2440}
2441
2442fn sql_data_type_to_paimon_type(
2447 sql_type: &datafusion::sql::sqlparser::ast::DataType,
2448 nullable: bool,
2449) -> DFResult<PaimonDataType> {
2450 use datafusion::sql::sqlparser::ast::{
2451 ArrayElemTypeDef, DataType as SqlType, ExactNumberInfo, TimezoneInfo,
2452 };
2453
2454 match sql_type {
2455 SqlType::Boolean => Ok(PaimonDataType::Boolean(BooleanType::with_nullable(
2456 nullable,
2457 ))),
2458 SqlType::TinyInt(_) => Ok(PaimonDataType::TinyInt(TinyIntType::with_nullable(
2459 nullable,
2460 ))),
2461 SqlType::SmallInt(_) => Ok(PaimonDataType::SmallInt(SmallIntType::with_nullable(
2462 nullable,
2463 ))),
2464 SqlType::Int(_) | SqlType::Integer(_) => {
2465 Ok(PaimonDataType::Int(IntType::with_nullable(nullable)))
2466 }
2467 SqlType::BigInt(_) => Ok(PaimonDataType::BigInt(BigIntType::with_nullable(nullable))),
2468 SqlType::Float(_) | SqlType::Real => {
2469 Ok(PaimonDataType::Float(FloatType::with_nullable(nullable)))
2470 }
2471 SqlType::Double(_) | SqlType::DoublePrecision => {
2472 Ok(PaimonDataType::Double(DoubleType::with_nullable(nullable)))
2473 }
2474 SqlType::Char(length) | SqlType::Character(length) => Ok(PaimonDataType::Char(
2475 CharType::with_nullable(nullable, character_length_or_default(length, 1)? as usize)
2476 .map_err(to_datafusion_error)?,
2477 )),
2478 SqlType::Varchar(length)
2479 | SqlType::Nvarchar(length)
2480 | SqlType::CharVarying(length)
2481 | SqlType::CharacterVarying(length) => Ok(PaimonDataType::VarChar(
2482 VarCharType::with_nullable(
2483 nullable,
2484 character_length_or_default(length, VarCharType::MAX_LENGTH)?,
2485 )
2486 .map_err(to_datafusion_error)?,
2487 )),
2488 SqlType::Text | SqlType::String(_) => Ok(PaimonDataType::VarChar(
2489 VarCharType::with_nullable(nullable, VarCharType::MAX_LENGTH)
2490 .map_err(to_datafusion_error)?,
2491 )),
2492 SqlType::Binary(length) => Ok(PaimonDataType::Binary(
2493 BinaryType::with_nullable(nullable, u64_length_or_default(*length, 1)? as usize)
2494 .map_err(to_datafusion_error)?,
2495 )),
2496 SqlType::Varbinary(length) => Ok(PaimonDataType::VarBinary(
2497 VarBinaryType::try_new(
2498 nullable,
2499 binary_length_or_default(length, VarBinaryType::MAX_LENGTH)?,
2500 )
2501 .map_err(to_datafusion_error)?,
2502 )),
2503 SqlType::Bytea => Ok(PaimonDataType::VarBinary(
2504 VarBinaryType::try_new(nullable, VarBinaryType::MAX_LENGTH)
2505 .map_err(to_datafusion_error)?,
2506 )),
2507 other if other.to_string().eq_ignore_ascii_case("BYTES") => Ok(PaimonDataType::VarBinary(
2508 VarBinaryType::try_new(nullable, VarBinaryType::MAX_LENGTH)
2509 .map_err(to_datafusion_error)?,
2510 )),
2511 SqlType::Blob(_) => Ok(PaimonDataType::Blob(BlobType::with_nullable(nullable))),
2512 SqlType::Custom(name, modifiers)
2513 if name.to_string().eq_ignore_ascii_case("VARIANT") && modifiers.is_empty() =>
2514 {
2515 Ok(PaimonDataType::Variant(VariantType::with_nullable(
2516 nullable,
2517 )))
2518 }
2519 SqlType::Date => Ok(PaimonDataType::Date(DateType::with_nullable(nullable))),
2520 SqlType::Timestamp(precision, tz_info) => {
2521 let precision = match precision {
2522 Some(0) => 0,
2523 Some(1..=3) | None => 3,
2524 Some(4..=6) => 6,
2525 _ => 9,
2526 };
2527 match tz_info {
2528 TimezoneInfo::None | TimezoneInfo::WithoutTimeZone => {
2529 Ok(PaimonDataType::Timestamp(
2530 TimestampType::with_nullable(nullable, precision)
2531 .map_err(to_datafusion_error)?,
2532 ))
2533 }
2534 _ => Ok(PaimonDataType::LocalZonedTimestamp(
2535 LocalZonedTimestampType::with_nullable(nullable, precision)
2536 .map_err(to_datafusion_error)?,
2537 )),
2538 }
2539 }
2540 SqlType::Decimal(info) => {
2541 let (precision, scale) = match info {
2542 ExactNumberInfo::PrecisionAndScale(precision, scale) => {
2543 (*precision as u32, *scale as u32)
2544 }
2545 ExactNumberInfo::Precision(precision) => (*precision as u32, 0),
2546 ExactNumberInfo::None => (10, 0),
2547 };
2548 Ok(PaimonDataType::Decimal(
2549 DecimalType::with_nullable(nullable, precision, scale)
2550 .map_err(to_datafusion_error)?,
2551 ))
2552 }
2553 SqlType::Array(elem_def) => {
2554 let element_type = match elem_def {
2555 ArrayElemTypeDef::AngleBracket(t)
2556 | ArrayElemTypeDef::SquareBracket(t, _)
2557 | ArrayElemTypeDef::Parenthesis(t) => sql_data_type_to_paimon_type(t, true)?,
2558 ArrayElemTypeDef::None => {
2559 return Err(DataFusionError::Plan(
2560 "ARRAY type requires an element type".to_string(),
2561 ));
2562 }
2563 };
2564 Ok(PaimonDataType::Array(PaimonArrayType::with_nullable(
2565 nullable,
2566 element_type,
2567 )))
2568 }
2569 SqlType::Map(key_type, value_type) => {
2570 let key = sql_data_type_to_paimon_type(key_type, false)?;
2571 let value = sql_data_type_to_paimon_type(value_type, true)?;
2572 Ok(PaimonDataType::Map(PaimonMapType::with_nullable(
2573 nullable, key, value,
2574 )))
2575 }
2576 SqlType::Struct(fields, _) => {
2577 let paimon_fields = fields
2578 .iter()
2579 .enumerate()
2580 .map(|(idx, field)| {
2581 let name = field
2582 .field_name
2583 .as_ref()
2584 .map(|n| n.value.clone())
2585 .unwrap_or_default();
2586 let data_type = sql_data_type_to_paimon_type(&field.field_type, true)?;
2587 Ok(PaimonDataField::new(idx as i32, name, data_type))
2588 })
2589 .collect::<DFResult<Vec<_>>>()?;
2590 Ok(PaimonDataType::Row(PaimonRowType::with_nullable(
2591 nullable,
2592 paimon_fields,
2593 )))
2594 }
2595 _ => Err(DataFusionError::Plan(format!(
2596 "Unsupported SQL data type: {sql_type}"
2597 ))),
2598 }
2599}
2600
2601fn object_name_to_string(name: &ObjectName) -> String {
2602 name.0
2603 .iter()
2604 .filter_map(|p| p.as_ident().map(|id| id.value.clone()))
2605 .collect::<Vec<_>>()
2606 .join(".")
2607}
2608
2609fn object_name_to_single_identifier(name: &ObjectName) -> DFResult<String> {
2610 match name.0.as_slice() {
2611 [part] => part
2612 .as_ident()
2613 .map(|id| id.value.clone())
2614 .ok_or_else(|| DataFusionError::Plan(format!("Invalid column name: {name}"))),
2615 _ => Err(DataFusionError::Plan(format!(
2616 "Expected a simple column name, got: {name}"
2617 ))),
2618 }
2619}
2620
2621fn extract_options(opts: &CreateTableOptions) -> DFResult<Vec<(String, String)>> {
2623 let sql_options = match opts {
2624 CreateTableOptions::With(options)
2625 | CreateTableOptions::Options(options)
2626 | CreateTableOptions::TableProperties(options)
2627 | CreateTableOptions::Plain(options) => options,
2628 CreateTableOptions::None => return Ok(Vec::new()),
2629 };
2630 sql_options
2631 .iter()
2632 .map(|opt| match opt {
2633 SqlOption::KeyValue { key, value } => {
2634 let v = value.to_string();
2635 let v = v
2637 .strip_prefix('\'')
2638 .and_then(|s| s.strip_suffix('\''))
2639 .unwrap_or(&v)
2640 .to_string();
2641 Ok((key.value.clone(), v))
2642 }
2643 other => Err(DataFusionError::Plan(format!(
2644 "Unsupported table option: {other}"
2645 ))),
2646 })
2647 .collect()
2648}
2649
2650fn is_table_not_exist(e: &paimon::Error) -> bool {
2651 matches!(e, paimon::Error::TableNotExist { .. })
2652}
2653
2654fn parse_partition_values(
2660 exprs: &[SqlExpr],
2661 all_fields: &[PaimonDataField],
2662 partition_keys: &[String],
2663) -> DFResult<Vec<HashMap<String, Option<Datum>>>> {
2664 let field_map: HashMap<&str, &PaimonDataField> =
2665 all_fields.iter().map(|f| (f.name(), f)).collect();
2666
2667 let mut partition = HashMap::new();
2668 for expr in exprs {
2669 let (col_name, val_expr) = match expr {
2670 SqlExpr::BinaryOp {
2671 left,
2672 op: datafusion::sql::sqlparser::ast::BinaryOperator::Eq,
2673 right,
2674 } => {
2675 let col = match left.as_ref() {
2676 SqlExpr::Identifier(ident) => ident.value.clone(),
2677 other => {
2678 return Err(DataFusionError::Plan(format!(
2679 "Expected column name in partition spec, got: {other}"
2680 )))
2681 }
2682 };
2683 (col, right.as_ref())
2684 }
2685 other => {
2686 return Err(DataFusionError::Plan(format!(
2687 "Expected 'column = value' in partition spec, got: {other}"
2688 )))
2689 }
2690 };
2691
2692 if !partition_keys.iter().any(|k| k == &col_name) {
2693 return Err(DataFusionError::Plan(format!(
2694 "Column '{col_name}' is not a partition column"
2695 )));
2696 }
2697
2698 let field = field_map.get(col_name.as_str()).ok_or_else(|| {
2699 DataFusionError::Plan(format!("Column '{col_name}' not found in table schema"))
2700 })?;
2701 let datum = sql_expr_to_datum(val_expr, field.data_type())?;
2702 partition.insert(col_name, Some(datum));
2703 }
2704
2705 let missing: Vec<&str> = partition_keys
2706 .iter()
2707 .filter(|k| !partition.contains_key(k.as_str()))
2708 .map(|k| k.as_str())
2709 .collect();
2710 if !missing.is_empty() {
2711 return Err(DataFusionError::Plan(format!(
2712 "Incomplete partition spec: missing keys [{}]. All partition columns must be specified.",
2713 missing.join(", ")
2714 )));
2715 }
2716
2717 Ok(vec![partition])
2718}
2719
2720fn parse_static_partitions(
2724 exprs: &[SqlExpr],
2725 partition_fields: &[PaimonDataField],
2726 all_fields: &[PaimonDataField],
2727) -> DFResult<HashMap<String, Option<Datum>>> {
2728 let mut result = HashMap::new();
2729 let field_map: HashMap<&str, &PaimonDataField> =
2730 all_fields.iter().map(|f| (f.name(), f)).collect();
2731 let partition_names: Vec<&str> = partition_fields.iter().map(|f| f.name()).collect();
2732
2733 for expr in exprs {
2734 let (col_name, val_expr) = match expr {
2735 SqlExpr::BinaryOp {
2736 left,
2737 op: datafusion::sql::sqlparser::ast::BinaryOperator::Eq,
2738 right,
2739 } => {
2740 let col = match left.as_ref() {
2741 SqlExpr::Identifier(ident) => ident.value.clone(),
2742 other => {
2743 return Err(DataFusionError::Plan(format!(
2744 "Expected column name in PARTITION clause, got: {other}"
2745 )))
2746 }
2747 };
2748 (col, right.as_ref())
2749 }
2750 SqlExpr::Identifier(ident) => {
2753 let col_name = &ident.value;
2754 if !partition_names.contains(&col_name.as_str()) {
2755 return Err(DataFusionError::Plan(format!(
2756 "Column '{col_name}' is not a partition column"
2757 )));
2758 }
2759 continue;
2760 }
2761 other => {
2762 return Err(DataFusionError::Plan(format!(
2763 "Unsupported expression in PARTITION clause: {other}"
2764 )))
2765 }
2766 };
2767
2768 if !partition_names.contains(&col_name.as_str()) {
2769 return Err(DataFusionError::Plan(format!(
2770 "Column '{col_name}' is not a partition column"
2771 )));
2772 }
2773
2774 let field = field_map.get(col_name.as_str()).ok_or_else(|| {
2775 DataFusionError::Plan(format!("Column '{col_name}' not found in table schema"))
2776 })?;
2777 let datum = sql_expr_to_datum(val_expr, field.data_type())?;
2778 result.insert(col_name, Some(datum));
2779 }
2780
2781 Ok(result)
2782}
2783
2784fn sql_expr_to_datum(expr: &SqlExpr, data_type: &PaimonDataType) -> DFResult<Datum> {
2786 let (value, negate) = match expr {
2787 SqlExpr::Value(v) => (&v.value, false),
2788 SqlExpr::UnaryOp {
2789 op: datafusion::sql::sqlparser::ast::UnaryOperator::Minus,
2790 expr: inner,
2791 } => {
2792 if let SqlExpr::Value(v) = inner.as_ref() {
2793 (&v.value, true)
2794 } else {
2795 return Err(DataFusionError::Plan(format!(
2796 "Unsupported partition value expression: {expr}"
2797 )));
2798 }
2799 }
2800 other => {
2801 return Err(DataFusionError::Plan(format!(
2802 "Unsupported partition value expression: {other}"
2803 )))
2804 }
2805 };
2806
2807 match (value, data_type) {
2808 (SqlValue::Number(n, _), _) => parse_number_datum(n, data_type, negate),
2809 (SqlValue::SingleQuotedString(s), PaimonDataType::VarChar(_)) if !negate => {
2810 Ok(Datum::String(s.clone()))
2811 }
2812 (SqlValue::SingleQuotedString(s), PaimonDataType::Date(_)) if !negate => {
2813 let date = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d")
2814 .map_err(|e| DataFusionError::Plan(format!("Invalid DATE '{s}': {e}")))?;
2815 let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
2816 Ok(Datum::Date((date - epoch).num_days() as i32))
2817 }
2818 (SqlValue::Boolean(b), PaimonDataType::Boolean(_)) if !negate => Ok(Datum::Bool(*b)),
2819 _ if negate => Err(DataFusionError::Plan(format!(
2820 "Cannot negate value for type {data_type:?}"
2821 ))),
2822 _ => Err(DataFusionError::Plan(format!(
2823 "Cannot convert {value} to {data_type:?}"
2824 ))),
2825 }
2826}
2827
2828fn parse_number_datum(n: &str, data_type: &PaimonDataType, negate: bool) -> DFResult<Datum> {
2829 let s: String = if negate {
2830 format!("-{n}")
2831 } else {
2832 n.to_string()
2833 };
2834 match data_type {
2835 PaimonDataType::TinyInt(_) => {
2836 Ok(Datum::TinyInt(s.parse::<i8>().map_err(|e| {
2837 DataFusionError::Plan(format!("Invalid TINYINT: {e}"))
2838 })?))
2839 }
2840 PaimonDataType::SmallInt(_) => {
2841 Ok(Datum::SmallInt(s.parse::<i16>().map_err(|e| {
2842 DataFusionError::Plan(format!("Invalid SMALLINT: {e}"))
2843 })?))
2844 }
2845 PaimonDataType::Int(_) => {
2846 Ok(Datum::Int(s.parse::<i32>().map_err(|e| {
2847 DataFusionError::Plan(format!("Invalid INT: {e}"))
2848 })?))
2849 }
2850 PaimonDataType::BigInt(_) => {
2851 Ok(Datum::Long(s.parse::<i64>().map_err(|e| {
2852 DataFusionError::Plan(format!("Invalid BIGINT: {e}"))
2853 })?))
2854 }
2855 PaimonDataType::Float(_) => {
2856 Ok(Datum::Float(s.parse::<f32>().map_err(|e| {
2857 DataFusionError::Plan(format!("Invalid FLOAT: {e}"))
2858 })?))
2859 }
2860 PaimonDataType::Double(_) => {
2861 Ok(Datum::Double(s.parse::<f64>().map_err(|e| {
2862 DataFusionError::Plan(format!("Invalid DOUBLE: {e}"))
2863 })?))
2864 }
2865 _ => Err(DataFusionError::Plan(format!(
2866 "Cannot convert {n} to {data_type:?}"
2867 ))),
2868 }
2869}
2870
2871fn append_partition_columns(
2873 batch: &RecordBatch,
2874 partitions: &HashMap<String, Option<Datum>>,
2875 expected_source_cols: usize,
2876 all_fields: &[PaimonDataField],
2877) -> DFResult<RecordBatch> {
2878 let num_rows = batch.num_rows();
2879
2880 let mut columns: Vec<(String, ArrayRef)> = Vec::with_capacity(all_fields.len());
2881
2882 let mut source_col_idx = 0;
2883 for field in all_fields {
2884 let name = field.name().to_string();
2885 if let Some(datum_opt) = partitions.get(&name) {
2886 let array = datum_to_constant_array(datum_opt, field.data_type(), num_rows)?;
2887 columns.push((name, array));
2888 } else {
2889 if source_col_idx >= batch.num_columns() {
2890 return Err(DataFusionError::Plan(format!(
2891 "Source query has fewer columns than expected non-partition columns. \
2892 Expected column '{name}' at position {source_col_idx}"
2893 )));
2894 }
2895 let col = batch.column(source_col_idx).clone();
2896 let target_type = paimon::arrow::paimon_type_to_arrow(field.data_type())
2897 .map_err(to_datafusion_error)?;
2898 let col = if col.data_type() != &target_type {
2899 cast(&col, &target_type).map_err(|e| {
2900 DataFusionError::Plan(format!(
2901 "Cannot cast column '{name}' from {:?} to {:?}: {e}",
2902 col.data_type(),
2903 target_type
2904 ))
2905 })?
2906 } else {
2907 col
2908 };
2909 columns.push((name, col));
2910 source_col_idx += 1;
2911 }
2912 }
2913
2914 if source_col_idx != batch.num_columns() || source_col_idx != expected_source_cols {
2915 return Err(DataFusionError::Plan(format!(
2916 "Source query has {} columns, but expected {} non-partition columns",
2917 batch.num_columns(),
2918 expected_source_cols
2919 )));
2920 }
2921
2922 let fields: Vec<Field> = columns
2923 .iter()
2924 .map(|(name, arr)| Field::new(name, arr.data_type().clone(), true))
2925 .collect();
2926 let schema = Arc::new(Schema::new(fields));
2927 let arrays: Vec<ArrayRef> = columns.into_iter().map(|(_, arr)| arr).collect();
2928 RecordBatch::try_new(schema, arrays).map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
2929}
2930
2931fn datum_to_constant_array(
2934 datum: &Option<Datum>,
2935 data_type: &PaimonDataType,
2936 num_rows: usize,
2937) -> DFResult<ArrayRef> {
2938 match datum {
2939 None => {
2940 let arrow_type =
2941 paimon::arrow::paimon_type_to_arrow(data_type).map_err(to_datafusion_error)?;
2942 Ok(new_null_array(&arrow_type, num_rows))
2943 }
2944 Some(d) => match d {
2945 Datum::Bool(v) => Ok(Arc::new(BooleanArray::from(vec![*v; num_rows]))),
2946 Datum::TinyInt(v) => Ok(Arc::new(Int8Array::from(vec![*v; num_rows]))),
2947 Datum::SmallInt(v) => Ok(Arc::new(Int16Array::from(vec![*v; num_rows]))),
2948 Datum::Int(v) => Ok(Arc::new(Int32Array::from(vec![*v; num_rows]))),
2949 Datum::Long(v) => Ok(Arc::new(Int64Array::from(vec![*v; num_rows]))),
2950 Datum::Float(v) => Ok(Arc::new(Float32Array::from(vec![*v; num_rows]))),
2951 Datum::Double(v) => Ok(Arc::new(Float64Array::from(vec![*v; num_rows]))),
2952 Datum::String(v) => Ok(Arc::new(StringArray::from(vec![v.as_str(); num_rows]))),
2953 Datum::Date(v) => Ok(Arc::new(Date32Array::from(vec![*v; num_rows]))),
2954 Datum::Time(_)
2955 | Datum::Timestamp { .. }
2956 | Datum::LocalZonedTimestamp { .. }
2957 | Datum::Decimal { .. }
2958 | Datum::Bytes(_)
2959 | Datum::Variant { .. } => Err(DataFusionError::Plan(format!(
2960 "Unsupported datum type for partition column: {d}"
2961 ))),
2962 },
2963 }
2964}
2965
2966struct VersionAsOfInfo {
2967 table_name: String,
2968 version: String,
2969 clause_range: (usize, usize),
2971}
2972
2973struct TimestampAsOfInfo {
2974 table_name: String,
2975 timestamp: String,
2976 clause_range: (usize, usize),
2978}
2979
2980fn contains_time_travel_keyword(sql: &str) -> bool {
2984 let lower = sql.to_lowercase();
2985 let bytes = lower.as_bytes();
2986 let len = bytes.len();
2987 let mut i = 0;
2988 while i < len {
2989 match bytes[i] {
2990 b'\'' => {
2991 i += 1;
2993 while i < len {
2994 if bytes[i] == b'\'' {
2995 i += 1;
2996 if i < len && bytes[i] == b'\'' {
2997 i += 1; } else {
2999 break;
3000 }
3001 } else {
3002 i += 1;
3003 }
3004 }
3005 }
3006 b'-' if i + 1 < len && bytes[i + 1] == b'-' => {
3007 i += 2;
3009 while i < len && bytes[i] != b'\n' {
3010 i += 1;
3011 }
3012 }
3013 b'/' if i + 1 < len && bytes[i + 1] == b'*' => {
3014 i += 2;
3016 while i + 1 < len {
3017 if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3018 i += 2;
3019 break;
3020 }
3021 i += 1;
3022 }
3023 }
3024 _ => {
3025 if i + 14 <= len && bytes[i..i + 14].eq_ignore_ascii_case(b"version as of ") {
3027 return true;
3028 }
3029 if i + 16 <= len && bytes[i..i + 16].eq_ignore_ascii_case(b"timestamp as of ") {
3030 return true;
3031 }
3032 i += 1;
3033 }
3034 }
3035 }
3036 false
3037}
3038
3039fn extract_all_version_as_of(sql: &str) -> Vec<VersionAsOfInfo> {
3042 let lower = sql.to_lowercase();
3043 let bytes = lower.as_bytes();
3044 let len = bytes.len();
3045 let sql_bytes = sql.as_bytes();
3046 let mut i = 0;
3047 let mut results = Vec::new();
3048
3049 while i < len {
3050 match bytes[i] {
3051 b'\'' => {
3052 i += 1;
3054 while i < len {
3055 if sql_bytes[i] == b'\'' {
3056 i += 1;
3057 if i < len && sql_bytes[i] == b'\'' {
3058 i += 1; } else {
3060 break;
3061 }
3062 } else {
3063 i += 1;
3064 }
3065 }
3066 }
3067 b'-' if i + 1 < len && bytes[i + 1] == b'-' => {
3068 i += 2;
3070 while i < len && bytes[i] != b'\n' {
3071 i += 1;
3072 }
3073 }
3074 b'/' if i + 1 < len && bytes[i + 1] == b'*' => {
3075 i += 2;
3077 while i + 1 < len {
3078 if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3079 i += 2;
3080 break;
3081 }
3082 i += 1;
3083 }
3084 }
3085 _ => {
3086 if i + 14 <= len && bytes[i..i + 14].eq_ignore_ascii_case(b"version as of ") {
3087 let kw_start = i;
3088 let val_start = i + 14;
3089 let remaining = &sql[val_start..];
3090
3091 let version = if let Some(after_quote) = remaining.strip_prefix('\'') {
3093 if let Some(close_quote) = after_quote.find('\'') {
3095 after_quote[..close_quote].to_string()
3096 } else {
3097 i += 1;
3098 continue;
3099 }
3100 } else {
3101 let v: String = remaining
3103 .chars()
3104 .take_while(|c| c.is_ascii_digit())
3105 .collect();
3106 if v.is_empty() {
3107 i += 1;
3108 continue;
3109 }
3110 v
3111 };
3112
3113 let is_quoted = remaining.starts_with('\'');
3114 let val_end = if is_quoted {
3115 val_start + version.len() + 2 } else {
3117 val_start + version.len()
3118 };
3119
3120 let table_end = sql[..kw_start].trim_end_matches(' ').len();
3122 let table_start = sql[..table_end]
3123 .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(')
3124 .map(|idx| idx + 1)
3125 .unwrap_or(0);
3126 let table_name = sql[table_start..table_end].to_string();
3127
3128 if !table_name.is_empty() {
3129 results.push(VersionAsOfInfo {
3130 table_name,
3131 version,
3132 clause_range: (table_start, val_end),
3133 });
3134 }
3135
3136 i = val_end;
3137 } else {
3138 i += 1;
3139 }
3140 }
3141 }
3142 }
3143
3144 results
3145}
3146
3147fn extract_all_timestamp_as_of(sql: &str) -> Vec<TimestampAsOfInfo> {
3150 let lower = sql.to_lowercase();
3151 let bytes = lower.as_bytes();
3152 let len = bytes.len();
3153 let sql_bytes = sql.as_bytes();
3154 let mut i = 0;
3155 let mut results = Vec::new();
3156
3157 while i < len {
3158 match bytes[i] {
3159 b'\'' => {
3160 i += 1;
3162 while i < len {
3163 if sql_bytes[i] == b'\'' {
3164 i += 1;
3165 if i < len && sql_bytes[i] == b'\'' {
3166 i += 1; } else {
3168 break;
3169 }
3170 } else {
3171 i += 1;
3172 }
3173 }
3174 }
3175 b'-' if i + 1 < len && bytes[i + 1] == b'-' => {
3176 i += 2;
3178 while i < len && bytes[i] != b'\n' {
3179 i += 1;
3180 }
3181 }
3182 b'/' if i + 1 < len && bytes[i + 1] == b'*' => {
3183 i += 2;
3185 while i + 1 < len {
3186 if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3187 i += 2;
3188 break;
3189 }
3190 i += 1;
3191 }
3192 }
3193 _ => {
3194 if i + 16 <= len && bytes[i..i + 16].eq_ignore_ascii_case(b"timestamp as of ") {
3195 let kw_start = i;
3196 let val_start = i + 16;
3197 let remaining = &sql[val_start..];
3198
3199 if !remaining.starts_with('\'') {
3201 i += 1;
3202 continue;
3203 }
3204 if let Some(close_quote) = remaining[1..].find('\'') {
3205 let timestamp = remaining[1..close_quote + 1].to_string();
3206 let val_end = val_start + close_quote + 2; let table_end = sql[..kw_start].trim_end_matches(' ').len();
3210 let table_start = sql[..table_end]
3211 .rfind(|c: char| c.is_whitespace() || c == ',' || c == '(')
3212 .map(|idx| idx + 1)
3213 .unwrap_or(0);
3214 let table_name = sql[table_start..table_end].to_string();
3215
3216 if !table_name.is_empty() {
3217 results.push(TimestampAsOfInfo {
3218 table_name,
3219 timestamp,
3220 clause_range: (table_start, val_end),
3221 });
3222 }
3223
3224 i = val_end;
3225 } else {
3226 i += 1;
3227 }
3228 } else {
3229 i += 1;
3230 }
3231 }
3232 }
3233 }
3234
3235 results
3236}
3237
3238fn ok_result(ctx: &SessionContext) -> DFResult<DataFrame> {
3240 let schema = Arc::new(Schema::new(vec![Field::new(
3241 "result",
3242 ArrowDataType::Utf8,
3243 false,
3244 )]));
3245 let batch = RecordBatch::try_new(
3246 schema.clone(),
3247 vec![Arc::new(StringArray::from(vec!["OK"]))],
3248 )?;
3249 let df = ctx.read_batch(batch)?;
3250 Ok(df)
3251}
3252
3253fn register_table_functions(
3257 ctx: &SessionContext,
3258 catalog: &Arc<dyn Catalog>,
3259 default_database: &str,
3260) {
3261 crate::blob_view::register_blob_view(ctx, Arc::clone(catalog), default_database);
3262 crate::vector_search::register_vector_search(ctx, Arc::clone(catalog), default_database);
3263 #[cfg(feature = "fulltext")]
3264 crate::full_text_search::register_full_text_search(ctx, Arc::clone(catalog), default_database);
3265 crate::hybrid_search::register_hybrid_search(ctx, Arc::clone(catalog), default_database);
3266}
3267
3268#[cfg(test)]
3269mod tests {
3270 use super::*;
3271 use std::collections::HashMap;
3272 use std::sync::Mutex;
3273
3274 use async_trait::async_trait;
3275 use datafusion::arrow::array::StringViewArray;
3276 use paimon::catalog::Database;
3277 use paimon::spec::{
3278 DataField as PaimonDataField, DataType as PaimonDataType, IntType, Schema as PaimonSchema,
3279 };
3280 use paimon::table::Table;
3281
3282 #[allow(clippy::enum_variant_names)]
3285 #[derive(Debug)]
3286 enum CatalogCall {
3287 CreateTable {
3288 identifier: Identifier,
3289 schema: PaimonSchema,
3290 ignore_if_exists: bool,
3291 },
3292 AlterTable {
3293 identifier: Identifier,
3294 changes: Vec<SchemaChange>,
3295 ignore_if_not_exists: bool,
3296 },
3297 RenameTable {
3298 from: Identifier,
3299 to: Identifier,
3300 ignore_if_not_exists: bool,
3301 },
3302 }
3303
3304 struct MockCatalog {
3305 calls: Mutex<Vec<CatalogCall>>,
3306 existing_table: Mutex<Option<Table>>,
3307 functions: Mutex<HashMap<Identifier, paimon::catalog::Function>>,
3308 views: Mutex<HashMap<Identifier, paimon::catalog::View>>,
3309 drop_view_supported: bool,
3310 }
3311
3312 impl MockCatalog {
3313 fn new() -> Self {
3314 Self {
3315 calls: Mutex::new(Vec::new()),
3316 existing_table: Mutex::new(None),
3317 functions: Mutex::new(HashMap::new()),
3318 views: Mutex::new(HashMap::new()),
3319 drop_view_supported: true,
3320 }
3321 }
3322
3323 fn without_drop_view_support() -> Self {
3324 Self {
3325 drop_view_supported: false,
3326 ..Self::new()
3327 }
3328 }
3329
3330 fn take_calls(&self) -> Vec<CatalogCall> {
3331 std::mem::take(&mut *self.calls.lock().unwrap())
3332 }
3333
3334 fn add_function(&self, function: paimon::catalog::Function) {
3335 self.functions
3336 .lock()
3337 .unwrap()
3338 .insert(function.identifier().clone(), function);
3339 }
3340
3341 fn add_view(&self, view: paimon::catalog::View) {
3342 self.views
3343 .lock()
3344 .unwrap()
3345 .insert(view.identifier().clone(), view);
3346 }
3347 }
3348
3349 #[async_trait]
3350 impl Catalog for MockCatalog {
3351 async fn list_databases(&self) -> paimon::Result<Vec<String>> {
3352 Ok(vec![])
3353 }
3354 async fn create_database(
3355 &self,
3356 _name: &str,
3357 _ignore_if_exists: bool,
3358 _properties: HashMap<String, String>,
3359 ) -> paimon::Result<()> {
3360 Ok(())
3361 }
3362 async fn get_database(&self, _name: &str) -> paimon::Result<Database> {
3363 Ok(Database::new(_name.to_string(), HashMap::new(), None))
3364 }
3365 async fn drop_database(
3366 &self,
3367 _name: &str,
3368 _ignore_if_not_exists: bool,
3369 _cascade: bool,
3370 ) -> paimon::Result<()> {
3371 Ok(())
3372 }
3373 async fn get_table(&self, _identifier: &Identifier) -> paimon::Result<Table> {
3374 if let Some(table) = self.existing_table.lock().unwrap().clone() {
3375 return Ok(table);
3376 }
3377 Err(paimon::Error::TableNotExist {
3378 full_name: _identifier.to_string(),
3379 })
3380 }
3381 async fn list_tables(&self, _database_name: &str) -> paimon::Result<Vec<String>> {
3382 Ok(vec![])
3383 }
3384 async fn create_table(
3385 &self,
3386 identifier: &Identifier,
3387 creation: PaimonSchema,
3388 ignore_if_exists: bool,
3389 ) -> paimon::Result<()> {
3390 self.calls.lock().unwrap().push(CatalogCall::CreateTable {
3391 identifier: identifier.clone(),
3392 schema: creation,
3393 ignore_if_exists,
3394 });
3395 Ok(())
3396 }
3397 async fn drop_table(
3398 &self,
3399 _identifier: &Identifier,
3400 _ignore_if_not_exists: bool,
3401 ) -> paimon::Result<()> {
3402 Ok(())
3403 }
3404 async fn rename_table(
3405 &self,
3406 from: &Identifier,
3407 to: &Identifier,
3408 ignore_if_not_exists: bool,
3409 ) -> paimon::Result<()> {
3410 self.calls.lock().unwrap().push(CatalogCall::RenameTable {
3411 from: from.clone(),
3412 to: to.clone(),
3413 ignore_if_not_exists,
3414 });
3415 Ok(())
3416 }
3417 async fn alter_table(
3418 &self,
3419 identifier: &Identifier,
3420 changes: Vec<SchemaChange>,
3421 ignore_if_not_exists: bool,
3422 ) -> paimon::Result<()> {
3423 self.calls.lock().unwrap().push(CatalogCall::AlterTable {
3424 identifier: identifier.clone(),
3425 changes,
3426 ignore_if_not_exists,
3427 });
3428 Ok(())
3429 }
3430
3431 async fn list_functions(&self, database_name: &str) -> paimon::Result<Vec<String>> {
3432 Ok(self
3433 .functions
3434 .lock()
3435 .unwrap()
3436 .keys()
3437 .filter(|identifier| identifier.database() == database_name)
3438 .map(|identifier| identifier.object().to_string())
3439 .collect())
3440 }
3441
3442 async fn create_function(
3443 &self,
3444 function: &paimon::catalog::Function,
3445 ignore_if_exists: bool,
3446 ) -> paimon::Result<()> {
3447 let mut functions = self.functions.lock().unwrap();
3448 if functions.contains_key(function.identifier()) {
3449 if ignore_if_exists {
3450 return Ok(());
3451 }
3452 return Err(paimon::Error::FunctionAlreadyExist {
3453 full_name: function.full_name(),
3454 });
3455 }
3456 functions.insert(function.identifier().clone(), function.clone());
3457 Ok(())
3458 }
3459
3460 async fn get_function(
3461 &self,
3462 identifier: &Identifier,
3463 ) -> paimon::Result<paimon::catalog::Function> {
3464 self.functions
3465 .lock()
3466 .unwrap()
3467 .get(identifier)
3468 .cloned()
3469 .ok_or_else(|| paimon::Error::FunctionNotExist {
3470 full_name: identifier.full_name(),
3471 })
3472 }
3473
3474 async fn list_views(&self, database_name: &str) -> paimon::Result<Vec<String>> {
3475 Ok(self
3476 .views
3477 .lock()
3478 .unwrap()
3479 .keys()
3480 .filter(|identifier| identifier.database() == database_name)
3481 .map(|identifier| identifier.object().to_string())
3482 .collect())
3483 }
3484
3485 async fn get_view(&self, identifier: &Identifier) -> paimon::Result<paimon::catalog::View> {
3486 self.views
3487 .lock()
3488 .unwrap()
3489 .get(identifier)
3490 .cloned()
3491 .ok_or_else(|| paimon::Error::ViewNotExist {
3492 full_name: identifier.full_name(),
3493 })
3494 }
3495
3496 async fn create_view(
3497 &self,
3498 identifier: &Identifier,
3499 schema: paimon::catalog::ViewSchema,
3500 ignore_if_exists: bool,
3501 ) -> paimon::Result<()> {
3502 let mut views = self.views.lock().unwrap();
3503 if views.contains_key(identifier) {
3504 if ignore_if_exists {
3505 return Ok(());
3506 }
3507 return Err(paimon::Error::ViewAlreadyExist {
3508 full_name: identifier.full_name(),
3509 });
3510 }
3511 views.insert(
3512 identifier.clone(),
3513 paimon::catalog::View::new(identifier.clone(), schema),
3514 );
3515 Ok(())
3516 }
3517
3518 async fn drop_view(
3519 &self,
3520 identifier: &Identifier,
3521 ignore_if_not_exists: bool,
3522 ) -> paimon::Result<()> {
3523 if !self.drop_view_supported {
3524 return Err(paimon::Error::Unsupported {
3525 message: "Catalog does not support views".to_string(),
3526 });
3527 }
3528 if self.views.lock().unwrap().remove(identifier).is_some() || ignore_if_not_exists {
3529 Ok(())
3530 } else {
3531 Err(paimon::Error::ViewNotExist {
3532 full_name: identifier.full_name(),
3533 })
3534 }
3535 }
3536 }
3537
3538 async fn make_sql_context(catalog: Arc<MockCatalog>) -> SQLContext {
3539 let mut ctx = SQLContext::new();
3540 ctx.register_catalog("paimon", catalog).await.unwrap();
3541 ctx
3542 }
3543
3544 fn add_unary_sql_function(
3545 catalog: &MockCatalog,
3546 name: &str,
3547 definition: &str,
3548 deterministic: bool,
3549 ) {
3550 add_unary_sql_function_in_database(catalog, "default", name, definition, deterministic);
3551 }
3552
3553 fn add_unary_sql_function_in_database(
3554 catalog: &MockCatalog,
3555 database: &str,
3556 name: &str,
3557 definition: &str,
3558 deterministic: bool,
3559 ) {
3560 let input_params: Vec<PaimonDataField> = serde_json::from_value(serde_json::json!([
3561 {"id": 0, "name": "x", "type": "BIGINT"}
3562 ]))
3563 .unwrap();
3564 let return_params: Vec<PaimonDataField> = serde_json::from_value(serde_json::json!([
3565 {"id": 0, "name": "result", "type": "BIGINT"}
3566 ]))
3567 .unwrap();
3568 catalog.add_function(paimon::catalog::Function::new(
3569 Identifier::new(database, name),
3570 Some(input_params),
3571 Some(return_params),
3572 deterministic,
3573 HashMap::from([(
3574 "datafusion".to_string(),
3575 paimon::catalog::FunctionDefinition::Sql {
3576 definition: definition.to_string(),
3577 },
3578 )]),
3579 None,
3580 HashMap::new(),
3581 ));
3582 }
3583
3584 fn add_plus_one_function(catalog: &MockCatalog) {
3585 add_unary_sql_function(catalog, "plus_one", "x + 1", true);
3586 }
3587
3588 fn add_constant_view(catalog: &MockCatalog) {
3589 let schema = serde_json::from_value(serde_json::json!({
3590 "fields": [
3591 {"id": 0, "name": "answer", "type": "BIGINT"}
3592 ],
3593 "query": "SELECT CAST(0 AS BIGINT) AS answer",
3594 "dialects": {
3595 "datafusion": "SELECT CAST(42 AS INT) AS source_answer"
3596 },
3597 "comment": null,
3598 "options": {}
3599 }))
3600 .unwrap();
3601 catalog.add_view(paimon::catalog::View::new(
3602 Identifier::new("default", "answer_view"),
3603 schema,
3604 ));
3605 }
3606
3607 fn add_bigint_view(catalog: &MockCatalog, database: &str, name: &str, query: &str) {
3608 let schema = serde_json::from_value(serde_json::json!({
3609 "fields": [
3610 {"id": 0, "name": "answer", "type": "BIGINT"}
3611 ],
3612 "query": query,
3613 "dialects": {},
3614 "comment": null,
3615 "options": {}
3616 }))
3617 .unwrap();
3618 catalog.add_view(paimon::catalog::View::new(
3619 Identifier::new(database, name),
3620 schema,
3621 ));
3622 }
3623
3624 #[tokio::test]
3625 async fn persistent_rest_catalog_view_can_be_created_and_read() {
3626 let catalog = Arc::new(MockCatalog::new());
3627 let ctx = make_sql_context(catalog).await;
3628
3629 ctx.sql(
3630 "CREATE VIEW answer_view AS \
3631 SELECT CAST(42 AS BIGINT) AS answer",
3632 )
3633 .await
3634 .unwrap();
3635
3636 let batches = ctx
3637 .sql("SELECT * FROM answer_view")
3638 .await
3639 .unwrap()
3640 .collect()
3641 .await
3642 .unwrap();
3643 let answers = batches[0]
3644 .column_by_name("answer")
3645 .unwrap()
3646 .as_any()
3647 .downcast_ref::<Int64Array>()
3648 .unwrap();
3649 assert_eq!(answers.value(0), 42);
3650 }
3651
3652 #[tokio::test]
3653 async fn persistent_rest_catalog_view_can_be_dropped() {
3654 let catalog = Arc::new(MockCatalog::new());
3655 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3656 ctx.sql(
3657 "CREATE VIEW answer_view AS \
3658 SELECT CAST(42 AS BIGINT) AS answer",
3659 )
3660 .await
3661 .unwrap();
3662
3663 ctx.sql("DROP VIEW answer_view").await.unwrap();
3664
3665 assert!(matches!(
3666 catalog
3667 .get_view(&Identifier::new("default", "answer_view"))
3668 .await,
3669 Err(paimon::Error::ViewNotExist { .. })
3670 ));
3671 assert!(ctx.sql("SELECT * FROM answer_view").await.is_err());
3672 }
3673
3674 #[tokio::test]
3675 async fn persistent_rest_catalog_view_drop_honors_if_exists() {
3676 let catalog = Arc::new(MockCatalog::new());
3677 let ctx = make_sql_context(catalog).await;
3678
3679 let error = ctx.sql("DROP VIEW missing_view").await.unwrap_err();
3680 assert!(
3681 error
3682 .to_string()
3683 .contains("View default.missing_view does not exist"),
3684 "unexpected error: {error}"
3685 );
3686 ctx.sql("DROP VIEW IF EXISTS missing_view").await.unwrap();
3687 }
3688
3689 #[tokio::test]
3690 async fn persistent_rest_catalog_view_drop_resolves_supported_name_forms() {
3691 let catalog = Arc::new(MockCatalog::new());
3692 add_bigint_view(&catalog, "default", "bare_view", "SELECT 1 AS answer");
3693 add_bigint_view(&catalog, "other", "two_part", "SELECT 1 AS answer");
3694 add_bigint_view(&catalog, "other", "three_part", "SELECT 1 AS answer");
3695 add_bigint_view(&catalog, "default", "Quoted View", "SELECT 1 AS answer");
3696 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3697
3698 ctx.sql("DROP VIEW bare_view").await.unwrap();
3699 ctx.sql("DROP VIEW IF EXISTS other.two_part").await.unwrap();
3700 ctx.sql("DROP VIEW IF EXISTS paimon.other.three_part")
3701 .await
3702 .unwrap();
3703 ctx.sql("DROP VIEW \"Quoted View\"").await.unwrap();
3704
3705 for identifier in [
3706 Identifier::new("default", "bare_view"),
3707 Identifier::new("other", "two_part"),
3708 Identifier::new("other", "three_part"),
3709 Identifier::new("default", "Quoted View"),
3710 ] {
3711 assert!(matches!(
3712 catalog.get_view(&identifier).await,
3713 Err(paimon::Error::ViewNotExist { .. })
3714 ));
3715 }
3716 }
3717
3718 #[tokio::test]
3719 async fn persistent_rest_catalog_view_drop_rejects_unsupported_modifiers() {
3720 let cases = [
3721 ("DROP VIEW paimon.default.invalid_view CASCADE", "CASCADE"),
3722 ("DROP VIEW paimon.default.invalid_view RESTRICT", "RESTRICT"),
3723 ("DROP VIEW paimon.default.invalid_view PURGE", "PURGE"),
3724 (
3725 "DROP VIEW paimon.default.invalid_view ON default.target",
3726 "ON clauses",
3727 ),
3728 ];
3729
3730 for (sql, modifier) in cases {
3731 let catalog = Arc::new(MockCatalog::new());
3732 let ctx = make_sql_context(catalog).await;
3733 let error = ctx.sql(sql).await.unwrap_err();
3734 assert!(
3735 error.to_string().contains(modifier),
3736 "expected {modifier} error, got: {error}"
3737 );
3738 }
3739 }
3740
3741 #[tokio::test]
3742 async fn persistent_rest_catalog_view_drop_rejects_multiple_targets_before_deleting() {
3743 let catalog = Arc::new(MockCatalog::new());
3744 add_bigint_view(&catalog, "default", "first", "SELECT 1 AS answer");
3745 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3746
3747 let error = ctx
3748 .sql("DROP VIEW paimon.default.first, datafusion.public.second")
3749 .await
3750 .unwrap_err();
3751
3752 assert!(error
3753 .to_string()
3754 .contains("Persistent DROP VIEW does not support multiple views"));
3755 assert!(catalog
3756 .get_view(&Identifier::new("default", "first"))
3757 .await
3758 .is_ok());
3759 }
3760
3761 #[tokio::test]
3762 async fn persistent_rest_catalog_view_drop_propagates_unsupported_catalog() {
3763 let catalog = Arc::new(MockCatalog::without_drop_view_support());
3764 add_bigint_view(&catalog, "default", "answer_view", "SELECT 1 AS answer");
3765 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3766
3767 let error = ctx.sql("DROP VIEW answer_view").await.unwrap_err();
3768
3769 assert!(error.to_string().contains("Catalog does not support views"));
3770 assert!(catalog
3771 .get_view(&Identifier::new("default", "answer_view"))
3772 .await
3773 .is_ok());
3774 }
3775
3776 #[tokio::test]
3777 async fn persistent_rest_catalog_view_drop_delegates_non_paimon_catalog() {
3778 let catalog = Arc::new(MockCatalog::new());
3779 let ctx = make_sql_context(catalog).await;
3780 ctx.sql("CREATE VIEW datafusion.public.delegated_view AS SELECT 1 AS answer")
3781 .await
3782 .unwrap();
3783
3784 ctx.sql("DROP VIEW datafusion.public.delegated_view")
3785 .await
3786 .unwrap();
3787
3788 assert!(ctx
3789 .sql("SELECT * FROM datafusion.public.delegated_view")
3790 .await
3791 .is_err());
3792 }
3793
3794 #[tokio::test]
3795 async fn persistent_rest_catalog_function_can_be_created_and_called() {
3796 let catalog = Arc::new(MockCatalog::new());
3797 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3798
3799 ctx.sql(
3800 "CREATE FUNCTION plus_one(x BIGINT) RETURNS BIGINT \
3801 LANGUAGE SQL IMMUTABLE RETURN x + 1",
3802 )
3803 .await
3804 .unwrap();
3805
3806 let stored = catalog
3807 .get_function(&Identifier::new("default", "plus_one"))
3808 .await
3809 .unwrap();
3810 assert_eq!(stored.input_params().unwrap()[0].id(), 0);
3811 assert_eq!(stored.input_params().unwrap()[0].name(), "x");
3812 assert!(stored.input_params().unwrap()[0].data_type().is_nullable());
3813 assert_eq!(stored.return_params().unwrap()[0].id(), 0);
3814 assert_eq!(stored.return_params().unwrap()[0].name(), "result");
3815 assert!(stored.return_params().unwrap()[0].data_type().is_nullable());
3816
3817 let batches = ctx
3818 .sql("SELECT plus_one(41) AS answer")
3819 .await
3820 .unwrap()
3821 .collect()
3822 .await
3823 .unwrap();
3824 let answers = batches[0]
3825 .column_by_name("answer")
3826 .unwrap()
3827 .as_any()
3828 .downcast_ref::<Int64Array>()
3829 .unwrap();
3830 assert_eq!(answers.value(0), 42);
3831 }
3832
3833 #[tokio::test]
3834 async fn persistent_rest_catalog_function_uses_databricks_default_sql_syntax() {
3835 let catalog = Arc::new(MockCatalog::new());
3836 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3837
3838 ctx.sql("CREATE FUNCTION plus_one(x BIGINT) RETURNS BIGINT RETURN x + 1")
3839 .await
3840 .unwrap();
3841
3842 let stored = catalog
3843 .get_function(&Identifier::new("default", "plus_one"))
3844 .await
3845 .unwrap();
3846 assert!(stored.is_deterministic());
3847
3848 let batches = ctx
3849 .sql("SELECT plus_one(41) AS answer")
3850 .await
3851 .unwrap()
3852 .collect()
3853 .await
3854 .unwrap();
3855 let answers = batches[0]
3856 .column_by_name("answer")
3857 .unwrap()
3858 .as_any()
3859 .downcast_ref::<Int64Array>()
3860 .unwrap();
3861 assert_eq!(answers.value(0), 42);
3862 }
3863
3864 #[tokio::test]
3865 async fn persistent_rest_catalog_function_supports_array_argument() {
3866 let catalog = Arc::new(MockCatalog::new());
3867 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3868
3869 ctx.sql(
3870 "CREATE FUNCTION array_answer(x ARRAY<BIGINT>) \
3871 RETURNS BIGINT RETURN 42",
3872 )
3873 .await
3874 .unwrap();
3875
3876 assert!(catalog
3877 .get_function(&Identifier::new("default", "array_answer"))
3878 .await
3879 .is_ok());
3880 }
3881
3882 #[tokio::test]
3883 async fn persistent_rest_catalog_function_supports_array_return_type() {
3884 let catalog = Arc::new(MockCatalog::new());
3885 let ctx = make_sql_context(catalog).await;
3886
3887 ctx.sql(
3888 "CREATE FUNCTION singleton(x BIGINT) \
3889 RETURNS ARRAY<BIGINT> RETURN make_array(x)",
3890 )
3891 .await
3892 .unwrap();
3893
3894 let batches = ctx
3895 .sql("SELECT singleton(42) AS answer")
3896 .await
3897 .unwrap()
3898 .collect()
3899 .await
3900 .unwrap();
3901 assert_eq!(batches[0].num_rows(), 1);
3902 assert!(matches!(
3903 batches[0]
3904 .schema()
3905 .field_with_name("answer")
3906 .unwrap()
3907 .data_type(),
3908 ArrowDataType::List(_)
3909 ));
3910 }
3911
3912 #[tokio::test]
3913 async fn persistent_rest_catalog_function_normalizes_unquoted_bare_name() {
3914 let catalog = Arc::new(MockCatalog::new());
3915 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3916
3917 ctx.sql(
3918 "CREATE FUNCTION PlusOne(X BIGINT) RETURNS BIGINT \
3919 LANGUAGE SQL IMMUTABLE RETURN X + 1",
3920 )
3921 .await
3922 .unwrap();
3923
3924 let stored = catalog
3925 .get_function(&Identifier::new("default", "plusone"))
3926 .await
3927 .unwrap();
3928 assert_eq!(stored.input_params().unwrap()[0].name(), "x");
3929 let batches = ctx
3930 .sql("SELECT plusone(41) AS answer")
3931 .await
3932 .unwrap()
3933 .collect()
3934 .await
3935 .unwrap();
3936 let answers = batches[0]
3937 .column_by_name("answer")
3938 .unwrap()
3939 .as_any()
3940 .downcast_ref::<Int64Array>()
3941 .unwrap();
3942 assert_eq!(answers.value(0), 42);
3943 }
3944
3945 #[tokio::test]
3946 async fn persistent_rest_catalog_function_preserves_quoted_names() {
3947 let catalog = Arc::new(MockCatalog::new());
3948 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3949
3950 ctx.sql(
3951 "CREATE FUNCTION \"PlusOne\"(\"Input\" BIGINT) RETURNS BIGINT \
3952 LANGUAGE SQL IMMUTABLE RETURN \"Input\" + 1",
3953 )
3954 .await
3955 .unwrap();
3956
3957 let stored = catalog
3958 .get_function(&Identifier::new("default", "PlusOne"))
3959 .await
3960 .unwrap();
3961 assert_eq!(stored.input_params().unwrap()[0].name(), "Input");
3962 let batches = ctx
3963 .sql("SELECT \"PlusOne\"(41) AS answer")
3964 .await
3965 .unwrap()
3966 .collect()
3967 .await
3968 .unwrap();
3969 let answers = batches[0]
3970 .column_by_name("answer")
3971 .unwrap()
3972 .as_any()
3973 .downcast_ref::<Int64Array>()
3974 .unwrap();
3975 assert_eq!(answers.value(0), 42);
3976 }
3977
3978 #[tokio::test]
3979 async fn persistent_rest_catalog_function_if_not_exists_preserves_existing_function() {
3980 let catalog = Arc::new(MockCatalog::new());
3981 let ctx = make_sql_context(Arc::clone(&catalog)).await;
3982 ctx.sql(
3983 "CREATE FUNCTION answer() RETURNS BIGINT \
3984 LANGUAGE SQL IMMUTABLE RETURN 1",
3985 )
3986 .await
3987 .unwrap();
3988
3989 ctx.sql(
3990 "CrEaTe /* keep comments */ FuNcTiOn IF /* gap */ NOT EXISTS answer() \
3991 RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN 2",
3992 )
3993 .await
3994 .unwrap();
3995
3996 let stored = catalog
3997 .get_function(&Identifier::new("default", "answer"))
3998 .await
3999 .unwrap();
4000 assert_eq!(
4001 stored
4002 .definition("datafusion")
4003 .and_then(paimon::catalog::FunctionDefinition::sql),
4004 Some("1")
4005 );
4006 }
4007
4008 #[tokio::test]
4009 async fn persistent_rest_catalog_function_if_not_exists_validates_proposed_body() {
4010 let catalog = Arc::new(MockCatalog::new());
4011 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4012 ctx.sql(
4013 "CREATE FUNCTION answer() RETURNS BIGINT \
4014 LANGUAGE SQL IMMUTABLE RETURN 1",
4015 )
4016 .await
4017 .unwrap();
4018
4019 let error = ctx
4020 .sql(
4021 "CREATE FUNCTION IF NOT EXISTS answer() RETURNS BIGINT \
4022 LANGUAGE SQL IMMUTABLE RETURN undeclared + 1",
4023 )
4024 .await
4025 .unwrap_err();
4026
4027 assert!(error.to_string().contains("undeclared identifier"));
4028 let stored = catalog
4029 .get_function(&Identifier::new("default", "answer"))
4030 .await
4031 .unwrap();
4032 assert_eq!(
4033 stored
4034 .definition("datafusion")
4035 .and_then(paimon::catalog::FunctionDefinition::sql),
4036 Some("1")
4037 );
4038 }
4039
4040 #[tokio::test]
4041 async fn persistent_rest_catalog_function_uses_owning_database_for_dependencies() {
4042 let catalog = Arc::new(MockCatalog::new());
4043 add_unary_sql_function_in_database(&catalog, "default", "plus_one", "x + 1", true);
4044 add_unary_sql_function_in_database(&catalog, "other", "plus_one", "x + 100", true);
4045 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4046 ctx.set_current_database("other").await.unwrap();
4047
4048 ctx.sql(
4049 "CREATE FUNCTION paimon.default.wrapper(x BIGINT) RETURNS BIGINT \
4050 LANGUAGE SQL IMMUTABLE RETURN plus_one(x)",
4051 )
4052 .await
4053 .unwrap();
4054
4055 let batches = ctx
4056 .sql("SELECT paimon.default.wrapper(41) AS answer")
4057 .await
4058 .unwrap()
4059 .collect()
4060 .await
4061 .unwrap();
4062 let answers = batches[0]
4063 .column_by_name("answer")
4064 .unwrap()
4065 .as_any()
4066 .downcast_ref::<Int64Array>()
4067 .unwrap();
4068 assert_eq!(answers.value(0), 42);
4069 }
4070
4071 #[tokio::test]
4072 async fn persistent_rest_catalog_function_rejects_nondeterministic_dependency() {
4073 let catalog = Arc::new(MockCatalog::new());
4074 add_unary_sql_function(&catalog, "unstable", "x + 1", false);
4075 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4076
4077 let error = ctx
4078 .sql(
4079 "CREATE FUNCTION wrapper(x BIGINT) RETURNS BIGINT \
4080 LANGUAGE SQL IMMUTABLE RETURN unstable(x)",
4081 )
4082 .await
4083 .unwrap_err();
4084
4085 assert!(error.to_string().contains("non-deterministic"));
4086 assert!(matches!(
4087 catalog
4088 .get_function(&Identifier::new("default", "wrapper"))
4089 .await,
4090 Err(paimon::Error::FunctionNotExist { .. })
4091 ));
4092 }
4093
4094 #[tokio::test]
4095 async fn persistent_rest_catalog_function_rejects_direct_recursion_before_create() {
4096 let catalog = Arc::new(MockCatalog::new());
4097 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4098
4099 let error = ctx
4100 .sql(
4101 "CREATE FUNCTION abs(x BIGINT) RETURNS BIGINT \
4102 LANGUAGE SQL IMMUTABLE RETURN abs(x)",
4103 )
4104 .await
4105 .unwrap_err();
4106
4107 assert!(
4108 error.to_string().contains("recursive REST SQL function"),
4109 "unexpected error: {error}"
4110 );
4111 assert!(matches!(
4112 catalog
4113 .get_function(&Identifier::new("default", "abs"))
4114 .await,
4115 Err(paimon::Error::FunctionNotExist { .. })
4116 ));
4117 }
4118
4119 #[tokio::test]
4120 async fn persistent_rest_catalog_function_rejects_indirect_recursion_before_create() {
4121 let catalog = Arc::new(MockCatalog::new());
4122 add_unary_sql_function(&catalog, "existing", "candidate(x)", true);
4123 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4124
4125 let error = ctx
4126 .sql(
4127 "CREATE FUNCTION candidate(x BIGINT) RETURNS BIGINT \
4128 LANGUAGE SQL IMMUTABLE RETURN existing(x)",
4129 )
4130 .await
4131 .unwrap_err();
4132
4133 assert!(
4134 error.to_string().contains("recursive REST SQL function"),
4135 "unexpected error: {error}"
4136 );
4137 assert!(matches!(
4138 catalog
4139 .get_function(&Identifier::new("default", "candidate"))
4140 .await,
4141 Err(paimon::Error::FunctionNotExist { .. })
4142 ));
4143 }
4144
4145 #[tokio::test]
4146 async fn persistent_rest_catalog_function_rejects_volatile_datafusion_function() {
4147 let catalog = Arc::new(MockCatalog::new());
4148 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4149
4150 let error = ctx
4151 .sql(
4152 "CREATE FUNCTION random_value() RETURNS DOUBLE \
4153 LANGUAGE SQL IMMUTABLE RETURN random()",
4154 )
4155 .await
4156 .unwrap_err();
4157
4158 assert!(
4159 error
4160 .to_string()
4161 .contains("non-immutable function 'random'"),
4162 "unexpected error: {error}"
4163 );
4164 assert!(matches!(
4165 catalog
4166 .get_function(&Identifier::new("default", "random_value"))
4167 .await,
4168 Err(paimon::Error::FunctionNotExist { .. })
4169 ));
4170 }
4171
4172 #[tokio::test]
4173 async fn persistent_rest_catalog_function_rejects_unsupported_clauses() {
4174 let cases = [
4175 (
4176 "CREATE OR REPLACE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN 1",
4177 "OR REPLACE",
4178 ),
4179 (
4180 "CREATE OR ALTER FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN 1",
4181 "OR ALTER",
4182 ),
4183 (
4184 "CREATE TEMPORARY FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN 1",
4185 "TEMPORARY",
4186 ),
4187 (
4188 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE CALLED ON NULL INPUT RETURN 1",
4189 "NULL INPUT",
4190 ),
4191 (
4192 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE PARALLEL SAFE RETURN 1",
4193 "PARALLEL",
4194 ),
4195 (
4196 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE SECURITY INVOKER RETURN 1",
4197 "SECURITY",
4198 ),
4199 (
4200 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE SET search_path TO public RETURN 1",
4201 "SET",
4202 ),
4203 ];
4204
4205 for (sql, clause) in cases {
4206 let catalog = Arc::new(MockCatalog::new());
4207 let ctx = make_sql_context(catalog).await;
4208 let error = match ctx.sql(sql).await {
4209 Ok(_) => panic!("expected error for {clause}: {sql}"),
4210 Err(error) => error,
4211 };
4212 assert!(
4213 error.to_string().contains(clause),
4214 "expected error for {clause}, got: {error}"
4215 );
4216 }
4217 }
4218
4219 #[tokio::test]
4220 async fn persistent_rest_catalog_function_rejects_non_scalar_bodies() {
4221 let cases = [
4222 (
4223 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN count(1)",
4224 "aggregate and window",
4225 ),
4226 (
4227 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN row_number() OVER ()",
4228 "aggregate and window",
4229 ),
4230 (
4231 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN (SELECT 1)",
4232 "subqueries",
4233 ),
4234 ];
4235
4236 for (sql, expected) in cases {
4237 let catalog = Arc::new(MockCatalog::new());
4238 let ctx = make_sql_context(catalog).await;
4239 let error = match ctx.sql(sql).await {
4240 Ok(_) => panic!("expected error containing {expected}: {sql}"),
4241 Err(error) => error,
4242 };
4243 assert!(
4244 error.to_string().contains(expected),
4245 "expected {expected}, got: {error}"
4246 );
4247 }
4248 }
4249
4250 #[tokio::test]
4251 async fn persistent_rest_catalog_function_rejects_invalid_signature_and_body_forms() {
4252 let cases = [
4253 (
4254 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE PYTHON RETURN 1",
4255 "LANGUAGE SQL",
4256 ),
4257 (
4258 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL STABLE RETURN 1",
4259 "deterministic SQL",
4260 ),
4261 (
4262 "CREATE FUNCTION invalid() RETURNS BIGINT LANGUAGE SQL IMMUTABLE AS '1'",
4263 "RETURN expression",
4264 ),
4265 (
4266 "CREATE FUNCTION invalid() RETURNS SETOF BIGINT LANGUAGE SQL IMMUTABLE RETURN 1",
4267 "SETOF",
4268 ),
4269 (
4270 "CREATE FUNCTION invalid(BIGINT) RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN 1",
4271 "must have names",
4272 ),
4273 (
4274 "CREATE FUNCTION invalid(IN x BIGINT) RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN x",
4275 "modes and defaults",
4276 ),
4277 (
4278 "CREATE FUNCTION invalid(x BIGINT = 1) RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN x",
4279 "modes and defaults",
4280 ),
4281 (
4282 "CREATE FUNCTION invalid(X BIGINT, x BIGINT) RETURNS BIGINT LANGUAGE SQL IMMUTABLE RETURN x",
4283 "duplicate function argument",
4284 ),
4285 ];
4286
4287 for (sql, expected) in cases {
4288 let catalog = Arc::new(MockCatalog::new());
4289 let ctx = make_sql_context(catalog).await;
4290 let error = match ctx.sql(sql).await {
4291 Ok(_) => panic!("expected error containing {expected}: {sql}"),
4292 Err(error) => error,
4293 };
4294 assert!(
4295 error.to_string().contains(expected),
4296 "expected {expected}, got: {error}"
4297 );
4298 }
4299 }
4300
4301 #[tokio::test]
4302 async fn persistent_rest_catalog_function_rejects_incompatible_return_type() {
4303 let catalog = Arc::new(MockCatalog::new());
4304 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4305
4306 let error = ctx
4307 .sql(
4308 "CREATE FUNCTION invalid() RETURNS BIGINT \
4309 LANGUAGE SQL IMMUTABLE RETURN named_struct('value', 1)",
4310 )
4311 .await
4312 .unwrap_err();
4313
4314 assert!(
4315 error.to_string().to_ascii_lowercase().contains("cast"),
4316 "unexpected error: {error}"
4317 );
4318 assert!(matches!(
4319 catalog
4320 .get_function(&Identifier::new("default", "invalid"))
4321 .await,
4322 Err(paimon::Error::FunctionNotExist { .. })
4323 ));
4324 }
4325
4326 #[tokio::test]
4327 async fn persistent_rest_catalog_view_infers_type_and_nullability() {
4328 let catalog = Arc::new(MockCatalog::new());
4329 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4330
4331 ctx.sql(
4332 "CREATE VIEW paimon.default.inferred_view AS \
4333 SELECT CAST(1 AS BIGINT) AS required, CAST(NULL AS BIGINT) AS optional",
4334 )
4335 .await
4336 .unwrap();
4337
4338 let view = catalog
4339 .get_view(&Identifier::new("default", "inferred_view"))
4340 .await
4341 .unwrap();
4342 let fields = view.schema().fields();
4343 assert!(matches!(fields[0].data_type(), PaimonDataType::BigInt(_)));
4344 assert!(!fields[0].data_type().is_nullable());
4345 assert!(matches!(fields[1].data_type(), PaimonDataType::BigInt(_)));
4346 assert!(fields[1].data_type().is_nullable());
4347 }
4348
4349 #[tokio::test]
4350 async fn persistent_rest_catalog_view_expands_function_in_owning_database() {
4351 let catalog = Arc::new(MockCatalog::new());
4352 add_unary_sql_function_in_database(&catalog, "default", "plus_one", "x + 1", true);
4353 add_unary_sql_function_in_database(&catalog, "other", "plus_one", "x + 100", true);
4354 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4355 ctx.set_current_database("other").await.unwrap();
4356
4357 ctx.sql(
4358 "CREATE VIEW paimon.default.function_view AS \
4359 SELECT plus_one(41) AS answer",
4360 )
4361 .await
4362 .unwrap();
4363
4364 let view = catalog
4365 .get_view(&Identifier::new("default", "function_view"))
4366 .await
4367 .unwrap();
4368 assert!(view.query_for("datafusion").contains("plus_one(41)"));
4369 let batches = ctx
4370 .sql("SELECT * FROM paimon.default.function_view")
4371 .await
4372 .unwrap()
4373 .collect()
4374 .await
4375 .unwrap();
4376 let answers = batches[0]
4377 .column_by_name("answer")
4378 .unwrap()
4379 .as_any()
4380 .downcast_ref::<Int64Array>()
4381 .unwrap();
4382 assert_eq!(answers.value(0), 42);
4383 }
4384
4385 #[tokio::test]
4386 async fn persistent_rest_catalog_view_binds_to_owning_database() {
4387 let catalog = Arc::new(MockCatalog::new());
4388 add_bigint_view(
4389 &catalog,
4390 "default",
4391 "base_view",
4392 "SELECT CAST(42 AS BIGINT) AS answer",
4393 );
4394 let other_schema = serde_json::from_value(serde_json::json!({
4395 "fields": [
4396 {"id": 0, "name": "answer", "type": "BIGINT"},
4397 {"id": 1, "name": "extra", "type": "BIGINT"}
4398 ],
4399 "query": "SELECT CAST(7 AS BIGINT) AS answer, CAST(8 AS BIGINT) AS extra",
4400 "dialects": {},
4401 "comment": null,
4402 "options": {}
4403 }))
4404 .unwrap();
4405 catalog.add_view(paimon::catalog::View::new(
4406 Identifier::new("other", "base_view"),
4407 other_schema,
4408 ));
4409 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4410 ctx.set_current_database("other").await.unwrap();
4411
4412 ctx.sql("CREATE VIEW paimon.default.created_view AS SELECT * FROM base_view")
4413 .await
4414 .unwrap();
4415
4416 let view = catalog
4417 .get_view(&Identifier::new("default", "created_view"))
4418 .await
4419 .unwrap();
4420 assert_eq!(view.schema().fields().len(), 1);
4421 assert_eq!(view.schema().fields()[0].name(), "answer");
4422 }
4423
4424 #[tokio::test]
4425 async fn persistent_rest_catalog_view_columns_override_names() {
4426 let catalog = Arc::new(MockCatalog::new());
4427 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4428
4429 ctx.sql(
4430 "CREATE VIEW paimon.default.named_view (renamed) \
4431 AS SELECT CAST(42 AS BIGINT) AS answer",
4432 )
4433 .await
4434 .unwrap();
4435
4436 let view = catalog
4437 .get_view(&Identifier::new("default", "named_view"))
4438 .await
4439 .unwrap();
4440 assert_eq!(view.schema().fields()[0].name(), "renamed");
4441 }
4442
4443 #[tokio::test]
4444 async fn persistent_rest_catalog_view_rejects_column_count_mismatch() {
4445 let catalog = Arc::new(MockCatalog::new());
4446 let ctx = make_sql_context(catalog).await;
4447
4448 let error = ctx
4449 .sql(
4450 "CREATE VIEW paimon.default.invalid_view (first, second) \
4451 AS SELECT CAST(42 AS BIGINT) AS answer",
4452 )
4453 .await
4454 .unwrap_err();
4455
4456 assert!(error
4457 .to_string()
4458 .contains("view column list has 2 columns but query produces 1 columns"));
4459 }
4460
4461 #[tokio::test]
4462 async fn persistent_rest_catalog_view_rejects_duplicate_column_names() {
4463 let catalog = Arc::new(MockCatalog::new());
4464 let ctx = make_sql_context(catalog).await;
4465
4466 let error = ctx
4467 .sql(
4468 "CREATE VIEW paimon.default.invalid_view (duplicate, duplicate) \
4469 AS SELECT CAST(42 AS BIGINT), CAST(7 AS BIGINT)",
4470 )
4471 .await
4472 .unwrap_err();
4473
4474 assert!(
4475 error
4476 .to_string()
4477 .contains("duplicate view column name 'duplicate'"),
4478 "unexpected error: {error}"
4479 );
4480 }
4481
4482 #[tokio::test]
4483 async fn persistent_rest_catalog_view_rejects_duplicate_inferred_names() {
4484 let catalog = Arc::new(MockCatalog::new());
4485 let ctx = make_sql_context(catalog).await;
4486
4487 let error = ctx
4488 .sql(
4489 "CREATE VIEW paimon.default.invalid_view AS \
4490 SELECT CAST(42 AS BIGINT) AS duplicate, \
4491 CAST(7 AS BIGINT) AS duplicate",
4492 )
4493 .await
4494 .unwrap_err();
4495
4496 assert!(
4497 error
4498 .to_string()
4499 .contains("Projections require unique expression names"),
4500 "unexpected error: {error}"
4501 );
4502 }
4503
4504 #[tokio::test]
4505 async fn persistent_rest_catalog_view_if_not_exists_preserves_existing_view() {
4506 let catalog = Arc::new(MockCatalog::new());
4507 let ctx = make_sql_context(Arc::clone(&catalog)).await;
4508
4509 ctx.sql(
4510 "CREATE VIEW paimon.default.existing_view \
4511 AS SELECT CAST(1 AS BIGINT) AS answer",
4512 )
4513 .await
4514 .unwrap();
4515 ctx.sql(
4516 "CREATE VIEW IF NOT EXISTS paimon.default.existing_view \
4517 AS SELECT CAST(2 AS BIGINT) AS answer",
4518 )
4519 .await
4520 .unwrap();
4521
4522 let view = catalog
4523 .get_view(&Identifier::new("default", "existing_view"))
4524 .await
4525 .unwrap();
4526 assert!(view.query_for("datafusion").contains("CAST(1 AS BIGINT)"));
4527 }
4528
4529 #[tokio::test]
4530 async fn persistent_rest_catalog_view_rejects_or_replace() {
4531 let catalog = Arc::new(MockCatalog::new());
4532 let ctx = make_sql_context(catalog).await;
4533
4534 let error = ctx
4535 .sql(
4536 "CREATE OR REPLACE VIEW paimon.default.invalid_view \
4537 AS SELECT CAST(1 AS BIGINT) AS answer",
4538 )
4539 .await
4540 .unwrap_err();
4541
4542 assert!(error
4543 .to_string()
4544 .contains("CREATE OR REPLACE VIEW is not supported"));
4545 }
4546
4547 #[tokio::test]
4548 async fn persistent_rest_catalog_view_rejects_unsupported_clauses() {
4549 let cases = [
4550 (
4551 "CREATE OR ALTER VIEW paimon.default.invalid_view AS SELECT 1",
4552 "OR ALTER",
4553 ),
4554 (
4555 "CREATE SECURE VIEW paimon.default.invalid_view AS SELECT 1",
4556 "SECURE",
4557 ),
4558 (
4559 "CREATE VIEW paimon.default.invalid_view COPY GRANTS AS SELECT 1",
4560 "COPY GRANTS",
4561 ),
4562 (
4563 "CREATE VIEW paimon.default.invalid_view IF NOT EXISTS AS SELECT 1",
4564 "name before IF NOT EXISTS",
4565 ),
4566 (
4567 "CREATE VIEW paimon.default.invalid_view WITH ('key' = 'value') AS SELECT 1",
4568 "WITH options",
4569 ),
4570 (
4571 "CREATE VIEW paimon.default.invalid_view OPTIONS(key = 'value') AS SELECT 1",
4572 "OPTIONS",
4573 ),
4574 (
4575 "CREATE VIEW paimon.default.invalid_view COMMENT = 'comment' AS SELECT 1",
4576 "COMMENT",
4577 ),
4578 (
4579 "CREATE VIEW paimon.default.invalid_view CLUSTER BY (answer) AS SELECT 1 AS answer",
4580 "CLUSTER BY",
4581 ),
4582 (
4583 "CREATE VIEW paimon.default.invalid_view TO default.sink AS SELECT 1",
4584 "TO",
4585 ),
4586 (
4587 "CREATE VIEW paimon.default.invalid_view AS SELECT 1 WITH NO SCHEMA BINDING",
4588 "WITH NO SCHEMA BINDING",
4589 ),
4590 (
4591 "CREATE ALGORITHM = MERGE VIEW paimon.default.invalid_view AS SELECT 1",
4592 "view parameters",
4593 ),
4594 (
4595 "CREATE VIEW paimon.default.invalid_view (answer COMMENT 'comment') AS SELECT 1",
4596 "column options",
4597 ),
4598 ];
4599
4600 for (sql, clause) in cases {
4601 let catalog = Arc::new(MockCatalog::new());
4602 let ctx = make_sql_context(catalog).await;
4603 let error = ctx.sql(sql).await.unwrap_err();
4604 assert!(
4605 error.to_string().contains(clause),
4606 "expected error for {clause}, got: {error}"
4607 );
4608 }
4609 }
4610
4611 #[tokio::test]
4612 async fn rest_catalog_view_is_planned_lazily() {
4613 let catalog = Arc::new(MockCatalog::new());
4614 add_constant_view(&catalog);
4615 let ctx = make_sql_context(catalog).await;
4616
4617 let batches = ctx
4618 .sql("SELECT * FROM answer_view")
4619 .await
4620 .unwrap()
4621 .collect()
4622 .await
4623 .unwrap();
4624 let answers = batches[0]
4625 .column_by_name("answer")
4626 .unwrap()
4627 .as_any()
4628 .downcast_ref::<Int64Array>()
4629 .unwrap();
4630 assert_eq!(answers.value(0), 42);
4631 }
4632
4633 #[tokio::test]
4634 async fn rest_catalog_view_can_call_bare_sql_function() {
4635 let catalog = Arc::new(MockCatalog::new());
4636 add_plus_one_function(&catalog);
4637 let schema = serde_json::from_value(serde_json::json!({
4638 "fields": [
4639 {"id": 0, "name": "answer", "type": "BIGINT"}
4640 ],
4641 "query": "SELECT plus_one(41) AS answer",
4642 "dialects": {},
4643 "comment": null,
4644 "options": {}
4645 }))
4646 .unwrap();
4647 catalog.add_view(paimon::catalog::View::new(
4648 Identifier::new("default", "function_view"),
4649 schema,
4650 ));
4651 let ctx = make_sql_context(catalog).await;
4652
4653 let batches = ctx
4654 .sql("SELECT * FROM function_view")
4655 .await
4656 .unwrap()
4657 .collect()
4658 .await
4659 .unwrap();
4660 let answers = batches[0]
4661 .column_by_name("answer")
4662 .unwrap()
4663 .as_any()
4664 .downcast_ref::<Int64Array>()
4665 .unwrap();
4666 assert_eq!(answers.value(0), 42);
4667 }
4668
4669 #[tokio::test]
4670 async fn rest_catalog_view_is_discoverable_from_schema_provider() {
4671 let catalog = Arc::new(MockCatalog::new());
4672 add_constant_view(&catalog);
4673 let ctx = make_sql_context(catalog).await;
4674
4675 let schema = ctx
4676 .ctx
4677 .catalog("paimon")
4678 .unwrap()
4679 .schema("default")
4680 .unwrap();
4681
4682 assert!(schema.table_names().contains(&"answer_view".to_string()));
4683 assert!(schema.table_exist("answer_view"));
4684 }
4685
4686 #[tokio::test]
4687 async fn nested_rest_catalog_view_binds_bare_names_to_owning_database() {
4688 let catalog = Arc::new(MockCatalog::new());
4689 add_bigint_view(
4690 &catalog,
4691 "default",
4692 "base_view",
4693 "SELECT CAST(42 AS BIGINT) AS answer",
4694 );
4695 add_bigint_view(
4696 &catalog,
4697 "other",
4698 "base_view",
4699 "SELECT CAST(7 AS BIGINT) AS answer",
4700 );
4701 add_bigint_view(&catalog, "default", "outer_view", "SELECT * FROM base_view");
4702 let ctx = make_sql_context(catalog).await;
4703 ctx.set_current_database("other").await.unwrap();
4704
4705 let batches = ctx
4706 .sql("SELECT * FROM paimon.default.outer_view")
4707 .await
4708 .unwrap()
4709 .collect()
4710 .await
4711 .unwrap();
4712 let answers = batches[0]
4713 .column_by_name("answer")
4714 .unwrap()
4715 .as_any()
4716 .downcast_ref::<Int64Array>()
4717 .unwrap();
4718 assert_eq!(answers.value(0), 42);
4719 }
4720
4721 #[tokio::test]
4722 async fn recursive_rest_catalog_views_are_rejected() {
4723 let catalog = Arc::new(MockCatalog::new());
4724 add_bigint_view(
4725 &catalog,
4726 "default",
4727 "first_view",
4728 "SELECT * FROM second_view",
4729 );
4730 add_bigint_view(
4731 &catalog,
4732 "default",
4733 "second_view",
4734 "SELECT * FROM first_view",
4735 );
4736 let ctx = make_sql_context(catalog).await;
4737
4738 let error = tokio::time::timeout(
4739 std::time::Duration::from_secs(2),
4740 ctx.sql("SELECT * FROM first_view"),
4741 )
4742 .await
4743 .expect("recursive view planning should terminate")
4744 .unwrap_err()
4745 .to_string();
4746
4747 assert!(
4748 error.contains("recursive REST catalog view"),
4749 "unexpected error: {error}"
4750 );
4751 }
4752
4753 #[tokio::test]
4754 async fn rest_catalog_view_allows_cte_with_same_name() {
4755 let catalog = Arc::new(MockCatalog::new());
4756 add_bigint_view(
4757 &catalog,
4758 "default",
4759 "cte_view",
4760 "WITH wrapper AS (\
4761 WITH cte_view AS (SELECT CAST(42 AS BIGINT) AS answer) \
4762 SELECT * FROM cte_view\
4763 ) SELECT * FROM wrapper",
4764 );
4765 let ctx = make_sql_context(catalog).await;
4766
4767 let batches = ctx
4768 .sql("SELECT * FROM cte_view")
4769 .await
4770 .unwrap()
4771 .collect()
4772 .await
4773 .unwrap();
4774 let answers = batches[0]
4775 .column_by_name("answer")
4776 .unwrap()
4777 .as_any()
4778 .downcast_ref::<Int64Array>()
4779 .unwrap();
4780 assert_eq!(answers.value(0), 42);
4781 }
4782
4783 #[tokio::test]
4784 async fn rest_catalog_view_normalizes_cte_identifiers() {
4785 let catalog = Arc::new(MockCatalog::new());
4786 add_bigint_view(
4787 &catalog,
4788 "default",
4789 "cte_view",
4790 "WITH cte_view AS (SELECT CAST(42 AS BIGINT) AS answer) \
4791 SELECT * FROM \"cte_view\"",
4792 );
4793 let ctx = make_sql_context(catalog).await;
4794
4795 let batches = ctx
4796 .sql("SELECT * FROM cte_view")
4797 .await
4798 .unwrap()
4799 .collect()
4800 .await
4801 .unwrap();
4802 let answers = batches[0]
4803 .column_by_name("answer")
4804 .unwrap()
4805 .as_any()
4806 .downcast_ref::<Int64Array>()
4807 .unwrap();
4808 assert_eq!(answers.value(0), 42);
4809 }
4810
4811 #[tokio::test]
4812 async fn rest_catalog_view_rejects_non_query_sql() {
4813 let catalog = Arc::new(MockCatalog::new());
4814 add_bigint_view(
4815 &catalog,
4816 "default",
4817 "unsafe_view",
4818 "DELETE FROM missing_table",
4819 );
4820 let ctx = make_sql_context(catalog).await;
4821
4822 let error = ctx
4823 .sql("SELECT * FROM unsafe_view")
4824 .await
4825 .unwrap_err()
4826 .to_string();
4827
4828 assert!(
4829 error.contains("read-only query"),
4830 "unexpected error: {error}"
4831 );
4832 }
4833
4834 #[tokio::test]
4835 async fn bare_rest_sql_function_is_expanded() {
4836 let catalog = Arc::new(MockCatalog::new());
4837 add_plus_one_function(&catalog);
4838 let ctx = make_sql_context(catalog).await;
4839
4840 let batches = ctx
4841 .sql("SELECT plus_one(41) AS answer")
4842 .await
4843 .unwrap()
4844 .collect()
4845 .await
4846 .unwrap();
4847 let answers = batches[0]
4848 .column_by_name("answer")
4849 .unwrap()
4850 .as_any()
4851 .downcast_ref::<Int64Array>()
4852 .unwrap();
4853 assert_eq!(answers.value(0), 42);
4854 }
4855
4856 #[tokio::test]
4857 async fn rest_sql_function_normalizes_call_identifiers() {
4858 let catalog = Arc::new(MockCatalog::new());
4859 add_plus_one_function(&catalog);
4860 let ctx = make_sql_context(catalog).await;
4861
4862 let batches = ctx
4863 .sql("SELECT PAIMON.DEFAULT.PLUS_ONE(41) AS answer")
4864 .await
4865 .unwrap()
4866 .collect()
4867 .await
4868 .unwrap();
4869 let answers = batches[0]
4870 .column_by_name("answer")
4871 .unwrap()
4872 .as_any()
4873 .downcast_ref::<Int64Array>()
4874 .unwrap();
4875 assert_eq!(answers.value(0), 42);
4876 }
4877
4878 #[tokio::test]
4879 async fn rest_sql_function_accepts_outer_column_argument() {
4880 let catalog = Arc::new(MockCatalog::new());
4881 add_plus_one_function(&catalog);
4882 let ctx = make_sql_context(catalog).await;
4883
4884 let batches = ctx
4885 .sql("SELECT plus_one(x) AS answer FROM (VALUES (41)) AS t(x)")
4886 .await
4887 .unwrap()
4888 .collect()
4889 .await
4890 .unwrap();
4891 let answers = batches[0]
4892 .column_by_name("answer")
4893 .unwrap()
4894 .as_any()
4895 .downcast_ref::<Int64Array>()
4896 .unwrap();
4897 assert_eq!(answers.value(0), 42);
4898 }
4899
4900 #[tokio::test]
4901 async fn fully_qualified_rest_sql_function_is_expanded() {
4902 let catalog = Arc::new(MockCatalog::new());
4903 add_plus_one_function(&catalog);
4904 let ctx = make_sql_context(catalog).await;
4905
4906 let batches = ctx
4907 .sql("SELECT paimon.default.plus_one(41) AS answer")
4908 .await
4909 .unwrap()
4910 .collect()
4911 .await
4912 .unwrap();
4913 let answers = batches[0]
4914 .column_by_name("answer")
4915 .unwrap()
4916 .as_any()
4917 .downcast_ref::<Int64Array>()
4918 .unwrap();
4919 assert_eq!(answers.value(0), 42);
4920 }
4921
4922 #[tokio::test]
4923 async fn rest_sql_function_result_is_cast_to_declared_type() {
4924 let catalog = Arc::new(MockCatalog::new());
4925 add_unary_sql_function(&catalog, "narrow_body", "CAST(x AS INT)", true);
4926 let ctx = make_sql_context(catalog).await;
4927
4928 let batches = ctx
4929 .sql("SELECT narrow_body(41) AS answer")
4930 .await
4931 .unwrap()
4932 .collect()
4933 .await
4934 .unwrap();
4935 let answers = batches[0]
4936 .column_by_name("answer")
4937 .unwrap()
4938 .as_any()
4939 .downcast_ref::<Int64Array>()
4940 .unwrap();
4941 assert_eq!(answers.value(0), 41);
4942 }
4943
4944 #[tokio::test]
4945 async fn rest_sql_function_normalizes_definition_parameters() {
4946 let catalog = Arc::new(MockCatalog::new());
4947 add_unary_sql_function(&catalog, "uppercase_parameter", "X + 1", true);
4948 let ctx = make_sql_context(catalog).await;
4949
4950 let batches = ctx
4951 .sql("SELECT uppercase_parameter(41) AS answer")
4952 .await
4953 .unwrap()
4954 .collect()
4955 .await
4956 .unwrap();
4957 let answers = batches[0]
4958 .column_by_name("answer")
4959 .unwrap()
4960 .as_any()
4961 .downcast_ref::<Int64Array>()
4962 .unwrap();
4963 assert_eq!(answers.value(0), 42);
4964 }
4965
4966 #[tokio::test]
4967 async fn rest_sql_function_preserves_quoted_metadata_parameter() {
4968 let catalog = Arc::new(MockCatalog::new());
4969 let input_params: Vec<PaimonDataField> = serde_json::from_value(serde_json::json!([
4970 {"id": 0, "name": "X", "type": "BIGINT"}
4971 ]))
4972 .unwrap();
4973 let return_params: Vec<PaimonDataField> = serde_json::from_value(serde_json::json!([
4974 {"id": 0, "name": "result", "type": "BIGINT"}
4975 ]))
4976 .unwrap();
4977 catalog.add_function(paimon::catalog::Function::new(
4978 Identifier::new("default", "quoted_parameter"),
4979 Some(input_params),
4980 Some(return_params),
4981 true,
4982 HashMap::from([(
4983 "datafusion".to_string(),
4984 paimon::catalog::FunctionDefinition::Sql {
4985 definition: "\"X\" + 1".to_string(),
4986 },
4987 )]),
4988 None,
4989 HashMap::new(),
4990 ));
4991 let ctx = make_sql_context(catalog).await;
4992
4993 let batches = ctx
4994 .sql("SELECT quoted_parameter(41) AS answer")
4995 .await
4996 .unwrap()
4997 .collect()
4998 .await
4999 .unwrap();
5000 let answers = batches[0]
5001 .column_by_name("answer")
5002 .unwrap()
5003 .as_any()
5004 .downcast_ref::<Int64Array>()
5005 .unwrap();
5006 assert_eq!(answers.value(0), 42);
5007 }
5008
5009 #[tokio::test]
5010 async fn rest_sql_function_is_expanded_in_explain() {
5011 let catalog = Arc::new(MockCatalog::new());
5012 add_plus_one_function(&catalog);
5013 let ctx = make_sql_context(catalog).await;
5014
5015 ctx.sql("EXPLAIN SELECT plus_one(1)")
5016 .await
5017 .unwrap()
5018 .collect()
5019 .await
5020 .unwrap();
5021 }
5022
5023 #[tokio::test]
5024 async fn rest_sql_function_is_expanded_in_time_travel_query() {
5025 let temp_dir = tempfile::tempdir().unwrap();
5026 let mut options = paimon::Options::new();
5027 options.set(
5028 paimon::CatalogOptions::WAREHOUSE,
5029 temp_dir.path().to_string_lossy(),
5030 );
5031 let storage_catalog = Arc::new(paimon::FileSystemCatalog::new(options).unwrap());
5032 let mut setup = SQLContext::new();
5033 setup
5034 .register_catalog("paimon", storage_catalog.clone())
5035 .await
5036 .unwrap();
5037 setup
5038 .sql("CREATE TABLE paimon.default.time_travel_source (id INT)")
5039 .await
5040 .unwrap();
5041 setup
5042 .sql("INSERT INTO paimon.default.time_travel_source VALUES (41)")
5043 .await
5044 .unwrap()
5045 .collect()
5046 .await
5047 .unwrap();
5048
5049 let catalog = Arc::new(MockCatalog::new());
5050 *catalog.existing_table.lock().unwrap() = Some(
5051 storage_catalog
5052 .get_table(&Identifier::new("default", "time_travel_source"))
5053 .await
5054 .unwrap(),
5055 );
5056 add_plus_one_function(&catalog);
5057 let ctx = make_sql_context(catalog).await;
5058
5059 let batches = ctx
5060 .sql(
5061 "SELECT plus_one(id) AS answer \
5062 FROM time_travel_source VERSION AS OF 1",
5063 )
5064 .await
5065 .unwrap()
5066 .collect()
5067 .await
5068 .unwrap();
5069 let answers = batches[0]
5070 .column_by_name("answer")
5071 .unwrap()
5072 .as_any()
5073 .downcast_ref::<Int64Array>()
5074 .unwrap();
5075 assert_eq!(answers.value(0), 42);
5076 }
5077
5078 #[tokio::test]
5079 async fn rest_sql_function_rejects_undeclared_identifiers() {
5080 let catalog = Arc::new(MockCatalog::new());
5081 add_unary_sql_function(&catalog, "captures_column", "x + y", true);
5082 let ctx = make_sql_context(catalog).await;
5083
5084 let error = ctx
5085 .sql("SELECT captures_column(41)")
5086 .await
5087 .unwrap_err()
5088 .to_string();
5089
5090 assert!(
5091 error.contains("undeclared identifier 'y'"),
5092 "unexpected error: {error}"
5093 );
5094 }
5095
5096 #[tokio::test]
5097 async fn non_deterministic_rest_sql_function_is_rejected() {
5098 let catalog = Arc::new(MockCatalog::new());
5099 add_unary_sql_function(&catalog, "unsafe_plus_one", "x + 1", false);
5100 let ctx = make_sql_context(catalog).await;
5101
5102 let error = ctx
5103 .sql("SELECT unsafe_plus_one(41)")
5104 .await
5105 .unwrap_err()
5106 .to_string();
5107
5108 assert!(
5109 error.contains("non-deterministic"),
5110 "unexpected error: {error}"
5111 );
5112 }
5113
5114 #[tokio::test]
5115 async fn rest_sql_function_without_single_return_is_rejected() {
5116 let catalog = Arc::new(MockCatalog::new());
5117 let input_params: Vec<PaimonDataField> = serde_json::from_value(serde_json::json!([
5118 {"id": 0, "name": "x", "type": "BIGINT"}
5119 ]))
5120 .unwrap();
5121 catalog.add_function(paimon::catalog::Function::new(
5122 Identifier::new("default", "missing_return"),
5123 Some(input_params),
5124 None,
5125 true,
5126 HashMap::from([(
5127 "datafusion".to_string(),
5128 paimon::catalog::FunctionDefinition::Sql {
5129 definition: "x + 1".to_string(),
5130 },
5131 )]),
5132 None,
5133 HashMap::new(),
5134 ));
5135 let ctx = make_sql_context(catalog).await;
5136
5137 let error = ctx
5138 .sql("SELECT missing_return(41)")
5139 .await
5140 .unwrap_err()
5141 .to_string();
5142
5143 assert!(
5144 error.contains("exactly one return parameter"),
5145 "unexpected error: {error}"
5146 );
5147 }
5148
5149 #[tokio::test]
5150 async fn nested_rest_sql_functions_are_expanded() {
5151 let catalog = Arc::new(MockCatalog::new());
5152 add_plus_one_function(&catalog);
5153 add_unary_sql_function(&catalog, "plus_two", "plus_one(x) + 1", true);
5154 let ctx = make_sql_context(catalog).await;
5155
5156 let batches = ctx
5157 .sql("SELECT plus_two(40) AS answer")
5158 .await
5159 .unwrap()
5160 .collect()
5161 .await
5162 .unwrap();
5163 let answers = batches[0]
5164 .column_by_name("answer")
5165 .unwrap()
5166 .as_any()
5167 .downcast_ref::<Int64Array>()
5168 .unwrap();
5169 assert_eq!(answers.value(0), 42);
5170 }
5171
5172 #[tokio::test]
5173 async fn nested_rest_sql_function_binds_to_owning_database() {
5174 let catalog = Arc::new(MockCatalog::new());
5175 add_unary_sql_function(&catalog, "plus_one", "x + 100", true);
5176 add_unary_sql_function_in_database(&catalog, "other", "plus_one", "x + 1", true);
5177 add_unary_sql_function_in_database(&catalog, "other", "plus_two", "plus_one(x) + 1", true);
5178 let ctx = make_sql_context(catalog).await;
5179
5180 let batches = ctx
5181 .sql("SELECT paimon.other.plus_two(40) AS answer")
5182 .await
5183 .unwrap()
5184 .collect()
5185 .await
5186 .unwrap();
5187 let answers = batches[0]
5188 .column_by_name("answer")
5189 .unwrap()
5190 .as_any()
5191 .downcast_ref::<Int64Array>()
5192 .unwrap();
5193 assert_eq!(answers.value(0), 42);
5194 }
5195
5196 #[tokio::test]
5197 async fn nested_rest_sql_function_normalizes_owning_database_reference() {
5198 let catalog = Arc::new(MockCatalog::new());
5199 add_unary_sql_function(&catalog, "plus_one", "x + 100", true);
5200 add_unary_sql_function_in_database(&catalog, "other", "plus_one", "x + 1", true);
5201 add_unary_sql_function_in_database(&catalog, "other", "plus_two", "PLUS_ONE(x) + 1", true);
5202 let ctx = make_sql_context(catalog).await;
5203
5204 let batches = ctx
5205 .sql("SELECT paimon.other.plus_two(40) AS answer")
5206 .await
5207 .unwrap()
5208 .collect()
5209 .await
5210 .unwrap();
5211 let answers = batches[0]
5212 .column_by_name("answer")
5213 .unwrap()
5214 .as_any()
5215 .downcast_ref::<Int64Array>()
5216 .unwrap();
5217 assert_eq!(answers.value(0), 42);
5218 }
5219
5220 #[tokio::test]
5221 async fn recursive_rest_sql_functions_are_rejected() {
5222 let catalog = Arc::new(MockCatalog::new());
5223 add_unary_sql_function(&catalog, "first", "second(x)", true);
5224 add_unary_sql_function(&catalog, "second", "first(x)", true);
5225 let ctx = make_sql_context(catalog).await;
5226
5227 let error = ctx.sql("SELECT first(1)").await.unwrap_err().to_string();
5228
5229 assert!(error.contains("recursive"), "unexpected error: {error}");
5230 }
5231
5232 #[tokio::test]
5233 async fn branching_rest_sql_function_expansion_is_bounded() {
5234 let catalog = Arc::new(MockCatalog::new());
5235 for index in 0..3 {
5236 let next = index + 1;
5237 add_unary_sql_function(
5238 &catalog,
5239 &format!("f{index}"),
5240 &format!("f{next}(x) + f{next}(x)"),
5241 true,
5242 );
5243 }
5244 add_unary_sql_function(&catalog, "f3", "x", true);
5245 let ctx = make_sql_context(catalog).await;
5246 let statement = Parser::parse_sql(&GenericDialect {}, "SELECT f0(1)")
5247 .unwrap()
5248 .remove(0);
5249
5250 let error = crate::sql_function::expand_statement_with_budget(
5251 statement,
5252 &ctx.catalogs,
5253 &ctx.current_catalog_name(),
5254 "default",
5255 4,
5256 )
5257 .await
5258 .unwrap_err()
5259 .to_string();
5260
5261 assert!(
5262 error.contains("expansion budget"),
5263 "unexpected error: {error}"
5264 );
5265 }
5266
5267 struct ProbeTrackingCatalog {
5274 get_calls: std::sync::atomic::AtomicUsize,
5275 create_calls: std::sync::atomic::AtomicUsize,
5276 get_view_calls: std::sync::atomic::AtomicUsize,
5277 table_identifiers: std::sync::Mutex<Vec<String>>,
5278 }
5279
5280 impl ProbeTrackingCatalog {
5281 fn new() -> Self {
5282 Self {
5283 get_calls: std::sync::atomic::AtomicUsize::new(0),
5284 create_calls: std::sync::atomic::AtomicUsize::new(0),
5285 get_view_calls: std::sync::atomic::AtomicUsize::new(0),
5286 table_identifiers: std::sync::Mutex::new(Vec::new()),
5287 }
5288 }
5289 fn get_count(&self) -> usize {
5290 self.get_calls.load(std::sync::atomic::Ordering::SeqCst)
5291 }
5292 fn create_count(&self) -> usize {
5293 self.create_calls.load(std::sync::atomic::Ordering::SeqCst)
5294 }
5295 fn get_view_count(&self) -> usize {
5296 self.get_view_calls
5297 .load(std::sync::atomic::Ordering::SeqCst)
5298 }
5299 fn table_identifiers(&self) -> Vec<String> {
5300 self.table_identifiers.lock().unwrap().clone()
5301 }
5302 }
5303
5304 #[async_trait]
5305 impl Catalog for ProbeTrackingCatalog {
5306 async fn list_databases(&self) -> paimon::Result<Vec<String>> {
5307 Ok(vec![])
5308 }
5309 async fn create_database(
5310 &self,
5311 _name: &str,
5312 _ignore_if_exists: bool,
5313 _properties: HashMap<String, String>,
5314 ) -> paimon::Result<()> {
5315 self.create_calls
5316 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5317 Ok(())
5318 }
5319 async fn get_database(&self, name: &str) -> paimon::Result<Database> {
5320 self.get_calls
5321 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5322 if name == "analytics" {
5323 return Ok(Database::new(name.to_string(), HashMap::new(), None));
5324 }
5325 Err(paimon::Error::Unsupported {
5326 message: "simulated Forbidden".to_string(),
5327 })
5328 }
5329 async fn drop_database(
5330 &self,
5331 _name: &str,
5332 _ignore_if_not_exists: bool,
5333 _cascade: bool,
5334 ) -> paimon::Result<()> {
5335 Ok(())
5336 }
5337 async fn get_table(&self, identifier: &Identifier) -> paimon::Result<Table> {
5338 self.table_identifiers
5339 .lock()
5340 .unwrap()
5341 .push(identifier.to_string());
5342 Err(paimon::Error::TableNotExist {
5343 full_name: identifier.to_string(),
5344 })
5345 }
5346 async fn list_tables(&self, _database_name: &str) -> paimon::Result<Vec<String>> {
5347 Ok(vec![])
5348 }
5349 async fn create_table(
5350 &self,
5351 _identifier: &Identifier,
5352 _creation: PaimonSchema,
5353 _ignore_if_exists: bool,
5354 ) -> paimon::Result<()> {
5355 Ok(())
5356 }
5357 async fn drop_table(
5358 &self,
5359 _identifier: &Identifier,
5360 _ignore_if_not_exists: bool,
5361 ) -> paimon::Result<()> {
5362 Ok(())
5363 }
5364 async fn rename_table(
5365 &self,
5366 _from: &Identifier,
5367 _to: &Identifier,
5368 _ignore_if_not_exists: bool,
5369 ) -> paimon::Result<()> {
5370 Ok(())
5371 }
5372 async fn alter_table(
5373 &self,
5374 _identifier: &Identifier,
5375 _changes: Vec<SchemaChange>,
5376 _ignore_if_not_exists: bool,
5377 ) -> paimon::Result<()> {
5378 Ok(())
5379 }
5380
5381 async fn get_view(
5382 &self,
5383 _identifier: &Identifier,
5384 ) -> paimon::Result<paimon::catalog::View> {
5385 self.get_view_calls
5386 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
5387 Err(paimon::Error::RestApi {
5388 source: paimon::api::RestError::Forbidden {
5389 message: "view access is forbidden".to_string(),
5390 },
5391 })
5392 }
5393 }
5394
5395 #[tokio::test]
5396 async fn register_catalog_with_none_skips_default_db_probe() {
5397 let catalog = Arc::new(ProbeTrackingCatalog::new());
5398 let mut ctx = SQLContext::new();
5399 ctx.register_catalog_with_default_db("paimon", catalog.clone(), None)
5400 .await
5401 .expect("None must skip probe so Forbidden-shaped error never fires");
5402 assert_eq!(catalog.get_count(), 0, "get_database must not be called");
5403 assert_eq!(
5404 catalog.create_count(),
5405 0,
5406 "create_database must not be called"
5407 );
5408 assert_eq!(
5410 ctx.ctx().state().config().options().catalog.default_catalog,
5411 "paimon"
5412 );
5413 }
5414
5415 #[tokio::test]
5416 async fn register_catalog_with_some_default_propagates_probe_error() {
5417 let catalog = Arc::new(ProbeTrackingCatalog::new());
5418 let mut ctx = SQLContext::new();
5419 let err = ctx
5420 .register_catalog_with_default_db("paimon", catalog.clone(), Some("default"))
5421 .await
5422 .expect_err("non-DatabaseNotExist error from get_database must propagate");
5423 assert!(
5424 err.to_string().contains("simulated Forbidden"),
5425 "unexpected error: {err}"
5426 );
5427 assert_eq!(catalog.get_count(), 1);
5428 assert_eq!(catalog.create_count(), 0);
5429 }
5430
5431 #[tokio::test]
5432 async fn register_catalog_with_some_empty_string_is_rejected() {
5433 let catalog = Arc::new(ProbeTrackingCatalog::new());
5436 let mut ctx = SQLContext::new();
5437 let err = ctx
5438 .register_catalog_with_default_db("paimon", catalog.clone(), Some(""))
5439 .await
5440 .expect_err("empty default_db must be rejected at the API");
5441 assert!(
5442 err.to_string().contains("must not be empty"),
5443 "unexpected error: {err}"
5444 );
5445 assert_eq!(
5446 catalog.get_count(),
5447 0,
5448 "guard must short-circuit before any catalog call"
5449 );
5450 }
5451
5452 #[tokio::test]
5453 async fn register_catalog_with_none_table_function_resolves_bare_name_to_literal_default() {
5454 let catalog = Arc::new(ProbeTrackingCatalog::new());
5460 let mut ctx = SQLContext::new();
5461 ctx.register_catalog_with_default_db("paimon", catalog, None)
5462 .await
5463 .unwrap();
5464
5465 let err = ctx
5466 .sql("SELECT * FROM vector_search('bare', 'col', '[1.0]', 1)")
5467 .await
5468 .expect_err("bare name must error out — no `default.bare` table in mock catalog");
5469 let msg = err.to_string();
5470 assert!(
5471 msg.contains("default") && msg.contains("bare"),
5472 "error must surface the fallback 'default' namespace + bare name, got: {msg}"
5473 );
5474 }
5475
5476 #[tokio::test]
5477 async fn registered_table_function_skips_view_lookup_during_relation_preload() {
5478 let catalog = Arc::new(ProbeTrackingCatalog::new());
5479 let mut ctx = SQLContext::new();
5480 ctx.register_catalog_with_default_db("paimon", catalog.clone(), None)
5481 .await
5482 .unwrap();
5483 ctx.set_current_database("analytics").await.unwrap();
5484
5485 let _ = ctx
5486 .sql(
5487 "SELECT *
5488 FROM vector_search(
5489 'paimon.analytics.documents',
5490 'embedding',
5491 '[0.1, 0.2]',
5492 3
5493 )",
5494 )
5495 .await
5496 .expect_err("the target table is intentionally absent");
5497
5498 assert_eq!(
5499 catalog.table_identifiers(),
5500 vec![
5501 "analytics.vector_search",
5502 "analytics.documents"
5503 ],
5504 "DataFusion may preload the UDTF name, then the function must resolve its table argument"
5505 );
5506 assert_eq!(
5507 catalog.get_view_count(),
5508 0,
5509 "registered table functions must not be resolved as REST catalog views"
5510 );
5511 }
5512
5513 #[tokio::test]
5514 async fn register_catalog_default_wrapper_uses_default_db() {
5515 let catalog = Arc::new(ProbeTrackingCatalog::new());
5518 let mut ctx = SQLContext::new();
5519 assert!(ctx
5520 .register_catalog("paimon", catalog.clone())
5521 .await
5522 .is_err());
5523 assert_eq!(catalog.get_count(), 1);
5524 }
5525
5526 fn assert_sql_type_to_paimon(
5527 sql_type: datafusion::sql::sqlparser::ast::DataType,
5528 expected: PaimonDataType,
5529 ) {
5530 assert_eq!(
5531 sql_data_type_to_paimon_type(&sql_type, true).unwrap(),
5532 expected
5533 );
5534 }
5535
5536 #[test]
5539 fn test_sql_type_boolean() {
5540 use datafusion::sql::sqlparser::ast::DataType as SqlType;
5541 assert_sql_type_to_paimon(
5542 SqlType::Boolean,
5543 PaimonDataType::Boolean(BooleanType::new()),
5544 );
5545 }
5546
5547 #[test]
5548 fn test_sql_type_integers() {
5549 use datafusion::sql::sqlparser::ast::DataType as SqlType;
5550 assert_sql_type_to_paimon(
5551 SqlType::TinyInt(None),
5552 PaimonDataType::TinyInt(TinyIntType::new()),
5553 );
5554 assert_sql_type_to_paimon(
5555 SqlType::SmallInt(None),
5556 PaimonDataType::SmallInt(SmallIntType::new()),
5557 );
5558 assert_sql_type_to_paimon(SqlType::Int(None), PaimonDataType::Int(IntType::new()));
5559 assert_sql_type_to_paimon(SqlType::Integer(None), PaimonDataType::Int(IntType::new()));
5560 assert_sql_type_to_paimon(
5561 SqlType::BigInt(None),
5562 PaimonDataType::BigInt(BigIntType::new()),
5563 );
5564 }
5565
5566 #[test]
5567 fn test_sql_type_floats() {
5568 use datafusion::sql::sqlparser::ast::{DataType as SqlType, ExactNumberInfo};
5569 assert_sql_type_to_paimon(
5570 SqlType::Float(ExactNumberInfo::None),
5571 PaimonDataType::Float(FloatType::new()),
5572 );
5573 assert_sql_type_to_paimon(SqlType::Real, PaimonDataType::Float(FloatType::new()));
5574 assert_sql_type_to_paimon(
5575 SqlType::DoublePrecision,
5576 PaimonDataType::Double(DoubleType::new()),
5577 );
5578 }
5579
5580 #[test]
5581 fn test_sql_type_string_variants() {
5582 use datafusion::sql::sqlparser::ast::{CharacterLength, DataType as SqlType};
5583 for sql_type in [SqlType::Varchar(None), SqlType::Text, SqlType::String(None)] {
5584 assert_sql_type_to_paimon(
5585 sql_type.clone(),
5586 PaimonDataType::VarChar(
5587 VarCharType::with_nullable(true, VarCharType::MAX_LENGTH).unwrap(),
5588 ),
5589 );
5590 }
5591 assert_sql_type_to_paimon(
5592 SqlType::Char(Some(CharacterLength::IntegerLength {
5593 length: 7,
5594 unit: None,
5595 })),
5596 PaimonDataType::Char(CharType::with_nullable(true, 7).unwrap()),
5597 );
5598 assert_sql_type_to_paimon(
5599 SqlType::Varchar(Some(CharacterLength::IntegerLength {
5600 length: 42,
5601 unit: None,
5602 })),
5603 PaimonDataType::VarChar(VarCharType::with_nullable(true, 42).unwrap()),
5604 );
5605 }
5606
5607 #[test]
5608 fn test_sql_type_binary() {
5609 use datafusion::sql::sqlparser::ast::{BinaryLength, DataType as SqlType};
5610 assert_sql_type_to_paimon(
5611 SqlType::Bytea,
5612 PaimonDataType::VarBinary(
5613 VarBinaryType::try_new(true, VarBinaryType::MAX_LENGTH).unwrap(),
5614 ),
5615 );
5616 assert_sql_type_to_paimon(
5617 SqlType::Binary(Some(8)),
5618 PaimonDataType::Binary(BinaryType::with_nullable(true, 8).unwrap()),
5619 );
5620 assert_sql_type_to_paimon(
5621 SqlType::Varbinary(Some(BinaryLength::IntegerLength { length: 32 })),
5622 PaimonDataType::VarBinary(VarBinaryType::try_new(true, 32).unwrap()),
5623 );
5624 }
5625
5626 #[test]
5627 fn test_sql_type_variant() {
5628 use datafusion::sql::sqlparser::ast::{DataType as SqlType, Ident, ObjectName};
5629 assert_sql_type_to_paimon(
5630 SqlType::Custom(ObjectName::from(Ident::new("VARIANT")), vec![]),
5631 PaimonDataType::Variant(VariantType::new()),
5632 );
5633 }
5634
5635 #[test]
5636 fn test_sql_type_date() {
5637 use datafusion::sql::sqlparser::ast::DataType as SqlType;
5638 assert_sql_type_to_paimon(SqlType::Date, PaimonDataType::Date(DateType::new()));
5639 }
5640
5641 #[test]
5642 fn test_sql_type_timestamp_default() {
5643 use datafusion::sql::sqlparser::ast::{DataType as SqlType, TimezoneInfo};
5644 assert_sql_type_to_paimon(
5645 SqlType::Timestamp(None, TimezoneInfo::None),
5646 PaimonDataType::Timestamp(TimestampType::with_nullable(true, 3).unwrap()),
5647 );
5648 }
5649
5650 #[test]
5651 fn test_sql_type_timestamp_with_precision() {
5652 use datafusion::sql::sqlparser::ast::{DataType as SqlType, TimezoneInfo};
5653 assert_sql_type_to_paimon(
5654 SqlType::Timestamp(Some(0), TimezoneInfo::None),
5655 PaimonDataType::Timestamp(TimestampType::with_nullable(true, 0).unwrap()),
5656 );
5657 assert_sql_type_to_paimon(
5658 SqlType::Timestamp(Some(3), TimezoneInfo::None),
5659 PaimonDataType::Timestamp(TimestampType::with_nullable(true, 3).unwrap()),
5660 );
5661 assert_sql_type_to_paimon(
5662 SqlType::Timestamp(Some(6), TimezoneInfo::None),
5663 PaimonDataType::Timestamp(TimestampType::with_nullable(true, 6).unwrap()),
5664 );
5665 assert_sql_type_to_paimon(
5666 SqlType::Timestamp(Some(9), TimezoneInfo::None),
5667 PaimonDataType::Timestamp(TimestampType::with_nullable(true, 9).unwrap()),
5668 );
5669 }
5670
5671 #[test]
5672 fn test_sql_type_timestamp_with_tz() {
5673 use datafusion::sql::sqlparser::ast::{DataType as SqlType, TimezoneInfo};
5674 assert_sql_type_to_paimon(
5675 SqlType::Timestamp(None, TimezoneInfo::WithTimeZone),
5676 PaimonDataType::LocalZonedTimestamp(
5677 LocalZonedTimestampType::with_nullable(true, 3).unwrap(),
5678 ),
5679 );
5680 }
5681
5682 #[test]
5683 fn test_sql_type_decimal() {
5684 use datafusion::sql::sqlparser::ast::{DataType as SqlType, ExactNumberInfo};
5685 assert_sql_type_to_paimon(
5686 SqlType::Decimal(ExactNumberInfo::PrecisionAndScale(18, 2)),
5687 PaimonDataType::Decimal(DecimalType::with_nullable(true, 18, 2).unwrap()),
5688 );
5689 assert_sql_type_to_paimon(
5690 SqlType::Decimal(ExactNumberInfo::Precision(10)),
5691 PaimonDataType::Decimal(DecimalType::with_nullable(true, 10, 0).unwrap()),
5692 );
5693 assert_sql_type_to_paimon(
5694 SqlType::Decimal(ExactNumberInfo::None),
5695 PaimonDataType::Decimal(DecimalType::with_nullable(true, 10, 0).unwrap()),
5696 );
5697 }
5698
5699 #[test]
5700 fn test_sql_type_unsupported() {
5701 use datafusion::sql::sqlparser::ast::DataType as SqlType;
5702 assert!(sql_data_type_to_paimon_type(&SqlType::Regclass, true).is_err());
5703 }
5704
5705 #[test]
5706 fn test_sql_type_array() {
5707 use datafusion::sql::sqlparser::ast::{ArrayElemTypeDef, DataType as SqlType};
5708 assert_sql_type_to_paimon(
5709 SqlType::Array(ArrayElemTypeDef::AngleBracket(Box::new(SqlType::Int(None)))),
5710 PaimonDataType::Array(PaimonArrayType::with_nullable(
5711 true,
5712 PaimonDataType::Int(IntType::new()),
5713 )),
5714 );
5715 }
5716
5717 #[test]
5718 fn test_sql_type_array_no_element() {
5719 use datafusion::sql::sqlparser::ast::{ArrayElemTypeDef, DataType as SqlType};
5720 assert!(
5721 sql_data_type_to_paimon_type(&SqlType::Array(ArrayElemTypeDef::None), true).is_err()
5722 );
5723 }
5724
5725 #[test]
5726 fn test_sql_type_map() {
5727 use datafusion::sql::sqlparser::ast::DataType as SqlType;
5728 assert_sql_type_to_paimon(
5729 SqlType::Map(
5730 Box::new(SqlType::Varchar(None)),
5731 Box::new(SqlType::Int(None)),
5732 ),
5733 PaimonDataType::Map(PaimonMapType::with_nullable(
5734 true,
5735 PaimonDataType::VarChar(
5736 VarCharType::with_nullable(false, VarCharType::MAX_LENGTH).unwrap(),
5737 ),
5738 PaimonDataType::Int(IntType::new()),
5739 )),
5740 );
5741 }
5742
5743 #[test]
5744 fn test_sql_type_struct() {
5745 use datafusion::sql::sqlparser::ast::{
5746 DataType as SqlType, Ident, StructBracketKind, StructField,
5747 };
5748 assert_sql_type_to_paimon(
5749 SqlType::Struct(
5750 vec![
5751 StructField {
5752 field_name: Some(Ident::new("name")),
5753 field_type: SqlType::Varchar(None),
5754 options: None,
5755 },
5756 StructField {
5757 field_name: Some(Ident::new("age")),
5758 field_type: SqlType::Int(None),
5759 options: None,
5760 },
5761 ],
5762 StructBracketKind::AngleBrackets,
5763 ),
5764 PaimonDataType::Row(PaimonRowType::with_nullable(
5765 true,
5766 vec![
5767 PaimonDataField::new(
5768 0,
5769 "name".to_string(),
5770 PaimonDataType::VarChar(
5771 VarCharType::with_nullable(true, VarCharType::MAX_LENGTH).unwrap(),
5772 ),
5773 ),
5774 PaimonDataField::new(1, "age".to_string(), PaimonDataType::Int(IntType::new())),
5775 ],
5776 )),
5777 );
5778 }
5779
5780 #[tokio::test]
5783 async fn test_resolve_three_part_name() {
5784 let catalog = Arc::new(MockCatalog::new());
5785 let sql_context = make_sql_context(catalog).await;
5786 let dialect = GenericDialect {};
5787 let stmts = Parser::parse_sql(&dialect, "SELECT * FROM paimon.mydb.mytable").unwrap();
5788 if let Statement::Query(q) = &stmts[0] {
5789 if let datafusion::sql::sqlparser::ast::SetExpr::Select(sel) = q.body.as_ref() {
5790 if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
5791 &sel.from[0].relation
5792 {
5793 let id = sql_context.resolve_table_name(name).unwrap();
5794 assert_eq!(id.database(), "mydb");
5795 assert_eq!(id.object(), "mytable");
5796 }
5797 }
5798 }
5799 }
5800
5801 #[tokio::test]
5802 async fn test_resolve_two_part_name() {
5803 let catalog = Arc::new(MockCatalog::new());
5804 let sql_context = make_sql_context(catalog).await;
5805 let dialect = GenericDialect {};
5806 let stmts = Parser::parse_sql(&dialect, "SELECT * FROM mydb.mytable").unwrap();
5807 if let Statement::Query(q) = &stmts[0] {
5808 if let datafusion::sql::sqlparser::ast::SetExpr::Select(sel) = q.body.as_ref() {
5809 if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
5810 &sel.from[0].relation
5811 {
5812 let id = sql_context.resolve_table_name(name).unwrap();
5813 assert_eq!(id.database(), "mydb");
5814 assert_eq!(id.object(), "mytable");
5815 }
5816 }
5817 }
5818 }
5819
5820 #[tokio::test]
5821 async fn test_resolve_wrong_catalog_name() {
5822 let catalog = Arc::new(MockCatalog::new());
5823 let sql_context = make_sql_context(catalog).await;
5824 let dialect = GenericDialect {};
5825 let stmts = Parser::parse_sql(&dialect, "SELECT * FROM other.mydb.mytable").unwrap();
5826 if let Statement::Query(q) = &stmts[0] {
5827 if let datafusion::sql::sqlparser::ast::SetExpr::Select(sel) = q.body.as_ref() {
5828 if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
5829 &sel.from[0].relation
5830 {
5831 let err = sql_context.resolve_table_name(name).unwrap_err();
5832 assert!(err.to_string().contains("Unknown catalog"));
5833 }
5834 }
5835 }
5836 }
5837
5838 #[tokio::test]
5839 async fn test_resolve_single_part_name_uses_default_schema() {
5840 let catalog = Arc::new(MockCatalog::new());
5841 let sql_context = make_sql_context(catalog).await;
5842 let dialect = GenericDialect {};
5843 let stmts = Parser::parse_sql(&dialect, "SELECT * FROM mytable").unwrap();
5844 if let Statement::Query(q) = &stmts[0] {
5845 if let datafusion::sql::sqlparser::ast::SetExpr::Select(sel) = q.body.as_ref() {
5846 if let datafusion::sql::sqlparser::ast::TableFactor::Table { name, .. } =
5847 &sel.from[0].relation
5848 {
5849 let id = sql_context.resolve_table_name(name).unwrap();
5850 assert_eq!(id.database(), "default");
5851 assert_eq!(id.object(), "mytable");
5852 }
5853 }
5854 }
5855 }
5856
5857 #[test]
5860 fn test_extract_options_none() {
5861 let opts = extract_options(&CreateTableOptions::None).unwrap();
5862 assert!(opts.is_empty());
5863 }
5864
5865 #[test]
5866 fn test_extract_options_with_kv() {
5867 let dialect = GenericDialect {};
5869 let stmts =
5870 Parser::parse_sql(&dialect, "CREATE TABLE t (id INT) WITH ('bucket' = '4')").unwrap();
5871 if let Statement::CreateTable(ct) = &stmts[0] {
5872 let opts = extract_options(&ct.table_options).unwrap();
5873 assert_eq!(opts.len(), 1);
5874 assert_eq!(opts[0].0, "bucket");
5875 assert_eq!(opts[0].1, "4");
5876 } else {
5877 panic!("expected CreateTable");
5878 }
5879 }
5880
5881 #[tokio::test]
5884 async fn test_create_table_basic() {
5885 let catalog = Arc::new(MockCatalog::new());
5886 let sql_context = make_sql_context(catalog.clone()).await;
5887
5888 sql_context
5889 .sql("CREATE TABLE mydb.t1 (id INT NOT NULL, name VARCHAR, PRIMARY KEY (id))")
5890 .await
5891 .unwrap();
5892
5893 let calls = catalog.take_calls();
5894 assert_eq!(calls.len(), 1);
5895 if let CatalogCall::CreateTable {
5896 identifier,
5897 schema,
5898 ignore_if_exists,
5899 } = &calls[0]
5900 {
5901 assert_eq!(identifier.database(), "mydb");
5902 assert_eq!(identifier.object(), "t1");
5903 assert!(!ignore_if_exists);
5904 assert_eq!(schema.primary_keys(), &["id"]);
5905 } else {
5906 panic!("expected CreateTable call");
5907 }
5908 }
5909
5910 #[tokio::test]
5911 async fn test_create_table_if_not_exists() {
5912 let catalog = Arc::new(MockCatalog::new());
5913 let sql_context = make_sql_context(catalog.clone()).await;
5914
5915 sql_context
5916 .sql("CREATE TABLE IF NOT EXISTS mydb.t1 (id INT)")
5917 .await
5918 .unwrap();
5919
5920 let calls = catalog.take_calls();
5921 assert_eq!(calls.len(), 1);
5922 if let CatalogCall::CreateTable {
5923 ignore_if_exists, ..
5924 } = &calls[0]
5925 {
5926 assert!(ignore_if_exists);
5927 } else {
5928 panic!("expected CreateTable call");
5929 }
5930 }
5931
5932 #[tokio::test]
5933 async fn test_create_table_with_options() {
5934 let catalog = Arc::new(MockCatalog::new());
5935 let sql_context = make_sql_context(catalog.clone()).await;
5936
5937 sql_context
5938 .sql("CREATE TABLE mydb.t1 (id INT) WITH ('bucket' = '4', 'file.format' = 'parquet')")
5939 .await
5940 .unwrap();
5941
5942 let calls = catalog.take_calls();
5943 assert_eq!(calls.len(), 1);
5944 if let CatalogCall::CreateTable { schema, .. } = &calls[0] {
5945 let opts = schema.options();
5946 assert_eq!(opts.get("bucket").unwrap(), "4");
5947 assert_eq!(opts.get("file.format").unwrap(), "parquet");
5948 } else {
5949 panic!("expected CreateTable call");
5950 }
5951 }
5952
5953 #[tokio::test]
5954 async fn test_create_table_three_part_name() {
5955 let catalog = Arc::new(MockCatalog::new());
5956 let sql_context = make_sql_context(catalog.clone()).await;
5957
5958 sql_context
5959 .sql("CREATE TABLE paimon.mydb.t1 (id INT)")
5960 .await
5961 .unwrap();
5962
5963 let calls = catalog.take_calls();
5964 if let CatalogCall::CreateTable { identifier, .. } = &calls[0] {
5965 assert_eq!(identifier.database(), "mydb");
5966 assert_eq!(identifier.object(), "t1");
5967 } else {
5968 panic!("expected CreateTable call");
5969 }
5970 }
5971
5972 #[tokio::test]
5973 async fn test_create_table_blob_type_preserved() {
5974 let catalog = Arc::new(MockCatalog::new());
5975 let sql_context = make_sql_context(catalog.clone()).await;
5976
5977 sql_context
5978 .sql("CREATE TABLE mydb.t1 (id INT, payload BLOB NOT NULL) WITH ('data-evolution.enabled' = 'true')")
5979 .await
5980 .unwrap();
5981
5982 let calls = catalog.take_calls();
5983 assert_eq!(calls.len(), 1);
5984 if let CatalogCall::CreateTable { schema, .. } = &calls[0] {
5985 assert_eq!(schema.fields().len(), 2);
5986 assert!(matches!(
5987 schema.fields()[1].data_type(),
5988 PaimonDataType::Blob(_)
5989 ));
5990 assert!(!schema.fields()[1].data_type().is_nullable());
5991 } else {
5992 panic!("expected CreateTable call");
5993 }
5994 }
5995
5996 #[tokio::test]
5997 async fn test_create_table_blob_comment_directives() {
5998 let catalog = Arc::new(MockCatalog::new());
5999 let sql_context = make_sql_context(catalog.clone()).await;
6000
6001 sql_context
6002 .sql(
6003 "CREATE TABLE mydb.t1 (\
6004 id INT, \
6005 photo BYTES COMMENT '__BLOB_FIELD; raw photo', \
6006 thumb BINARY COMMENT '__BLOB_DESCRIPTOR_FIELD', \
6007 preview VARBINARY COMMENT '__BLOB_VIEW_FIELD; preview ref'\
6008 ) WITH ('data-evolution.enabled' = 'true')",
6009 )
6010 .await
6011 .unwrap();
6012
6013 let calls = catalog.take_calls();
6014 assert_eq!(calls.len(), 1);
6015 if let CatalogCall::CreateTable { schema, .. } = &calls[0] {
6016 assert!(matches!(
6017 schema.fields()[1].data_type(),
6018 PaimonDataType::Blob(_)
6019 ));
6020 assert!(matches!(
6021 schema.fields()[2].data_type(),
6022 PaimonDataType::Blob(_)
6023 ));
6024 assert!(matches!(
6025 schema.fields()[3].data_type(),
6026 PaimonDataType::Blob(_)
6027 ));
6028 assert_eq!(schema.fields()[1].description(), Some("raw photo"));
6029 assert_eq!(schema.fields()[2].description(), None);
6030 assert_eq!(schema.fields()[3].description(), Some("preview ref"));
6031 assert_eq!(
6032 schema.options().get("blob-field").map(String::as_str),
6033 Some("photo")
6034 );
6035 assert_eq!(
6036 schema
6037 .options()
6038 .get("blob-descriptor-field")
6039 .map(String::as_str),
6040 Some("thumb")
6041 );
6042 assert_eq!(
6043 schema.options().get("blob-view-field").map(String::as_str),
6044 Some("preview")
6045 );
6046 } else {
6047 panic!("expected CreateTable call");
6048 }
6049 }
6050
6051 #[tokio::test]
6052 async fn test_alter_table_add_column() {
6053 let catalog = Arc::new(MockCatalog::new());
6054 let sql_context = make_sql_context(catalog.clone()).await;
6055
6056 sql_context
6057 .sql("ALTER TABLE mydb.t1 ADD COLUMN age INT")
6058 .await
6059 .unwrap();
6060
6061 let calls = catalog.take_calls();
6062 assert_eq!(calls.len(), 1);
6063 if let CatalogCall::AlterTable {
6064 identifier,
6065 changes,
6066 ..
6067 } = &calls[0]
6068 {
6069 assert_eq!(identifier.database(), "mydb");
6070 assert_eq!(identifier.object(), "t1");
6071 assert_eq!(changes.len(), 1);
6072 assert!(
6073 matches!(&changes[0], SchemaChange::AddColumn { field_names, .. } if field_names.first().map(String::as_str) == Some("age"))
6074 );
6075 } else {
6076 panic!("expected AlterTable call");
6077 }
6078 }
6079
6080 #[tokio::test]
6081 async fn test_alter_table_add_blob_column() {
6082 let catalog = Arc::new(MockCatalog::new());
6083 let sql_context = make_sql_context(catalog.clone()).await;
6084
6085 sql_context
6086 .sql("ALTER TABLE mydb.t1 ADD COLUMN payload BLOB")
6087 .await
6088 .unwrap();
6089
6090 let calls = catalog.take_calls();
6091 assert_eq!(calls.len(), 1);
6092 if let CatalogCall::AlterTable { changes, .. } = &calls[0] {
6093 assert_eq!(changes.len(), 1);
6094 assert!(matches!(
6095 &changes[0],
6096 SchemaChange::AddColumn {
6097 field_names,
6098 data_type,
6099 ..
6100 } if field_names.first().map(String::as_str) == Some("payload") && matches!(data_type, PaimonDataType::Blob(_))
6101 ));
6102 } else {
6103 panic!("expected AlterTable call");
6104 }
6105 }
6106
6107 #[tokio::test]
6108 async fn test_alter_table_add_blob_comment_directive_passes_core_input() {
6109 let catalog = Arc::new(MockCatalog::new());
6110 let sql_context = make_sql_context(catalog.clone()).await;
6111
6112 sql_context
6113 .sql("ALTER TABLE mydb.t1 ADD COLUMN preview BYTES COMMENT '__BLOB_DESCRIPTOR_FIELD; preview descriptor'")
6114 .await
6115 .unwrap();
6116
6117 let calls = catalog.take_calls();
6118 assert_eq!(calls.len(), 1);
6119 if let CatalogCall::AlterTable { changes, .. } = &calls[0] {
6120 assert_eq!(changes.len(), 1);
6121 assert!(matches!(
6122 &changes[0],
6123 SchemaChange::AddColumn {
6124 field_names,
6125 data_type,
6126 comment,
6127 ..
6128 } if field_names.first().map(String::as_str) == Some("preview")
6129 && matches!(data_type, PaimonDataType::VarBinary(_))
6130 && comment.as_deref() == Some("__BLOB_DESCRIPTOR_FIELD; preview descriptor")
6131 ));
6132 } else {
6133 panic!("expected AlterTable call");
6134 }
6135 }
6136
6137 #[tokio::test]
6138 async fn test_alter_table_drop_column() {
6139 let catalog = Arc::new(MockCatalog::new());
6140 let sql_context = make_sql_context(catalog.clone()).await;
6141
6142 sql_context
6143 .sql("ALTER TABLE mydb.t1 DROP COLUMN age")
6144 .await
6145 .unwrap();
6146
6147 let calls = catalog.take_calls();
6148 assert_eq!(calls.len(), 1);
6149 if let CatalogCall::AlterTable { changes, .. } = &calls[0] {
6150 assert_eq!(changes.len(), 1);
6151 assert!(
6152 matches!(&changes[0], SchemaChange::DropColumn { field_names } if field_names.first().map(String::as_str) == Some("age"))
6153 );
6154 } else {
6155 panic!("expected AlterTable call");
6156 }
6157 }
6158
6159 #[tokio::test]
6160 async fn test_alter_table_rename_column() {
6161 let catalog = Arc::new(MockCatalog::new());
6162 let sql_context = make_sql_context(catalog.clone()).await;
6163
6164 sql_context
6165 .sql("ALTER TABLE mydb.t1 RENAME COLUMN old_name TO new_name")
6166 .await
6167 .unwrap();
6168
6169 let calls = catalog.take_calls();
6170 assert_eq!(calls.len(), 1);
6171 if let CatalogCall::AlterTable { changes, .. } = &calls[0] {
6172 assert_eq!(changes.len(), 1);
6173 assert!(matches!(
6174 &changes[0],
6175 SchemaChange::RenameColumn { field_names, new_name }
6176 if field_names.first().map(String::as_str) == Some("old_name") && new_name == "new_name"
6177 ));
6178 } else {
6179 panic!("expected AlterTable call");
6180 }
6181 }
6182
6183 #[tokio::test]
6184 async fn test_alter_table_rename_table() {
6185 let catalog = Arc::new(MockCatalog::new());
6186 let sql_context = make_sql_context(catalog.clone()).await;
6187
6188 sql_context
6189 .sql("ALTER TABLE mydb.t1 RENAME TO t2")
6190 .await
6191 .unwrap();
6192
6193 let calls = catalog.take_calls();
6194 assert_eq!(calls.len(), 1);
6195 if let CatalogCall::RenameTable { from, to, .. } = &calls[0] {
6196 assert_eq!(from.database(), "mydb");
6197 assert_eq!(from.object(), "t1");
6198 assert_eq!(to.database(), "mydb");
6199 assert_eq!(to.object(), "t2");
6200 } else {
6201 panic!("expected RenameTable call");
6202 }
6203 }
6204
6205 #[tokio::test]
6206 async fn test_alter_table_if_exists_add_column() {
6207 let catalog = Arc::new(MockCatalog::new());
6208 let sql_context = make_sql_context(catalog.clone()).await;
6209
6210 sql_context
6211 .sql("ALTER TABLE IF EXISTS mydb.t1 ADD COLUMN age INT")
6212 .await
6213 .unwrap();
6214
6215 let calls = catalog.take_calls();
6216 assert_eq!(calls.len(), 1);
6217 if let CatalogCall::AlterTable {
6218 ignore_if_not_exists,
6219 ..
6220 } = &calls[0]
6221 {
6222 assert!(ignore_if_not_exists);
6223 } else {
6224 panic!("expected AlterTable call");
6225 }
6226 }
6227
6228 #[tokio::test]
6229 async fn test_alter_table_without_if_exists() {
6230 let catalog = Arc::new(MockCatalog::new());
6231 let sql_context = make_sql_context(catalog.clone()).await;
6232
6233 sql_context
6234 .sql("ALTER TABLE mydb.t1 ADD COLUMN age INT")
6235 .await
6236 .unwrap();
6237
6238 let calls = catalog.take_calls();
6239 if let CatalogCall::AlterTable {
6240 ignore_if_not_exists,
6241 ..
6242 } = &calls[0]
6243 {
6244 assert!(!ignore_if_not_exists);
6245 } else {
6246 panic!("expected AlterTable call");
6247 }
6248 }
6249
6250 #[tokio::test]
6251 async fn test_alter_table_if_exists_rename() {
6252 let catalog = Arc::new(MockCatalog::new());
6253 let sql_context = make_sql_context(catalog.clone()).await;
6254
6255 sql_context
6256 .sql("ALTER TABLE IF EXISTS mydb.t1 RENAME TO t2")
6257 .await
6258 .unwrap();
6259
6260 let calls = catalog.take_calls();
6261 assert_eq!(calls.len(), 1);
6262 if let CatalogCall::RenameTable {
6263 from,
6264 to,
6265 ignore_if_not_exists,
6266 } = &calls[0]
6267 {
6268 assert!(ignore_if_not_exists);
6269 assert_eq!(from.object(), "t1");
6270 assert_eq!(to.object(), "t2");
6271 } else {
6272 panic!("expected RenameTable call");
6273 }
6274 }
6275
6276 #[tokio::test]
6277 async fn test_alter_table_rename_three_part_name() {
6278 let catalog = Arc::new(MockCatalog::new());
6279 let sql_context = make_sql_context(catalog.clone()).await;
6280
6281 sql_context
6282 .sql("ALTER TABLE paimon.mydb.t1 RENAME TO t2")
6283 .await
6284 .unwrap();
6285
6286 let calls = catalog.take_calls();
6287 assert_eq!(calls.len(), 1);
6288 if let CatalogCall::RenameTable { from, to, .. } = &calls[0] {
6289 assert_eq!(from.database(), "mydb");
6290 assert_eq!(from.object(), "t1");
6291 assert_eq!(to.database(), "mydb");
6292 assert_eq!(to.object(), "t2");
6293 } else {
6294 panic!("expected RenameTable call");
6295 }
6296 }
6297
6298 #[tokio::test]
6299 async fn test_sql_parse_error() {
6300 let catalog = Arc::new(MockCatalog::new());
6301 let sql_context = make_sql_context(catalog).await;
6302 let result = sql_context.sql("NOT VALID SQL !!!").await;
6303 assert!(result.is_err());
6304 assert!(result.unwrap_err().to_string().contains("SQL parse error"));
6305 }
6306
6307 #[tokio::test]
6308 async fn test_multiple_statements_error() {
6309 let catalog = Arc::new(MockCatalog::new());
6310 let sql_context = make_sql_context(catalog).await;
6311 let result = sql_context.sql("SELECT 1; SELECT 2").await;
6312 assert!(result.is_err());
6313 assert!(result
6314 .unwrap_err()
6315 .to_string()
6316 .contains("exactly one SQL statement"));
6317 }
6318
6319 #[tokio::test]
6320 async fn test_create_external_table_rejected() {
6321 let catalog = Arc::new(MockCatalog::new());
6322 let sql_context = make_sql_context(catalog).await;
6323 let result = sql_context
6324 .sql("CREATE EXTERNAL TABLE mydb.t1 (id INT) STORED AS PARQUET")
6325 .await;
6326 assert!(result.is_err());
6327 assert!(result
6328 .unwrap_err()
6329 .to_string()
6330 .contains("CREATE EXTERNAL TABLE is not supported"));
6331 }
6332
6333 #[tokio::test]
6334 async fn test_non_ddl_delegates_to_datafusion() {
6335 let catalog = Arc::new(MockCatalog::new());
6336 let sql_context = make_sql_context(catalog.clone()).await;
6337 let df = sql_context.sql("SELECT 1 AS x").await.unwrap();
6339 let batches = df.collect().await.unwrap();
6340 assert_eq!(batches.len(), 1);
6341 assert_eq!(batches[0].num_rows(), 1);
6342 assert!(catalog.take_calls().is_empty());
6344 }
6345
6346 #[test]
6349 fn test_extract_partition_by_no_clause() {
6350 let (rewritten, keys) = extract_partition_by("CREATE TABLE t (id INT)").unwrap();
6351 assert_eq!(rewritten, "CREATE TABLE t (id INT)");
6352 assert!(keys.is_empty());
6353 }
6354
6355 #[test]
6356 fn test_extract_partition_by_single_column() {
6357 let (rewritten, keys) = extract_partition_by(
6358 "CREATE TABLE t (id INT, dt STRING) PARTITIONED BY (dt) WITH ('k'='v')",
6359 )
6360 .unwrap();
6361 assert_eq!(keys, vec!["dt"]);
6362 assert!(!rewritten.contains("PARTITIONED"));
6363 assert!(rewritten.contains("WITH"));
6364 }
6365
6366 #[test]
6367 fn test_extract_partition_by_multiple_columns() {
6368 let (_, keys) =
6369 extract_partition_by("CREATE TABLE t (a INT, b INT, c INT) PARTITIONED BY (a, b)")
6370 .unwrap();
6371 assert_eq!(keys, vec!["a", "b"]);
6372 }
6373
6374 #[test]
6375 fn test_extract_partition_by_mixed_case() {
6376 let (_, keys) =
6377 extract_partition_by("CREATE TABLE t (dt INT) Partitioned by (dt)").unwrap();
6378 assert_eq!(keys, vec!["dt"]);
6379 }
6380
6381 #[test]
6382 fn test_extract_partition_by_rejects_typed_column() {
6383 let err = extract_partition_by("CREATE TABLE t (dt STRING) PARTITIONED BY (dt STRING)")
6384 .unwrap_err();
6385 assert!(err.to_string().contains("should not specify a type"));
6386 }
6387
6388 #[test]
6389 fn test_extract_partition_by_empty_parens() {
6390 let err = extract_partition_by("CREATE TABLE t (id INT) PARTITIONED BY ()").unwrap_err();
6391 assert!(err.to_string().contains("at least one column"));
6392 }
6393
6394 #[test]
6395 fn test_extract_partition_by_unmatched_paren() {
6396 let err = extract_partition_by("CREATE TABLE t (id INT) PARTITIONED BY (dt").unwrap_err();
6397 assert!(err.to_string().contains("Unmatched"));
6398 }
6399
6400 #[test]
6401 fn test_extract_partition_by_skips_string_literal() {
6402 let sql =
6403 "CREATE TABLE t (id INT) WITH ('note' = 'PARTITIONED BY (x)') PARTITIONED BY (id)";
6404 let (rewritten, keys) = extract_partition_by(sql).unwrap();
6405 assert_eq!(keys, vec!["id"]);
6406 assert!(rewritten.contains("WITH"));
6407 assert!(rewritten.contains("'PARTITIONED BY (x)'"));
6408 }
6409
6410 #[test]
6411 fn test_extract_partition_by_skips_line_comment() {
6412 let sql = "CREATE TABLE t (id INT) -- PARTITIONED BY (x)\nPARTITIONED BY (id)";
6413 let (_, keys) = extract_partition_by(sql).unwrap();
6414 assert_eq!(keys, vec!["id"]);
6415 }
6416
6417 #[test]
6418 fn test_extract_partition_by_double_quoted_identifier() {
6419 let (_, keys) =
6420 extract_partition_by("CREATE TABLE t (\"order\" INT) PARTITIONED BY (\"order\")")
6421 .unwrap();
6422 assert_eq!(keys, vec!["order"]);
6423 }
6424
6425 #[test]
6426 fn test_extract_partition_by_double_quoted_identifier_with_escaped_quote_and_comma() {
6427 let (_, keys) = extract_partition_by(
6428 "CREATE TABLE t (\"a\"\"b,c\" INT, `d``e,f` INT) \
6429 PARTITIONED BY (\"a\"\"b,c\", `d``e,f`)",
6430 )
6431 .unwrap();
6432 assert_eq!(keys, vec!["a\"b,c", "d`e,f"]);
6433 }
6434
6435 #[test]
6436 fn test_extract_partition_by_backtick_quoted_identifier() {
6437 let (_, keys) =
6438 extract_partition_by("CREATE TABLE t (`order` INT) PARTITIONED BY (`order`)").unwrap();
6439 assert_eq!(keys, vec!["order"]);
6440 }
6441
6442 #[test]
6443 fn test_extract_partition_by_no_paren_after_by() {
6444 let err = extract_partition_by("CREATE TABLE t (id INT) PARTITIONED BY dt").unwrap_err();
6445 assert!(err.to_string().contains("Expected '('"));
6446 }
6447
6448 #[test]
6449 fn test_extract_partition_by_only_partitioned_no_by() {
6450 let (rewritten, keys) = extract_partition_by("CREATE TABLE partitioned (id INT)").unwrap();
6451 assert_eq!(rewritten, "CREATE TABLE partitioned (id INT)");
6452 assert!(keys.is_empty());
6453 }
6454
6455 #[test]
6456 fn test_extract_partition_by_skips_block_comment() {
6457 let sql = "CREATE TABLE t (id INT) /* PARTITIONED BY (x) */ PARTITIONED BY (id)";
6458 let (rewritten, keys) = extract_partition_by(sql).unwrap();
6459 assert_eq!(keys, vec!["id"]);
6460 assert!(rewritten.contains("/* PARTITIONED BY (x) */"));
6461 }
6462
6463 #[test]
6464 fn test_looks_like_create_table() {
6465 assert!(looks_like_create_table("CREATE TABLE t (id INT)"));
6466 assert!(looks_like_create_table(" create table t (id INT)"));
6467 assert!(looks_like_create_table(
6468 "CREATE TABLE IF NOT EXISTS t (id INT)",
6469 ));
6470 assert!(looks_like_create_table(
6471 "/* note */ CREATE TABLE t (id INT)",
6472 ));
6473 assert!(looks_like_create_table(
6474 "-- comment\nCREATE TABLE t (id INT)",
6475 ));
6476 assert!(looks_like_create_table(
6477 "/* a */ /* b */ CREATE TABLE t (id INT)",
6478 ));
6479 assert!(!looks_like_create_table("ALTER TABLE t ADD COLUMN x INT"));
6480 assert!(!looks_like_create_table("SELECT 1"));
6481 assert!(!looks_like_create_table(
6482 "SELECT aaaaaaaaaaaaaaaaaaaa中文 FROM t",
6483 ));
6484 }
6485
6486 #[tokio::test]
6489 async fn test_create_table_partition_key_not_in_columns() {
6490 let catalog = Arc::new(MockCatalog::new());
6491 let sql_context = make_sql_context(catalog).await;
6492 let err = sql_context
6493 .sql("CREATE TABLE mydb.t (id INT, dt STRING) PARTITIONED BY (nonexistent)")
6494 .await
6495 .unwrap_err();
6496 assert!(err.to_string().contains("is not defined in the table"));
6497 }
6498
6499 #[tokio::test]
6500 async fn test_create_table_partition_key_matches_column() {
6501 let catalog = Arc::new(MockCatalog::new());
6502 let sql_context = make_sql_context(catalog.clone()).await;
6503 sql_context
6504 .sql("CREATE TABLE mydb.t (id INT, dt STRING) PARTITIONED BY (dt)")
6505 .await
6506 .unwrap();
6507 let calls = catalog.take_calls();
6508 assert_eq!(calls.len(), 1);
6509 if let CatalogCall::CreateTable { schema, .. } = &calls[0] {
6510 assert_eq!(schema.partition_keys(), &["dt"]);
6511 } else {
6512 panic!("expected CreateTable call");
6513 }
6514 }
6515
6516 #[tokio::test]
6519 async fn test_set_paimon_option() {
6520 let catalog = Arc::new(MockCatalog::new());
6521 let sql_context = make_sql_context(catalog).await;
6522 sql_context
6523 .sql("SET 'paimon.scan.version' = '1'")
6524 .await
6525 .unwrap();
6526 let opts = sql_context.dynamic_options().read().unwrap();
6527 assert_eq!(opts.get("scan.version").unwrap(), "1");
6528 }
6529
6530 #[tokio::test]
6531 async fn test_set_paimon_option_overwrites() {
6532 let catalog = Arc::new(MockCatalog::new());
6533 let sql_context = make_sql_context(catalog).await;
6534 sql_context
6535 .sql("SET 'paimon.scan.version' = '1'")
6536 .await
6537 .unwrap();
6538 sql_context
6539 .sql("SET 'paimon.scan.version' = '2'")
6540 .await
6541 .unwrap();
6542 let opts = sql_context.dynamic_options().read().unwrap();
6543 assert_eq!(opts.get("scan.version").unwrap(), "2");
6544 }
6545
6546 #[tokio::test]
6547 async fn test_reset_paimon_option() {
6548 let catalog = Arc::new(MockCatalog::new());
6549 let sql_context = make_sql_context(catalog).await;
6550 sql_context
6551 .sql("SET 'paimon.scan.version' = '1'")
6552 .await
6553 .unwrap();
6554 sql_context
6555 .sql("RESET 'paimon.scan.version'")
6556 .await
6557 .unwrap();
6558 let opts = sql_context.dynamic_options().read().unwrap();
6559 assert!(opts.get("scan.version").is_none());
6560 }
6561
6562 #[tokio::test]
6563 async fn test_set_non_paimon_option_delegates() {
6564 let catalog = Arc::new(MockCatalog::new());
6565 let sql_context = make_sql_context(catalog).await;
6566 let _ = sql_context
6569 .sql("SET datafusion.optimizer.max_passes = 3")
6570 .await;
6571 let opts = sql_context.dynamic_options().read().unwrap();
6572 assert!(opts.is_empty());
6573 }
6574
6575 #[tokio::test]
6576 async fn test_set_multiple_paimon_options() {
6577 let catalog = Arc::new(MockCatalog::new());
6578 let sql_context = make_sql_context(catalog).await;
6579 sql_context
6580 .sql("SET 'paimon.scan.version' = '1'")
6581 .await
6582 .unwrap();
6583 sql_context
6584 .sql("SET 'paimon.scan.timestamp-millis' = '1000'")
6585 .await
6586 .unwrap();
6587 let opts = sql_context.dynamic_options().read().unwrap();
6588 assert_eq!(opts.len(), 2);
6589 assert_eq!(opts.get("scan.version").unwrap(), "1");
6590 assert_eq!(opts.get("scan.timestamp-millis").unwrap(), "1000");
6591 }
6592
6593 #[tokio::test]
6594 async fn test_reset_nonexistent_paimon_option_is_noop() {
6595 let catalog = Arc::new(MockCatalog::new());
6596 let sql_context = make_sql_context(catalog).await;
6597 sql_context
6598 .sql("RESET 'paimon.scan.version'")
6599 .await
6600 .unwrap();
6601 let opts = sql_context.dynamic_options().read().unwrap();
6602 assert!(opts.is_empty());
6603 }
6604
6605 async fn setup_fs_sql_context() -> (tempfile::TempDir, SQLContext) {
6608 use paimon::{CatalogOptions, FileSystemCatalog, Options};
6609
6610 let temp_dir = tempfile::TempDir::new().unwrap();
6611 let warehouse = format!("file://{}", temp_dir.path().display());
6612 let mut options = Options::new();
6613 options.set(CatalogOptions::WAREHOUSE, warehouse);
6614 let catalog = Arc::new(FileSystemCatalog::new(options).unwrap());
6615
6616 let mut sql_context = SQLContext::new();
6617 sql_context
6618 .register_catalog("paimon", catalog.clone())
6619 .await
6620 .unwrap();
6621 sql_context
6622 .sql("CREATE SCHEMA paimon.test_db")
6623 .await
6624 .unwrap();
6625
6626 (temp_dir, sql_context)
6627 }
6628
6629 #[tokio::test]
6630 async fn test_dynamic_read_batch_size_overrides_table_option() {
6631 let (_tmp, sql_context) = setup_fs_sql_context().await;
6632
6633 sql_context
6634 .sql(
6635 "CREATE TABLE paimon.test_db.batch_size_t (id INT) \
6636 WITH ('read.batch-size' = '3')",
6637 )
6638 .await
6639 .unwrap();
6640 sql_context
6641 .sql("INSERT INTO paimon.test_db.batch_size_t VALUES (1), (2), (3), (4), (5)")
6642 .await
6643 .unwrap()
6644 .collect()
6645 .await
6646 .unwrap();
6647
6648 sql_context
6649 .sql("SET 'paimon.read.batch-size' = '2'")
6650 .await
6651 .unwrap();
6652 let batches = sql_context
6653 .sql("SELECT id FROM paimon.test_db.batch_size_t")
6654 .await
6655 .unwrap()
6656 .collect()
6657 .await
6658 .unwrap();
6659 assert_eq!(
6660 batches
6661 .iter()
6662 .map(|batch| batch.num_rows())
6663 .collect::<Vec<_>>(),
6664 vec![2, 2, 1]
6665 );
6666
6667 sql_context
6668 .sql("RESET 'paimon.read.batch-size'")
6669 .await
6670 .unwrap();
6671 let batches = sql_context
6672 .sql("SELECT id FROM paimon.test_db.batch_size_t")
6673 .await
6674 .unwrap()
6675 .collect()
6676 .await
6677 .unwrap();
6678 assert_eq!(
6679 batches
6680 .iter()
6681 .map(|batch| batch.num_rows())
6682 .collect::<Vec<_>>(),
6683 vec![3, 2]
6684 );
6685 }
6686
6687 #[tokio::test]
6688 async fn test_truncate_table() {
6689 let (_tmp, sql_context) = setup_fs_sql_context().await;
6690
6691 sql_context
6692 .sql("CREATE TABLE paimon.test_db.t1 (id INT, value INT)")
6693 .await
6694 .unwrap();
6695 sql_context
6696 .sql("INSERT INTO paimon.test_db.t1 VALUES (1, 10), (2, 20)")
6697 .await
6698 .unwrap()
6699 .collect()
6700 .await
6701 .unwrap();
6702
6703 sql_context
6704 .sql("TRUNCATE TABLE paimon.test_db.t1")
6705 .await
6706 .unwrap();
6707
6708 let batches = sql_context
6709 .sql("SELECT * FROM paimon.test_db.t1")
6710 .await
6711 .unwrap()
6712 .collect()
6713 .await
6714 .unwrap();
6715 let total: usize = batches.iter().map(|b| b.num_rows()).sum();
6716 assert_eq!(total, 0);
6717 }
6718
6719 #[tokio::test]
6720 async fn test_truncate_table_partition() {
6721 let (_tmp, sql_context) = setup_fs_sql_context().await;
6722
6723 sql_context
6724 .sql("CREATE TABLE paimon.test_db.t2 (pt VARCHAR, id INT) PARTITIONED BY (pt)")
6725 .await
6726 .unwrap();
6727 sql_context
6728 .sql("INSERT INTO paimon.test_db.t2 VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
6729 .await
6730 .unwrap()
6731 .collect()
6732 .await
6733 .unwrap();
6734
6735 sql_context
6736 .sql("TRUNCATE TABLE paimon.test_db.t2 PARTITION (pt = 'a')")
6737 .await
6738 .unwrap();
6739
6740 let batches = sql_context
6741 .sql("SELECT pt, id FROM paimon.test_db.t2 ORDER BY id")
6742 .await
6743 .unwrap()
6744 .collect()
6745 .await
6746 .unwrap();
6747
6748 let mut rows = Vec::new();
6749 for batch in &batches {
6750 let pts = batch
6751 .column(0)
6752 .as_any()
6753 .downcast_ref::<StringViewArray>()
6754 .unwrap();
6755 let ids = batch
6756 .column(1)
6757 .as_any()
6758 .downcast_ref::<Int32Array>()
6759 .unwrap();
6760 for i in 0..batch.num_rows() {
6761 rows.push((pts.value(i).to_string(), ids.value(i)));
6762 }
6763 }
6764 assert_eq!(rows, vec![("b".to_string(), 3), ("b".to_string(), 4)]);
6765 }
6766
6767 #[tokio::test]
6768 async fn test_alter_table_drop_partitions() {
6769 let (_tmp, sql_context) = setup_fs_sql_context().await;
6770
6771 sql_context
6772 .sql("CREATE TABLE paimon.test_db.t3 (pt VARCHAR, id INT) PARTITIONED BY (pt)")
6773 .await
6774 .unwrap();
6775 sql_context
6776 .sql("INSERT INTO paimon.test_db.t3 VALUES ('a', 1), ('a', 2), ('b', 3), ('b', 4)")
6777 .await
6778 .unwrap()
6779 .collect()
6780 .await
6781 .unwrap();
6782
6783 sql_context
6784 .sql("ALTER TABLE paimon.test_db.t3 DROP PARTITION (pt = 'b')")
6785 .await
6786 .unwrap();
6787
6788 let batches = sql_context
6789 .sql("SELECT pt, id FROM paimon.test_db.t3 ORDER BY id")
6790 .await
6791 .unwrap()
6792 .collect()
6793 .await
6794 .unwrap();
6795
6796 let mut rows = Vec::new();
6797 for batch in &batches {
6798 let pts = batch
6799 .column(0)
6800 .as_any()
6801 .downcast_ref::<StringViewArray>()
6802 .unwrap();
6803 let ids = batch
6804 .column(1)
6805 .as_any()
6806 .downcast_ref::<Int32Array>()
6807 .unwrap();
6808 for i in 0..batch.num_rows() {
6809 rows.push((pts.value(i).to_string(), ids.value(i)));
6810 }
6811 }
6812 assert_eq!(rows, vec![("a".to_string(), 1), ("a".to_string(), 2)]);
6813 }
6814
6815 #[tokio::test]
6816 async fn test_truncate_table_incomplete_partition_spec() {
6817 let (_tmp, sql_context) = setup_fs_sql_context().await;
6818
6819 sql_context
6820 .sql("CREATE TABLE paimon.test_db.t_multi (pt1 VARCHAR, pt2 VARCHAR, id INT) PARTITIONED BY (pt1, pt2)")
6821 .await
6822 .unwrap();
6823 sql_context
6824 .sql("INSERT INTO paimon.test_db.t_multi VALUES ('a', 'x', 1)")
6825 .await
6826 .unwrap()
6827 .collect()
6828 .await
6829 .unwrap();
6830
6831 let err = sql_context
6832 .sql("TRUNCATE TABLE paimon.test_db.t_multi PARTITION (pt1 = 'a')")
6833 .await
6834 .unwrap_err();
6835 assert!(
6836 err.to_string().contains("Incomplete partition spec"),
6837 "Expected incomplete partition spec error, got: {err}"
6838 );
6839 }
6840
6841 #[tokio::test]
6842 async fn test_truncate_table_if_exists_nonexistent() {
6843 let (_tmp, sql_context) = setup_fs_sql_context().await;
6844
6845 sql_context
6846 .sql("TRUNCATE TABLE IF EXISTS paimon.test_db.nonexistent")
6847 .await
6848 .unwrap();
6849 }
6850
6851 #[tokio::test]
6852 async fn test_truncate_table_nonexistent_without_if_exists() {
6853 let (_tmp, sql_context) = setup_fs_sql_context().await;
6854
6855 let err = sql_context
6856 .sql("TRUNCATE TABLE paimon.test_db.nonexistent")
6857 .await
6858 .unwrap_err();
6859 assert!(
6860 err.to_string().contains("does not exist"),
6861 "Expected table-not-exist error, got: {err}"
6862 );
6863 }
6864
6865 #[tokio::test]
6866 async fn test_alter_table_if_exists_drop_partition_nonexistent() {
6867 let (_tmp, sql_context) = setup_fs_sql_context().await;
6868
6869 sql_context
6870 .sql("ALTER TABLE IF EXISTS paimon.test_db.nonexistent DROP PARTITION (pt = 'a')")
6871 .await
6872 .unwrap();
6873 }
6874
6875 #[tokio::test]
6876 async fn test_drop_partition_incomplete_spec() {
6877 let (_tmp, sql_context) = setup_fs_sql_context().await;
6878
6879 sql_context
6880 .sql("CREATE TABLE paimon.test_db.t_dp (pt1 VARCHAR, pt2 VARCHAR, id INT) PARTITIONED BY (pt1, pt2)")
6881 .await
6882 .unwrap();
6883 sql_context
6884 .sql("INSERT INTO paimon.test_db.t_dp VALUES ('a', 'x', 1)")
6885 .await
6886 .unwrap()
6887 .collect()
6888 .await
6889 .unwrap();
6890
6891 let err = sql_context
6892 .sql("ALTER TABLE paimon.test_db.t_dp DROP PARTITION (pt1 = 'a')")
6893 .await
6894 .unwrap_err();
6895 assert!(
6896 err.to_string().contains("Incomplete partition spec"),
6897 "Expected incomplete partition spec error, got: {err}"
6898 );
6899 }
6900
6901 #[tokio::test]
6902 async fn test_create_temp_table_if_not_exists() {
6903 let catalog = Arc::new(MockCatalog::new());
6904 let sql_context = make_sql_context(catalog).await;
6905
6906 sql_context
6908 .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)")
6909 .await
6910 .unwrap();
6911
6912 let err = sql_context
6914 .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)")
6915 .await
6916 .unwrap_err();
6917 assert!(
6918 err.to_string().contains("already exists"),
6919 "Expected already-exists error, got: {err}"
6920 );
6921
6922 sql_context
6924 .sql("CREATE TEMPORARY TABLE IF NOT EXISTS mydb.t1 (id INT)")
6925 .await
6926 .unwrap();
6927 }
6928
6929 #[tokio::test]
6930 async fn test_create_temp_table_if_not_exists_as_select() {
6931 let catalog = Arc::new(MockCatalog::new());
6932 let sql_context = make_sql_context(catalog).await;
6933
6934 sql_context
6936 .sql("CREATE TEMPORARY TABLE mydb.t2 AS SELECT 1 AS id")
6937 .await
6938 .unwrap();
6939
6940 sql_context
6942 .sql("CREATE TEMPORARY TABLE IF NOT EXISTS mydb.t2 AS SELECT 2 AS id")
6943 .await
6944 .unwrap();
6945
6946 let df = sql_context.sql("SELECT * FROM mydb.t2").await.unwrap();
6948 let batches = df.collect().await.unwrap();
6949 let val = batches[0]
6950 .column(0)
6951 .as_any()
6952 .downcast_ref::<Int64Array>()
6953 .unwrap();
6954 assert_eq!(val.value(0), 1);
6955 }
6956
6957 #[tokio::test]
6958 async fn test_create_temp_view_if_not_exists() {
6959 let catalog = Arc::new(MockCatalog::new());
6960 let sql_context = make_sql_context(catalog).await;
6961
6962 sql_context
6964 .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 1 AS id")
6965 .await
6966 .unwrap();
6967
6968 let err = sql_context
6970 .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 2 AS id")
6971 .await
6972 .unwrap_err();
6973 assert!(
6974 err.to_string().contains("already exists"),
6975 "Expected already-exists error, got: {err}"
6976 );
6977
6978 sql_context
6980 .sql("CREATE TEMPORARY VIEW IF NOT EXISTS mydb.v1 AS SELECT 3 AS id")
6981 .await
6982 .unwrap();
6983
6984 let df = sql_context.sql("SELECT * FROM mydb.v1").await.unwrap();
6986 let batches = df.collect().await.unwrap();
6987 let val = batches[0]
6988 .column(0)
6989 .as_any()
6990 .downcast_ref::<Int64Array>()
6991 .unwrap();
6992 assert_eq!(val.value(0), 1);
6993 }
6994
6995 #[tokio::test]
6996 async fn test_drop_temp_table_if_exists() {
6997 let catalog = Arc::new(MockCatalog::new());
6998 let sql_context = make_sql_context(catalog).await;
6999
7000 let err = sql_context
7002 .sql("DROP TEMPORARY TABLE mydb.nonexistent")
7003 .await
7004 .unwrap_err();
7005 let msg = err.to_string();
7006 assert!(
7007 msg.contains("doesn't exist")
7008 || msg.contains("does not exist")
7009 || msg.contains("Unknown temp database"),
7010 "Expected table-not-exist error, got: {msg}"
7011 );
7012
7013 sql_context
7015 .sql("DROP TEMPORARY TABLE IF EXISTS mydb.nonexistent")
7016 .await
7017 .unwrap();
7018
7019 sql_context
7021 .sql("CREATE TEMPORARY TABLE mydb.t1 (id INT)")
7022 .await
7023 .unwrap();
7024
7025 sql_context
7026 .sql("DROP TEMPORARY TABLE IF EXISTS mydb.t1")
7027 .await
7028 .unwrap();
7029
7030 assert!(
7032 !sql_context.temp_table_exist("mydb.t1").unwrap(),
7033 "Expected temp table to be gone after DROP"
7034 );
7035 }
7036
7037 #[tokio::test]
7038 async fn test_drop_temp_view_if_exists() {
7039 let catalog = Arc::new(MockCatalog::new());
7040 let sql_context = make_sql_context(catalog).await;
7041
7042 let err = sql_context
7044 .sql("DROP TEMPORARY VIEW mydb.nonexistent")
7045 .await
7046 .unwrap_err();
7047 let msg = err.to_string();
7048 assert!(
7049 msg.contains("doesn't exist")
7050 || msg.contains("does not exist")
7051 || msg.contains("Unknown temp database"),
7052 "Expected view-not-exist error, got: {msg}"
7053 );
7054
7055 sql_context
7057 .sql("DROP TEMPORARY VIEW IF EXISTS mydb.nonexistent")
7058 .await
7059 .unwrap();
7060
7061 sql_context
7063 .sql("CREATE TEMPORARY VIEW mydb.v1 AS SELECT 1 AS id")
7064 .await
7065 .unwrap();
7066
7067 sql_context
7068 .sql("DROP TEMPORARY VIEW IF EXISTS mydb.v1")
7069 .await
7070 .unwrap();
7071
7072 assert!(
7074 !sql_context.temp_table_exist("mydb.v1").unwrap(),
7075 "Expected temp view to be gone after DROP"
7076 );
7077 }
7078
7079 #[test]
7080 fn test_extract_version_as_of() {
7081 let sql = "SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 1";
7082 let infos = extract_all_version_as_of(sql);
7083 assert_eq!(infos.len(), 1);
7084 let info = &infos[0];
7085 assert_eq!(info.version, "1");
7086 assert_eq!(info.table_name, "paimon.default.time_travel_table");
7087 let rewritten = format!(
7088 "{}__uuid{}",
7089 &sql[..info.clause_range.0],
7090 &sql[info.clause_range.1..]
7091 );
7092 assert_eq!(rewritten, "SELECT id, name FROM __uuid");
7093 }
7094
7095 #[test]
7096 fn test_extract_version_as_of_multi_digit() {
7097 let sql = "SELECT * FROM mydb.t VERSION AS OF 42";
7098 let infos = extract_all_version_as_of(sql);
7099 assert_eq!(infos.len(), 1);
7100 let info = &infos[0];
7101 assert_eq!(info.version, "42");
7102 assert_eq!(info.table_name, "mydb.t");
7103 let rewritten = format!(
7104 "{}__uuid{}",
7105 &sql[..info.clause_range.0],
7106 &sql[info.clause_range.1..]
7107 );
7108 assert_eq!(rewritten, "SELECT * FROM __uuid");
7109 }
7110
7111 #[test]
7112 fn test_extract_version_as_of_case_insensitive() {
7113 let sql = "SELECT * FROM t version as of 5";
7114 let infos = extract_all_version_as_of(sql);
7115 assert_eq!(infos.len(), 1);
7116 let info = &infos[0];
7117 assert_eq!(info.version, "5");
7118 assert_eq!(info.table_name, "t");
7119 let rewritten = format!(
7120 "{}__uuid{}",
7121 &sql[..info.clause_range.0],
7122 &sql[info.clause_range.1..]
7123 );
7124 assert_eq!(rewritten, "SELECT * FROM __uuid");
7125 }
7126
7127 #[test]
7128 fn test_extract_version_as_of_not_present() {
7129 let sql = "SELECT * FROM t";
7130 assert!(extract_all_version_as_of(sql).is_empty());
7131 }
7132
7133 #[test]
7134 fn test_extract_version_as_of_tag() {
7135 let sql = "SELECT id, name FROM paimon.default.t VERSION AS OF 'snapshot1'";
7136 let infos = extract_all_version_as_of(sql);
7137 assert_eq!(infos.len(), 1);
7138 let info = &infos[0];
7139 assert_eq!(info.version, "snapshot1");
7140 assert_eq!(info.table_name, "paimon.default.t");
7141 let rewritten = format!(
7142 "{}__uuid{}",
7143 &sql[..info.clause_range.0],
7144 &sql[info.clause_range.1..]
7145 );
7146 assert_eq!(rewritten, "SELECT id, name FROM __uuid");
7147 }
7148
7149 #[test]
7150 fn test_extract_version_as_of_tag_case_insensitive() {
7151 let sql = "SELECT * FROM t version as of 'my_tag'";
7152 let infos = extract_all_version_as_of(sql);
7153 assert_eq!(infos.len(), 1);
7154 let info = &infos[0];
7155 assert_eq!(info.version, "my_tag");
7156 assert_eq!(info.table_name, "t");
7157 let rewritten = format!(
7158 "{}__uuid{}",
7159 &sql[..info.clause_range.0],
7160 &sql[info.clause_range.1..]
7161 );
7162 assert_eq!(rewritten, "SELECT * FROM __uuid");
7163 }
7164
7165 #[test]
7166 fn test_extract_version_as_of_numeric_still_works() {
7167 let sql = "SELECT * FROM t VERSION AS OF 123";
7168 let infos = extract_all_version_as_of(sql);
7169 assert_eq!(infos.len(), 1);
7170 assert_eq!(infos[0].version, "123");
7171 assert_eq!(infos[0].table_name, "t");
7172 }
7173
7174 #[test]
7175 fn test_extract_version_as_of_multiple() {
7176 let sql = "SELECT * FROM t1 VERSION AS OF 1 JOIN t2 VERSION AS OF 2 ON t1.id = t2.id";
7178 let infos = extract_all_version_as_of(sql);
7179 assert_eq!(infos.len(), 2);
7180 assert_eq!(infos[0].version, "1");
7181 assert_eq!(infos[0].table_name, "t1");
7182 assert_eq!(infos[1].version, "2");
7183 assert_eq!(infos[1].table_name, "t2");
7184 }
7185
7186 #[test]
7187 fn test_extract_version_as_of_skips_string_literal() {
7188 let sql = "SELECT * FROM t WHERE note = 'version as of 1'";
7189 let infos = extract_all_version_as_of(sql);
7190 assert!(infos.is_empty());
7191 }
7192
7193 #[test]
7194 fn test_extract_version_as_of_skips_comment() {
7195 let sql = "SELECT * FROM t -- version as of 1\n WHERE id > 0";
7196 let infos = extract_all_version_as_of(sql);
7197 assert!(infos.is_empty());
7198 }
7199
7200 #[test]
7201 fn test_contains_time_travel_keyword() {
7202 assert!(contains_time_travel_keyword(
7203 "SELECT * FROM t VERSION AS OF 1"
7204 ));
7205 assert!(contains_time_travel_keyword(
7206 "SELECT * FROM t TIMESTAMP AS OF '2024-01-01 00:00:00'"
7207 ));
7208 assert!(!contains_time_travel_keyword(
7210 "SELECT * FROM t WHERE note = 'version as of 1'"
7211 ));
7212 assert!(!contains_time_travel_keyword(
7214 "SELECT * FROM t -- version as of 1"
7215 ));
7216 assert!(!contains_time_travel_keyword(
7217 "SELECT * FROM t /* timestamp as of now */ WHERE id > 0"
7218 ));
7219 assert!(!contains_time_travel_keyword("SELECT * FROM t"));
7221 }
7222
7223 #[test]
7224 fn test_extract_timestamp_as_of() {
7225 let sql = "SELECT * FROM paimon.default.t TIMESTAMP AS OF '2024-01-15 10:30:00'";
7226 let infos = extract_all_timestamp_as_of(sql);
7227 assert_eq!(infos.len(), 1);
7228 let info = &infos[0];
7229 assert_eq!(info.timestamp, "2024-01-15 10:30:00");
7230 assert_eq!(info.table_name, "paimon.default.t");
7231 let rewritten = format!(
7232 "{}__uuid{}",
7233 &sql[..info.clause_range.0],
7234 &sql[info.clause_range.1..]
7235 );
7236 assert_eq!(rewritten, "SELECT * FROM __uuid");
7237 }
7238
7239 #[test]
7240 fn test_extract_timestamp_as_of_case_insensitive() {
7241 let sql = "SELECT * FROM t timestamp as of '2024-06-01 00:00:00'";
7242 let infos = extract_all_timestamp_as_of(sql);
7243 assert_eq!(infos.len(), 1);
7244 let info = &infos[0];
7245 assert_eq!(info.timestamp, "2024-06-01 00:00:00");
7246 assert_eq!(info.table_name, "t");
7247 let rewritten = format!(
7248 "{}__uuid{}",
7249 &sql[..info.clause_range.0],
7250 &sql[info.clause_range.1..]
7251 );
7252 assert_eq!(rewritten, "SELECT * FROM __uuid");
7253 }
7254
7255 #[test]
7256 fn test_extract_timestamp_as_of_not_present() {
7257 let sql = "SELECT * FROM t";
7258 assert!(extract_all_timestamp_as_of(sql).is_empty());
7259 }
7260}