feat: basic support for executing prepared statements #13242

jonahgao · 2024-11-04T06:37:23Z

Which issue does this PR close?

Closes #4549.

Rationale for this change

Store the logical plan of the prepared statement in the session state. During EXECUTE, retrieve the prepared logical plan, replace the placeholders with execution parameters, and then execute.

What changes are included in this PR?

Support executing prepared statements.

DataFusion CLI v42.2.0
> prepare p(int) as select 100+$1;
0 row(s) fetched.
Elapsed 0.004 seconds.

> execute p(10);
+-----------------+
| Int64(100) + $1 |
+-----------------+
| 110             |
+-----------------+
1 row(s) fetched.
Elapsed 0.006 seconds.

Are these changes tested?

Yes

Are there any user-facing changes?

Yes.

ctx.sql("Prepare p(int) as query") now returns an empty DataFrame, similar to Create View and other DDL statements. Previously, it would return a DataFrame containing the Prepare logical plan.

jonahgao · 2024-11-04T06:41:35Z

datafusion/core/tests/sql/select.rs

@@ -176,27 +175,6 @@ async fn prepared_statement_type_coercion() -> Result<()> {
    Ok(())
 }

-#[tokio::test]
-async fn prepared_statement_invalid_types() -> Result<()> {


Remove this test because now parameters will be cast to the target type.
This is also the behavior of PostgreSQL.

psql=> prepare p(int) as select 100+$1; PREPARE psql=> execute p('100'); ?column? ---------- 200 (1 row) psql=> execute p(20.12); ?column? ---------- 120 (1 row)

jonahgao · 2024-11-04T06:47:49Z

datafusion/core/tests/sql/select.rs

@@ -106,9 +105,9 @@ async fn test_prepare_statement() -> Result<()> {
    let ctx = create_ctx_with_partition(&tmp_dir, partition_count).await?;

    // sql to statement then to prepare logical plan with parameters
-    // c1 defined as UINT32, c2 defined as UInt64 but the params are Int32 and Float64
-    let dataframe =
-        ctx.sql("PREPARE my_plan(INT, DOUBLE) AS SELECT c1, c2 FROM test WHERE c1 > $2 AND c1 < $1").await?;


PREPARE now returns an empty DataFrame and can't call with_param_values on it. I think this change is reasonable because the result of PREPARE is not a real relation, as it doesn't contain valid rows. Its result is more like DDL.

I agree -- PREPARE is a statement in my mind (and has no results)

alamb

Thank you so much @jonahgao -- this PR looks basically perfect. My only concern is the security implications of executing prepared statements when the SQL is defined to be "read only"

It would also be neat to have some way to deallocate prepared statements, but that would be a natural follow on ticket (perhaps adding support for DEALLOCATE for example)

alamb · 2024-11-05T18:07:43Z

datafusion/core/src/execution/context/mod.rs

@@ -687,7 +688,31 @@ impl SessionContext {
            LogicalPlan::Statement(Statement::SetVariable(stmt)) => {
                self.set_variable(stmt).await
            }
-
+            LogicalPlan::Prepare(Prepare {


I am worried that LogicalPlan::Prepare seems to be runnable even if we run sql_with_options and disable statements.

datafusion/datafusion/core/src/execution/context/mod.rs

Lines 605 to 614 in aad7744

pub async fn sql_with_options(

&self,

sql: &str,

options: SQLOptions,

) -> Result<DataFrame> {

let plan = self.state().create_logical_plan(sql).await?;

options.verify_plan(&plan)?;

self.execute_logical_plan(plan).await

}

I realize that this PR does not change if Prepare is a statement or not, but it is the first that actually makes Prepare do something

Could you please add a test like this (but for PREPARE statements) and make sure the statement can't be executed if statements are disabled?

datafusion/datafusion/core/tests/sql/sql_api.rs

Lines 99 to 114 in 6692382

async fn unsupported_statement_returns_error() {

let ctx = SessionContext::new();

ctx.sql("CREATE TABLE test (x int)").await.unwrap();

let options = SQLOptions::new().with_allow_statements(false);

let sql = "set datafusion.execution.batch_size = 5";

let df = ctx.sql_with_options(sql, options).await;

assert_eq!(

df.unwrap_err().strip_backtrace(),

"Error during planning: Statement not supported: SetVariable"

);

let options = options.with_allow_statements(true);

ctx.sql_with_options(sql, options).await.unwrap();

}

Perhaps (as a follow on PR) we can make LogicalPlan::Prepare a statement https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Statement.html

Temporarily added PREPARE to BadPlanVisitor and added a test for this in 4729766
I will make LogicalPlan::Prepare a statement in a follow-up PR.

alamb · 2024-11-05T18:08:15Z

datafusion/core/src/execution/context/mod.rs

+                .collect::<Result<_>>()?;
+        }
+
+        let params = ParamValues::List(params);


alamb · 2024-11-05T18:09:29Z

datafusion/core/src/execution/session_state.rs

+                Ok(())
+            }
+            Entry::Occupied(e) => {
+                exec_err!("Prepared statement '{}' already exists", e.key())


Since this errors if an existing prepared statement should we add some way to erase / clear prepared plans? Now there is no way to avoid accumulating them over time

I plan to implement DEALLOCATE in next PR.

alamb · 2024-11-05T18:09:55Z

datafusion/core/tests/sql/select.rs

@@ -106,9 +105,9 @@ async fn test_prepare_statement() -> Result<()> {
    let ctx = create_ctx_with_partition(&tmp_dir, partition_count).await?;

    // sql to statement then to prepare logical plan with parameters
-    // c1 defined as UINT32, c2 defined as UInt64 but the params are Int32 and Float64
-    let dataframe =
-        ctx.sql("PREPARE my_plan(INT, DOUBLE) AS SELECT c1, c2 FROM test WHERE c1 > $2 AND c1 < $1").await?;


I agree -- PREPARE is a statement in my mind (and has no results)

alamb · 2024-11-05T18:10:22Z

datafusion/sqllogictest/test_files/prepare.slt

@@ -30,56 +30,175 @@ select * from person;
 # Error due to syntax and semantic violation

 # Syntax error: no name specified after the keyword prepare
-statement error
+statement error DataFusion error: SQL error: ParserError


alamb · 2024-11-05T18:11:01Z

datafusion/sqllogictest/test_files/prepare.slt

+2 Bar
+
+
+# Test issue: https://github.com/apache/datafusion/issues/12294


alamb

Thank you @jonahgao -- this looks great

alamb · 2024-11-06T15:49:35Z

datafusion/core/src/execution/context/mod.rs

@@ -1705,6 +1773,14 @@ impl<'n, 'a> TreeNodeVisitor<'n> for BadPlanVisitor<'a> {
            LogicalPlan::Statement(stmt) if !self.options.allow_statements => {
                plan_err!("Statement not supported: {}", stmt.name())
            }
+            // TODO: Implement PREPARE as a LogicalPlan::Statement
+            LogicalPlan::Prepare(_) if !self.options.allow_statements => {


jonahgao added 6 commits November 1, 2024 21:16

feat: basic support for executing prepared statements

b75c3c7

Merge branch 'main' into prepare

31f7316

Improve execute_prepared

317dd84

Fix tests

01f5b69

Merge branch 'main' into prepare

35a682f

Update doc

c81c1b4

github-actions bot added logical-expr Logical plan and expressions core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Nov 4, 2024

jonahgao commented Nov 4, 2024

View reviewed changes

jonahgao added 2 commits November 4, 2024 14:54

Add test

aad7744

Add issue test

b9a9674

alamb reviewed Nov 5, 2024

View reviewed changes

jonahgao added 2 commits November 6, 2024 10:40

Respect allow_statements option

4729766

Merge branch 'main' into prepare

0dc4860

alamb approved these changes Nov 6, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: basic support for executing prepared statements #13242

feat: basic support for executing prepared statements #13242

jonahgao commented Nov 4, 2024 •

edited

Loading

jonahgao Nov 4, 2024

jonahgao Nov 4, 2024 •

edited

Loading

alamb Nov 5, 2024

alamb left a comment

alamb Nov 5, 2024

jonahgao Nov 6, 2024

alamb Nov 5, 2024

alamb Nov 5, 2024

jonahgao Nov 6, 2024

alamb Nov 5, 2024

alamb Nov 5, 2024

alamb Nov 5, 2024

alamb left a comment

alamb Nov 6, 2024

	pub async fn sql_with_options(
	&self,
	sql: &str,
	options: SQLOptions,
	) -> Result<DataFrame> {
	let plan = self.state().create_logical_plan(sql).await?;
	options.verify_plan(&plan)?;

	self.execute_logical_plan(plan).await
	}

	async fn unsupported_statement_returns_error() {
	let ctx = SessionContext::new();
	ctx.sql("CREATE TABLE test (x int)").await.unwrap();

	let options = SQLOptions::new().with_allow_statements(false);

	let sql = "set datafusion.execution.batch_size = 5";
	let df = ctx.sql_with_options(sql, options).await;
	assert_eq!(
	df.unwrap_err().strip_backtrace(),
	"Error during planning: Statement not supported: SetVariable"
	);

	let options = options.with_allow_statements(true);
	ctx.sql_with_options(sql, options).await.unwrap();
	}

		2 Bar


		# Test issue: https://github.com/apache/datafusion/issues/12294

feat: basic support for executing prepared statements #13242

Are you sure you want to change the base?

feat: basic support for executing prepared statements #13242

Conversation

jonahgao commented Nov 4, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

jonahgao Nov 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonahgao commented Nov 4, 2024 •

edited

Loading

jonahgao Nov 4, 2024 •

edited

Loading