chore: add better regclass and oid cast support (#2795)

partially addresses #2746

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(_)

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<V>(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<Self::Break> {
@@ -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.

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)?;
         }
 

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);

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