refactor: move oid name-lookup machinery out of the coercion rule

datafusion-pg-catalog/src/pg_catalog.rs

@@ -12,10 +12,12 @@ use datafusion::arrow::ipc::reader::FileReader;
 use datafusion::catalog::streaming::StreamingTable;
 use datafusion::catalog::{MemTable, SchemaProvider, TableFunctionImpl};
 use datafusion::common::utils::SingleRowListArrayBuilder;
-use datafusion::datasource::TableProvider;
+use datafusion::common::{Column, ScalarValue, Spans};
+use datafusion::datasource::{TableProvider, provider_as_source};
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_expr::{
-    ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+    BinaryExpr, ColumnarValue, LogicalPlanBuilder, Operator, ScalarUDF, ScalarUDFImpl, Signature,
+    Subquery, TypeSignature, Volatility,
 };
 use datafusion::physical_plan::streaming::PartitionStream;
 use datafusion::prelude::{Expr, SessionContext, create_udf};
@@ -229,16 +231,70 @@ impl<C: CatalogInfo, P: PgCatalogContextProvider> SchemaProvider for PgCatalogSc
     }
 }
 
-/// Backs the oid-coercion analyzer rule: build a pg_catalog table provider
-/// synchronously so the rule can construct name->oid lookup subqueries at
-/// plan time (data is fetched lazily at execution).
+/// The `(table, name-column)` backing a name->oid lookup for each oid-alias
+/// kind, or `None` for kinds without name-resolution support (e.g. bare `oid`).
+fn oid_lookup_target(kind: &str) -> Option<(&'static str, &'static str)> {
+    match kind {
+        oid_field::kind::REGNAMESPACE => Some((PG_CATALOG_TABLE_PG_NAMESPACE, "nspname")),
+        oid_field::kind::REGCLASS => Some((PG_CATALOG_TABLE_PG_CLASS, "relname")),
+        oid_field::kind::REGTYPE => Some((PG_CATALOG_TABLE_PG_TYPE, "typname")),
+        oid_field::kind::REGPROC => Some((PG_CATALOG_TABLE_PG_PROC, "proname")),
+        _ => None,
+    }
+}
+
+/// Build a scalar subquery `(SELECT oid FROM <table> WHERE <name_col> = '<name>')`
+/// that resolves a name to its oid at execution time. Data is fetched lazily
+/// by the table's stream, so oids assigned at execution (dynamic tables like
+/// `pg_namespace`) resolve correctly.
+///
+/// Returns `None` if the kind has no backing table or the provider cannot
+/// build it, in which case the coercion rule leaves the comparison untouched.
+fn build_oid_name_subquery(
+    kind: &str,
+    name: &str,
+    build_table: impl Fn(&str) -> Result<Option<Arc<dyn TableProvider>>>,
+) -> Result<Option<Expr>> {
+    let Some((table, name_col)) = oid_lookup_target(kind) else {
+        return Ok(None);
+    };
+    let Some(provider) = build_table(table)? else {
+        return Ok(None);
+    };
+
+    let source = provider_as_source(provider);
+    let plan = LogicalPlanBuilder::scan(table.to_string(), source, None)?
+        .filter(Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column::new_unqualified(name_col))),
+            op: Operator::Eq,
+            right: Box::new(Expr::Literal(
+                ScalarValue::Utf8(Some(name.to_string())),
+                None,
+            )),
+        }))?
+        .project(vec![Expr::Column(Column::new_unqualified("oid"))])?
+        .build()?;
+    Ok(Some(Expr::ScalarSubquery(Subquery {
+        subquery: Arc::new(plan),
+        outer_ref_columns: vec![],
+        spans: Spans::default(),
+    })))
+}
+
+/// Backs the oid-coercion analyzer rule: resolve a name string to an
+/// oid-valued expression by looking it up in the matching pg_catalog table.
+/// The lookup logic (which table backs which kind, and the subquery
+/// construction) lives here rather than in the rule, so the rule stays a pure
+/// expression-rewriter with no dependency on `TableProvider` or plan building.
 impl<C: CatalogInfo, P: PgCatalogContextProvider> oid_coercion_rule::OidLookupProvider
     for PgCatalogSchemaProvider<C, P>
 {
-    fn build_table_provider(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
-        self.build_table_by_name(name)?
-            .map(|t| t.try_into_table_provider())
-            .transpose()
+    fn resolve_name_expr(&self, kind: &str, name: &str) -> Result<Option<Expr>> {
+        build_oid_name_subquery(kind, name, |table| {
+            self.build_table_by_name(table)?
+                .map(|t| t.try_into_table_provider())
+                .transpose()
+        })
     }
 }
 
@@ -1541,6 +1597,30 @@ where
 mod test {
     use super::*;
 
+    #[test]
+    fn oid_lookup_target_maps_kinds_to_backing_tables() {
+        // Each name-resolvable oid-alias kind maps to the pg_catalog table
+        // (and name column) that backs its name->oid lookup.
+        assert_eq!(
+            oid_lookup_target(oid_field::kind::REGNAMESPACE),
+            Some((PG_CATALOG_TABLE_PG_NAMESPACE, "nspname"))
+        );
+        assert_eq!(
+            oid_lookup_target(oid_field::kind::REGCLASS),
+            Some((PG_CATALOG_TABLE_PG_CLASS, "relname"))
+        );
+        assert_eq!(
+            oid_lookup_target(oid_field::kind::REGTYPE),
+            Some((PG_CATALOG_TABLE_PG_TYPE, "typname"))
+        );
+        assert_eq!(
+            oid_lookup_target(oid_field::kind::REGPROC),
+            Some((PG_CATALOG_TABLE_PG_PROC, "proname"))
+        );
+        // bare `oid` has no backing table -> name resolution unsupported.
+        assert_eq!(oid_lookup_target(oid_field::kind::OID), None);
+    }
+
     #[tokio::test]
     async fn current_database_returns_the_configured_catalog_name() {
         use crate::pg_catalog::context::EmptyContextProvider;

datafusion-pg-catalog/src/pg_catalog/oid_coercion_rule.rs

@@ -19,10 +19,11 @@
 //! against a string literal, the string is resolved to an oid using Postgres
 //! semantics:
 //!
-//! * a **numeric** string (`'16417'`) becomes a literal `Int32` oid;
-//! * a **name** string (`'public'`) becomes a scalar subquery that looks the
-//!   name up in the backing catalog table (`pg_namespace`, `pg_class`,
-//!   `pg_proc`).
+//! * a **numeric** string (`'16417'`) becomes a literal `Int32` oid (handled
+//!   inline -- no catalog access needed);
+//! * a **name** string (`'public'`) is resolved by the [`OidLookupProvider`]
+//!   (e.g. into a scalar subquery against the backing catalog table). The rule
+//!   itself stays free of any data-access concerns.
 //!
 //! The rule is order-independent: it handles both the pre-coercion form
 //! `col = 'x'` and the post-coercion form `CAST(col AS Utf8) = 'x'`.
@@ -33,25 +34,26 @@ use std::sync::Arc;
 use datafusion::common::Column;
 use datafusion::common::config::ConfigOptions;
 use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
-use datafusion::common::{DFSchemaRef, ExprSchema, Result, ScalarValue, Spans};
-use datafusion::datasource::{TableProvider, provider_as_source};
+use datafusion::common::{DFSchemaRef, ExprSchema, Result, ScalarValue};
 use datafusion::logical_expr::expr::InList;
-use datafusion::logical_expr::{
-    BinaryExpr, Cast, Expr, LogicalPlan, LogicalPlanBuilder, Operator, Subquery, TryCast,
-};
+use datafusion::logical_expr::{BinaryExpr, Cast, Expr, LogicalPlan, Operator, TryCast};
 use datafusion::optimizer::analyzer::AnalyzerRule;
 
-use crate::pg_catalog::oid_field::{self, OID_ALIAS_KEY};
+use crate::pg_catalog::oid_field::OID_ALIAS_KEY;
 
-/// Synchronous accessor for the pg_catalog tables backing name->oid lookups.
+/// Resolves name strings to oids for the [`OidStringCoercion`] analyzer rule.
 ///
-/// Implemented by `PgCatalogSchemaProvider`, which can build a table provider
-/// without async I/O (dynamic tables defer data fetch to execution time via a
-/// streaming source).
+/// Implemented by the pg_catalog provider, which owns the catalog tables
+/// backing each name->oid lookup. Encapsulating the lookup behind this trait
+/// keeps the rule a pure expression-rewriter -- it has no direct dependency on
+/// `TableProvider` or logical-plan construction, and the rule does not need to
+/// know which pg_catalog table backs a given oid-alias kind.
 pub trait OidLookupProvider: Send + Sync + Debug {
-    /// Build the named pg_catalog table as a [`TableProvider`], or `None` if it
-    /// does not exist.
-    fn build_table_provider(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>>;
+    /// Resolve `name` to an oid-valued [`Expr`] for the given oid-alias
+    /// `kind` (`regnamespace`, `regclass`, `regtype`, `regproc`). Returns
+    /// `None` when the kind has no backing catalog table or the table is not
+    /// available, so the rule leaves the comparison untouched.
+    fn resolve_name_expr(&self, kind: &str, name: &str) -> Result<Option<Expr>>;
 }
 
 /// Rewrite oid / oid-alias column comparisons against string literals to use
@@ -294,26 +296,12 @@ fn as_str_literal(e: &Expr) -> Option<&str> {
     }
 }
 
-/// The `(table, name_column)` used to resolve a name to an oid for a given
-/// oid-alias kind. Returns `None` for kinds without a backing catalog table
-/// (e.g. bare `oid`, or `regtype` where `pg_type` is unavailable) -- those only
-/// support numeric resolution.
-fn lookup_target(kind: &str) -> Option<(&'static str, &'static str)> {
-    match kind {
-        oid_field::kind::REGNAMESPACE => Some(("pg_namespace", "nspname")),
-        oid_field::kind::REGCLASS => Some(("pg_class", "relname")),
-        oid_field::kind::REGTYPE => Some(("pg_type", "typname")),
-        oid_field::kind::REGPROC => Some(("pg_proc", "proname")),
-        _ => None,
-    }
-}
-
 /// Resolve a string operand to an oid-valued expression.
 ///
-/// * numeric strings -> literal `Int32` (Postgres behavior for every
-///   oid-alias type);
-/// * names -> a scalar subquery `SELECT oid FROM <table> WHERE <namecol> = $s`
-///   for kinds that have a backing table.
+/// * numeric strings -> a literal `Int32` oid (valid for every oid-alias
+///   kind, and the only form of "coercion" that needs no catalog access);
+/// * name strings -> delegated to the [`OidLookupProvider`], which owns the
+///   name->oid lookup and may return a subquery / literal / etc.
 fn resolve_operand(
     kind: &str,
     s: &str,
@@ -327,58 +315,16 @@ fn resolve_operand(
             .map(|n| Expr::Literal(ScalarValue::Int32(Some(n)), None)));
     }
 
-    // Name -> oid lookup subquery (only for kinds with a backing table).
-    let Some((table, name_col)) = lookup_target(kind) else {
-        return Ok(None);
-    };
-    let Some(provider) = provider.build_table_provider(table)? else {
-        return Ok(None);
-    };
-
-    let source = provider_as_source(provider);
-    let plan = LogicalPlanBuilder::scan(table.to_string(), source, None)?
-        .filter(eq_str(name_col, s))?
-        .project(vec![Expr::Column(Column::new_unqualified("oid"))])?
-        .build()?;
-    Ok(Some(Expr::ScalarSubquery(Subquery {
-        subquery: Arc::new(plan),
-        outer_ref_columns: vec![],
-        spans: Spans::default(),
-    })))
-}
-
-/// Build `<name_col> = '<s>'` against a freshly-scanned table.
-fn eq_str(name_col: &str, s: &str) -> Expr {
-    Expr::BinaryExpr(BinaryExpr {
-        left: Box::new(Expr::Column(Column::new_unqualified(name_col))),
-        op: Operator::Eq,
-        right: Box::new(Expr::Literal(ScalarValue::Utf8(Some(s.to_string())), None)),
-    })
+    // Name string -> the provider resolves it. The rule deliberately does not
+    // know which catalog table backs each kind.
+    provider.resolve_name_expr(kind, s)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use datafusion::arrow::datatypes::DataType;
 
-    #[test]
-    fn lookup_target_only_for_backed_kinds() {
-        assert_eq!(
-            lookup_target(oid_field::kind::REGNAMESPACE),
-            Some(("pg_namespace", "nspname"))
-        );
-        assert_eq!(
-            lookup_target(oid_field::kind::REGCLASS),
-            Some(("pg_class", "relname"))
-        );
-        assert_eq!(
-            lookup_target(oid_field::kind::REGPROC),
-            Some(("pg_proc", "proname"))
-        );
-        // bare oid has no backing table -> name resolution unsupported
-        assert_eq!(lookup_target(oid_field::kind::OID), None);
-    }
-
     #[test]
     fn unwrap_column_handles_bare_and_cast_forms() {
         let bare = Expr::Column(Column::new_unqualified("relnamespace"));