fix: memory context release error when drop fdw modify state (supabase#536)

* fix: fdw modify state drop issue

* Update supabase-wrappers/src/memctx.rs

Co-authored-by: Copilot <[email protected]>

* add bigquery fdw test case for unsupoorted type

---------

Co-authored-by: Copilot <[email protected]>

Author: burmecia

supabase-wrappers/src/modify.rs

@@ -369,15 +369,15 @@ pub(super) extern "C-unwind" fn exec_foreign_insert<
             (*rinfo).ri_FdwState as *mut FdwModifyState<E, W>,
         );
 
-        PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
+        let result = PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
             let row = utils::tuple_table_slot_to_row(slot);
-            let result = state.insert(&row);
-            if result.is_err() {
-                drop_fdw_modify_state(state.as_ptr());
-                (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
-                result.report_unwrap();
-            }
+            state.insert(&row)
         });
+        if result.is_err() {
+            drop_fdw_modify_state(state.as_ptr());
+            (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
+            result.report_unwrap();
+        }
     }
 
     slot
@@ -408,17 +408,19 @@ pub(super) extern "C-unwind" fn exec_foreign_delete<
             (*rinfo).ri_FdwState as *mut FdwModifyState<E, W>,
         );
 
-        PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
+        let result = PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
             let cell = get_rowid_cell(&state, plan_slot);
             if let Some(rowid) = cell {
-                let result = state.delete(&rowid);
-                if result.is_err() {
-                    drop_fdw_modify_state(state.as_ptr());
-                    (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
-                    result.report_unwrap();
-                }
+                state.delete(&rowid)
+            } else {
+                Ok(())
             }
         });
+        if result.is_err() {
+            drop_fdw_modify_state(state.as_ptr());
+            (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
+            result.report_unwrap();
+        }
     }
 
     slot
@@ -440,7 +442,7 @@ pub(super) extern "C-unwind" fn exec_foreign_update<
             (*rinfo).ri_FdwState as *mut FdwModifyState<E, W>,
         );
 
-        PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
+        let result = PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| {
             let rowid_cell = get_rowid_cell(&state, plan_slot);
             if let Some(rowid) = rowid_cell {
                 let mut new_row = utils::tuple_table_slot_to_row(plan_slot);
@@ -465,14 +467,16 @@ pub(super) extern "C-unwind" fn exec_foreign_update<
                     }
                 });
 
-                let result = state.update(&rowid, &new_row);
-                if result.is_err() {
-                    drop_fdw_modify_state(state.as_ptr());
-                    (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
-                    result.report_unwrap();
-                }
+                state.update(&rowid, &new_row)
+            } else {
+                Ok(())
             }
         });
+        if result.is_err() {
+            drop_fdw_modify_state(state.as_ptr());
+            (*rinfo).ri_FdwState = ptr::null::<FdwModifyState<E, W>>() as _;
+            result.report_unwrap();
+        }
     }
 
     slot

wrappers/dockerfiles/bigquery/data.yaml

@@ -56,3 +56,13 @@ projects:
               attrs: '{"aa": 123.45}'
               signup_dt: "2023-10-21"
               created_at: "2023-10-21T00:00:00"
+        - id: test_unsupported_type
+          columns:
+            - name: id
+              type: INTEGER
+              mode: REQUIRED
+            - name: time_col
+              type: TIME
+          data:
+            - id: 1
+              time_col: "12:30:00.45"

wrappers/src/fdw/bigquery_fdw/tests.rs

@@ -152,4 +152,53 @@ mod tests {
             assert_eq!(results, vec!["foo", "bar"]);
         });
     }
+
+    #[pg_test]
+    #[should_panic]
+    fn bigquery_unsupported_type() {
+        Spi::connect_mut(|c| {
+            c.update(
+                r#"CREATE FOREIGN DATA WRAPPER bigquery_wrapper
+                         HANDLER big_query_fdw_handler VALIDATOR big_query_fdw_validator"#,
+                None,
+                &[],
+            )
+            .unwrap();
+            c.update(
+                r#"CREATE SERVER my_bigquery_server
+                         FOREIGN DATA WRAPPER bigquery_wrapper
+                         OPTIONS (
+                           project_id 'test_project',
+                           dataset_id 'test_dataset',
+                           api_endpoint 'http://localhost:9111',
+                           mock_auth 'true'
+                         )"#,
+                None,
+                &[],
+            )
+            .unwrap();
+            c.update(
+                r#"
+                  CREATE FOREIGN TABLE test_unsupported_type (
+                    id bigint,
+                    time_col time
+                  )
+                  SERVER my_bigquery_server
+                  OPTIONS (
+                    table 'test_unsupported_type',
+                    rowid_column 'id'
+                  )
+             "#,
+                None,
+                &[],
+            )
+            .unwrap();
+
+            let _results = c.update(
+                "INSERT INTO test_unsupported_type (id, time_col) VALUES (42, '12:30:00.45')",
+                None,
+                &[],
+            );
+        });
+    }
 }