diff --git a/crates/datafusion_ext/src/planner/mod.rs b/crates/datafusion_ext/src/planner/mod.rs index cc70333cc..62514a197 100644 --- a/crates/datafusion_ext/src/planner/mod.rs +++ b/crates/datafusion_ext/src/planner/mod.rs @@ -281,15 +281,17 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> { SQLDataType::Bytea => Ok(DataType::Binary), SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), SQLDataType::Custom(obj, _) => { - let obj = obj.to_string(); + let obj = obj.to_string().to_lowercase(); match obj.as_str() { // PSQL uses `pg_catalog.text` for `text` type in some cases "pg_catalog.text" => Ok(DataType::Utf8), - _ => Err(DataFusionError::NotImplemented(format!( - "Unsupported SQL type {sql_type:?}" - ))), + "oid" => Ok(DataType::Int64), + _ => not_impl_err!( + "Unsupported custom SQL type {sql_type:?}" + ), } } + SQLDataType::Regclass => Ok(DataType::Int64), // Explicitly list all other types so that if sqlparser // adds/changes the `SQLDataType` the compiler will tell us on upgrade // and avoid bugs like https://github.com/apache/arrow-datafusion/issues/3059 @@ -300,7 +302,6 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> { | SQLDataType::Varbinary(_) | SQLDataType::Blob(_) | SQLDataType::Datetime(_) - | SQLDataType::Regclass | SQLDataType::Array(_) | SQLDataType::Enum(_) | SQLDataType::Set(_) diff --git a/crates/sqlexec/src/planner/preprocess.rs b/crates/sqlexec/src/planner/preprocess.rs index c8e06ea5b..98aa15693 100644 --- a/crates/sqlexec/src/planner/preprocess.rs +++ b/crates/sqlexec/src/planner/preprocess.rs @@ -13,6 +13,9 @@ pub enum PreprocessError { #[error("Casting expressions to regclass unsupported")] ExprUnsupportedRegclassCast, + + #[error("Casting expressions to oid unsupported")] + ExprUnsupportedOIDCast, } pub fn preprocess(statement: &mut ast::Statement, visitor: &mut V) -> Result<(), PreprocessError> @@ -25,13 +28,13 @@ where } } -/// Replace `CAST('table_name' as REGCLASS)` expressions with the oid of the +/// Replace `CAST('table_name' as [REGCLASS | OID])` expressions with the oid of the /// table. -pub struct CastRegclassReplacer<'a> { +pub struct CastOIDReplacer<'a> { pub ctx: &'a LocalSessionContext, } -impl<'a> ast::VisitorMut for CastRegclassReplacer<'a> { +impl<'a> ast::VisitorMut for CastOIDReplacer<'a> { type Break = PreprocessError; fn post_visit_expr(&mut self, expr: &mut ast::Expr) -> ControlFlow { @@ -50,9 +53,14 @@ impl<'a> ast::VisitorMut for CastRegclassReplacer<'a> { let replace_expr = match expr { ast::Expr::Cast { expr: inner_expr, - data_type: ast::DataType::Regclass, + data_type, format: _, } => { + match data_type { + ast::DataType::Regclass => {} + ast::DataType::Custom(name, _) if name.to_string().to_lowercase() == "oid" => {} + _ => return ControlFlow::Continue(()), // Nothing to do. + } if let ast::Expr::Value(ast::Value::SingleQuotedString(relation)) = &**inner_expr { match find_oid(self.ctx, relation) { Some(oid) => ast::Expr::Value(ast::Value::Number(oid.to_string(), false)), @@ -63,8 +71,13 @@ impl<'a> ast::VisitorMut for CastRegclassReplacer<'a> { } } } else { - // We don't currently support any other casts to regclass. - return ControlFlow::Break(PreprocessError::ExprUnsupportedRegclassCast); + // We don't currently support any other casts to regclass or oid. + let e = match data_type { + ast::DataType::Regclass => PreprocessError::ExprUnsupportedRegclassCast, + ast::DataType::Custom(_, _) => PreprocessError::ExprUnsupportedOIDCast, + _ => unreachable!(), + }; + return ControlFlow::Break(e); } } _ => return ControlFlow::Continue(()), // Nothing to do. diff --git a/crates/sqlexec/src/planner/session_planner.rs b/crates/sqlexec/src/planner/session_planner.rs index f59b9df4c..f800e2714 100644 --- a/crates/sqlexec/src/planner/session_planner.rs +++ b/crates/sqlexec/src/planner/session_planner.rs @@ -171,7 +171,7 @@ use crate::planner::logical_plan::{ TransactionPlan, Update, }; -use crate::planner::preprocess::{preprocess, CastRegclassReplacer, EscapedStringToDoubleQuoted}; +use crate::planner::preprocess::{preprocess, CastOIDReplacer, EscapedStringToDoubleQuoted}; use crate::remote::table::StubRemoteTableProvider; use crate::resolve::{EntryResolver, ResolvedEntry}; @@ -226,7 +226,7 @@ impl<'a> SessionPlanner<'a> { // Run replacers as needed. if let StatementWithExtensions::Statement(inner) = &mut statement { - preprocess(inner, &mut CastRegclassReplacer { ctx: self.ctx })?; + preprocess(inner, &mut CastOIDReplacer { ctx: self.ctx })?; preprocess(inner, &mut EscapedStringToDoubleQuoted)?; } diff --git a/testdata/sqllogictests/cast/oid.slt b/testdata/sqllogictests/cast/oid.slt new file mode 100644 index 000000000..b77f85483 --- /dev/null +++ b/testdata/sqllogictests/cast/oid.slt @@ -0,0 +1,30 @@ + +statement error +select 'my_table'::oid; + +statement ok +create external table my_table from debug options(table_type = 'never_ending'); + +statement ok +select 'my_table'::oid; + +# make sure upper case works +statement ok +select 'my_table'::OID; + +statement ok +select cast('my_table' as oid); + +statement ok +select cast('my_table' as OID); + + +# We don't support expressions other than strings right now. +statement error Casting expressions to oid unsupported +select cast(1 as oid); + +statement error Casting expressions to oid unsupported +select cast(1 + 1 as oid); + +statement error Relation 'not_my_table' does not exist +select cast('not_my_table' as oid); \ No newline at end of file diff --git a/testdata/sqllogictests/cast/regclass.slt b/testdata/sqllogictests/cast/regclass.slt index fea155e8b..0e21f2105 100644 --- a/testdata/sqllogictests/cast/regclass.slt +++ b/testdata/sqllogictests/cast/regclass.slt @@ -24,6 +24,13 @@ select 'my_table'::regclass > 2000; ---- t +# make sure it works with the `CAST` syntax too +statement ok +select + cast('my_table' as regclass) as lower, + cast('my_table' as REGCLASS) as upper; + + # We don't support expressions other than strings right now. statement error