improve error reporting for multistatement sql

src/daft-sql/src/planner.rs

@@ -12,8 +12,8 @@ use daft_plan::{LogicalPlanBuilder, LogicalPlanRef};
 use sqlparser::{
     ast::{
         ArrayElemTypeDef, BinaryOperator, CastKind, ExactNumberInfo, GroupByExpr, Ident, Query,
-        SelectItem, StructField, Subscript, TableWithJoins, TimezoneInfo, UnaryOperator, Value,
-        WildcardAdditionalOptions,
+        SelectItem, Statement, StructField, Subscript, TableWithJoins, TimezoneInfo, UnaryOperator,
+        Value, WildcardAdditionalOptions,
     },
     dialect::GenericDialect,
     parser::{Parser, ParserOptions},
@@ -88,9 +88,18 @@ impl SQLPlanner {
 
         let statements = parser.parse_statements()?;
 
-        match statements.as_slice() {
-            [sqlparser::ast::Statement::Query(query)] => Ok(self.plan_query(query)?.build()),
-            other => unsupported_sql_err!("{}", other[0]),
+        match statements.len() {
+            1 => Ok(self.plan_statement(&statements[0])?),
+            other => {
+                unsupported_sql_err!("Only exactly one SQL statement allowed, found {}", other)
+            }
+        }
+    }
+
+    fn plan_statement(&mut self, statement: &Statement) -> SQLPlannerResult<LogicalPlanRef> {
+        match statement {
+            Statement::Query(query) => Ok(self.plan_query(query)?.build()),
+            other => unsupported_sql_err!("{}", other),
         }
     }
 

tests/sql/test_sql.py

@@ -149,3 +149,9 @@ def test_sql_count_star():
     actual = df2.collect().to_pydict()
     expected = df.agg(daft.col("b").count()).collect().to_pydict()
     assert actual == expected
+
+
+def test_sql_multi_statement_sql_error():
+    catalog = SQLCatalog({})
+    with pytest.raises(Exception, match="one SQL statement allowed"):
+        daft.sql("SELECT * FROM df; SELECT * FROM df", catalog)