diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 728428188..ba2b48d71 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -2102,6 +2102,11 @@ impl SQLContext { DataFusionError::Plan("TRUNCATE TABLE requires a table name".to_string()) })?; Self::ensure_main_branch_write_target(&target.name, "TRUNCATE TABLE")?; + if self.temp_table_for(&target.name).await?.is_some() { + return Err(DataFusionError::NotImplemented( + "TRUNCATE TABLE on temporary tables or views is not supported".into(), + )); + } let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&target.name)?; let table = match catalog.get_table(&identifier).await { Ok(t) => t, diff --git a/crates/integrations/datafusion/tests/sql_context_tests.rs b/crates/integrations/datafusion/tests/sql_context_tests.rs index 13e2a9d25..a9141f6e5 100644 --- a/crates/integrations/datafusion/tests/sql_context_tests.rs +++ b/crates/integrations/datafusion/tests/sql_context_tests.rs @@ -1076,6 +1076,64 @@ async fn test_update_and_delete_reject_shadowed_temporary_targets() { } } +#[tokio::test] +async fn test_truncate_rejects_shadowed_temporary_targets() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + ctx.sql("CREATE TABLE default.\"TruncateTarget\" (id INT, pt INT) PARTITIONED BY (pt)") + .await + .unwrap(); + ctx.sql("INSERT INTO default.\"TruncateTarget\" VALUES (1, 1), (2, 2)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + for kind in ["TABLE", "VIEW"] { + ctx.sql(&format!( + "CREATE TEMPORARY {kind} default.\"TruncateTarget\" AS SELECT CAST(999 AS INT) AS id" + )) + .await + .unwrap(); + assert_eq!( + collect_ids(&ctx, "SELECT id FROM \"TruncateTarget\"").await, + vec![999] + ); + for sql in [ + "TRUNCATE TABLE \"TruncateTarget\"", + "TRUNCATE TABLE IF EXISTS default.\"TruncateTarget\"", + "TRUNCATE TABLE paimon.default.\"TruncateTarget\" PARTITION (pt = 1)", + ] { + assert_sql_error_contains( + &ctx, + sql, + "TRUNCATE TABLE on temporary tables or views is not supported", + ) + .await; + } + ctx.sql(&format!("DROP TEMPORARY {kind} default.\"TruncateTarget\"")) + .await + .unwrap(); + assert_eq!( + collect_ids(&ctx, "SELECT id FROM \"TruncateTarget\"").await, + vec![1, 2] + ); + } + + ctx.sql("TRUNCATE TABLE \"TruncateTarget\" PARTITION (pt = 1)") + .await + .unwrap(); + assert_eq!( + collect_ids(&ctx, "SELECT id FROM \"TruncateTarget\"").await, + vec![2] + ); + ctx.sql("TRUNCATE TABLE \"TruncateTarget\"").await.unwrap(); + assert!(collect_ids(&ctx, "SELECT id FROM \"TruncateTarget\"") + .await + .is_empty()); +} + #[tokio::test] async fn test_scoped_show_accepts_information_schema() { let (_tmp, catalog) = create_test_env();