Merge pull request #315 from remilapeyre/import-foreign-schema

Change import_foreign_schema() signature.
supabase · Aug 2, 2024 · 9440656 · 9440656
2 parents 37f9bea + d0603c9
commit 9440656
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 13 deletions.
diff --git a/supabase-wrappers/src/import_foreign_schema.rs b/supabase-wrappers/src/import_foreign_schema.rs
@@ -1,10 +1,28 @@
-use pg_sys::AsPgCStr;
+use pg_sys::{AsPgCStr, Oid};
 use pgrx::pg_sys::panic::ErrorReport;
 use pgrx::{debug2, prelude::*, PgList};
+use std::marker::PhantomData;
 
+use crate::instance;
 use crate::options::options_to_hashmap;
 use crate::prelude::ForeignDataWrapper;
 
+// Fdw private state for import_foreign_schema
+struct FdwState<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> {
+    // foreign data wrapper instance
+    instance: W,
+    _phantom: PhantomData<E>,
+}
+
+impl<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> FdwState<E, W> {
+    unsafe fn new(foreignserverid: Oid) -> Self {
+        Self {
+            instance: instance::create_fdw_instance_from_server_id(foreignserverid),
+            _phantom: PhantomData,
+        }
+    }
+}
+
 #[repr(u32)]
 #[derive(Debug, Clone)]
 pub enum ListType {
@@ -30,10 +48,10 @@ pub(super) extern "C" fn import_foreign_schema<E: Into<ErrorReport>, W: ForeignD
 ) -> *mut pg_sys::List {
     debug2!("---> import_foreign_schema");
 
-    let import_foreign_schema_stmt: ImportForeignSchemaStmt;
+    let create_stmts: Vec<String>;
 
     unsafe {
-        import_foreign_schema_stmt = ImportForeignSchemaStmt {
+        let import_foreign_schema_stmt = ImportForeignSchemaStmt {
             server_name: std::ffi::CStr::from_ptr((*stmt).server_name)
                 .to_str()
                 .unwrap()
@@ -76,11 +94,16 @@ pub(super) extern "C" fn import_foreign_schema<E: Into<ErrorReport>, W: ForeignD
             },
 
             options: options_to_hashmap((*stmt).options).unwrap(),
-        }
+        };
+
+        let mut state = FdwState::<E, W>::new(server_oid);
+        create_stmts = state
+            .instance
+            .import_foreign_schema(import_foreign_schema_stmt);
     }
 
     let mut ret: PgList<i8> = PgList::new();
-    for command in W::import_foreign_schema(import_foreign_schema_stmt, server_oid) {
+    for command in create_stmts {
         ret.push(command.as_pg_cstr());
     }
 

diff --git a/supabase-wrappers/src/instance.rs b/supabase-wrappers/src/instance.rs
@@ -2,13 +2,26 @@ use crate::prelude::*;
 use pgrx::pg_sys::panic::ErrorReport;
 use pgrx::prelude::*;
 
-// create a fdw instance
-pub(super) unsafe fn create_fdw_instance<E: Into<ErrorReport>, W: ForeignDataWrapper<E>>(
-    ftable_id: pg_sys::Oid,
+// create a fdw instance from its id
+pub(super) unsafe fn create_fdw_instance_from_server_id<
+    E: Into<ErrorReport>,
+    W: ForeignDataWrapper<E>,
+>(
+    fserver_id: pg_sys::Oid,
 ) -> W {
-    let ftable = pg_sys::GetForeignTable(ftable_id);
-    let fserver = pg_sys::GetForeignServer((*ftable).serverid);
+    let fserver = pg_sys::GetForeignServer(fserver_id);
     let fserver_opts = options_to_hashmap((*fserver).options).report_unwrap();
     let wrapper = W::new(&fserver_opts);
     wrapper.report_unwrap()
 }
+
+// create a fdw instance from a foreign table id
+pub(super) unsafe fn create_fdw_instance_from_table_id<
+    E: Into<ErrorReport>,
+    W: ForeignDataWrapper<E>,
+>(
+    ftable_id: pg_sys::Oid,
+) -> W {
+    let ftable = pg_sys::GetForeignTable(ftable_id);
+    create_fdw_instance_from_server_id((*ftable).serverid)
+}
diff --git a/supabase-wrappers/src/interface.rs b/supabase-wrappers/src/interface.rs
@@ -628,8 +628,8 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
     ///
     /// [See more details](https://www.postgresql.org/docs/current/fdw-callbacks.html#FDW-CALLBACKS-IMPORT).
     fn import_foreign_schema(
+        &mut self,
         _stmt: crate::import_foreign_schema::ImportForeignSchemaStmt,
-        _server_oid: pg_sys::Oid,
     ) -> Vec<String> {
         Vec::new()
     }

diff --git a/supabase-wrappers/src/modify.rs b/supabase-wrappers/src/modify.rs
@@ -40,7 +40,7 @@ struct FdwModifyState<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> {
 impl<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> FdwModifyState<E, W> {
     unsafe fn new(foreigntableid: Oid, tmp_ctx: PgMemoryContexts) -> Self {
         Self {
-            instance: instance::create_fdw_instance(foreigntableid),
+            instance: instance::create_fdw_instance_from_table_id(foreigntableid),
             rowid_name: String::default(),
             rowid_attno: 0,
             rowid_typid: Oid::INVALID,

diff --git a/supabase-wrappers/src/scan.rs b/supabase-wrappers/src/scan.rs
@@ -55,7 +55,7 @@ struct FdwState<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> {
 impl<E: Into<ErrorReport>, W: ForeignDataWrapper<E>> FdwState<E, W> {
     unsafe fn new(foreigntableid: Oid, tmp_ctx: PgMemoryContexts) -> Self {
         Self {
-            instance: instance::create_fdw_instance(foreigntableid),
+            instance: instance::create_fdw_instance_from_table_id(foreigntableid),
             quals: Vec::new(),
             tgts: Vec::new(),
             sorts: Vec::new(),