Skip to content

Commit

Permalink
[BUG] improve error reporting for multistatement sql (#2916)
Browse files Browse the repository at this point in the history
Called it a bug because the error is confusing, this fixes error
behavior of multistatement sql. Before:

```py
daft.sql('''SELECT * FROM df; SELECT count(*) FROM df''', cat)
```

results in this error:
```py
daft.exceptions.InvalidSQLException: Unsupported SQL: 'SELECT * FROM df'
```

but of course that SQL is supported! With this change, error becomes:

```py
daft.exceptions.InvalidSQLException: Unsupported SQL: 'Only exactly one SQL statement allowed, found 2'
```

which I believe is at least currently correct, and any future support of
multiple statements would have to touch this area of code anyway.
  • Loading branch information
amitschang authored Sep 27, 2024
1 parent ab6d98d commit 2935abd
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 5 deletions.
19 changes: 14 additions & 5 deletions src/daft-sql/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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),
}
}

Expand Down
6 changes: 6 additions & 0 deletions tests/sql/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,3 +199,9 @@ def test_sql_function_raises_when_cant_get_frame(monkeypatch):
monkeypatch.setattr("inspect.currentframe", lambda: None)
with pytest.raises(DaftCoreException, match="Cannot get caller environment"):
daft.sql("SELECT * FROM df")


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)

0 comments on commit 2935abd

Please sign in to comment.