From 80a5a5aef2cdcb48dec7864486f38323b34d77ce Mon Sep 17 00:00:00 2001
From: Weijun-H <huangweijun1001@gmail.com>
Date: Fri, 25 Oct 2024 21:56:37 +0800
Subject: [PATCH] chore

---
 src/hooks/utility.rs      | 30 ++++++++++++++++--------------
 src/hooks/utility/view.rs | 23 ++++++++++++-----------
 tests/tests/scan.rs       |  3 ++-
 3 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/src/hooks/utility.rs b/src/hooks/utility.rs
index 4cb69e14..b42ef984 100644
--- a/src/hooks/utility.rs
+++ b/src/hooks/utility.rs
@@ -23,18 +23,14 @@ mod view;
 
 use std::ptr::null_mut;
 
+use super::query::*;
 use anyhow::{bail, Result};
+use explain::explain_query;
 use pgrx::{pg_sys, AllocatedByRust, HookResult, PgBox};
+use prepare::*;
 use sqlparser::{ast::Statement, dialect::PostgreSqlDialect, parser::Parser};
 use view::view_query;
 
-use pg_sys::NodeTag;
-
-use explain::explain_query;
-use prepare::*;
-
-use super::query::*;
-
 type ProcessUtilityHook = fn(
     pstmt: PgBox<pg_sys::PlannedStmt>,
     query_string: &core::ffi::CStr,
@@ -120,12 +116,18 @@ pub async fn process_utility_hook(
             pstmt.utilityStmt as *mut pg_sys::ExplainStmt,
             dest.as_ptr(),
         )?,
-        pg_sys::NodeTag::T_ViewStmt => view_query(
-            query_string,
-            pstmt.utilityStmt as *mut pg_sys::ViewStmt,
-            pstmt.stmt_location,
-            pstmt.stmt_len,
-        )?,
+        pg_sys::NodeTag::T_ViewStmt => {
+            let utility_stmt = unsafe {
+                pg_sys::copyObjectImpl(pstmt.utilityStmt as *const std::ffi::c_void)
+                    as *mut pg_sys::Node
+            };
+            view_query(
+                query_string,
+                utility_stmt as *mut pg_sys::ViewStmt,
+                pstmt.stmt_location,
+                pstmt.stmt_len,
+            )?
+        }
         _ => bail!("unexpected statement type in utility hook"),
     };
 
@@ -145,7 +147,7 @@ pub async fn process_utility_hook(
     Ok(())
 }
 
-fn is_support_utility(stmt_type: NodeTag) -> bool {
+fn is_support_utility(stmt_type: pg_sys::NodeTag) -> bool {
     stmt_type == pg_sys::NodeTag::T_ExplainStmt
         || stmt_type == pg_sys::NodeTag::T_ViewStmt
         || stmt_type == pg_sys::NodeTag::T_PrepareStmt
diff --git a/src/hooks/utility/view.rs b/src/hooks/utility/view.rs
index 5037b5da..3138faed 100644
--- a/src/hooks/utility/view.rs
+++ b/src/hooks/utility/view.rs
@@ -19,7 +19,10 @@ use std::ptr::null_mut;
 
 use anyhow::Result;
 
-use pgrx::{pg_sys, warning};
+use pgrx::{
+    pg_sys::{self},
+    warning,
+};
 
 use crate::{duckdb::connection::execute, hooks::query::is_duckdb_query};
 
@@ -32,18 +35,16 @@ pub fn view_query(
     stmt_len: i32,
 ) -> Result<bool> {
     // Perform parsing and analysis to get the Query
-    let query_list = unsafe {
-        let mut raw_stmt = pg_sys::RawStmt {
-            type_: pg_sys::NodeTag::T_RawStmt,
-            stmt: (*stmt).query,
-            stmt_location,
-            stmt_len,
-        };
+    let rewritten_queries = unsafe {
+        let mut raw_stmt = pgrx::PgBox::<pg_sys::RawStmt>::alloc_node(pg_sys::NodeTag::T_RawStmt);
+        raw_stmt.stmt = (*stmt).query;
+        raw_stmt.stmt_location = stmt_location;
+        raw_stmt.stmt_len = stmt_len;
 
         #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))]
         {
             pg_sys::pg_analyze_and_rewrite_fixedparams(
-                &mut raw_stmt,
+                raw_stmt.as_ptr(),
                 query_string.as_ptr(),
                 null_mut(),
                 0,
@@ -54,7 +55,7 @@ pub fn view_query(
         #[cfg(any(feature = "pg13", feature = "pg14"))]
         {
             pg_sys::pg_analyze_and_rewrite(
-                &mut raw_stmt,
+                raw_stmt.as_ptr(),
                 query_string.as_ptr(),
                 null_mut(),
                 0,
@@ -65,7 +66,7 @@ pub fn view_query(
 
     let plan_list = unsafe {
         pg_sys::pg_plan_queries(
-            query_list,
+            rewritten_queries,
             query_string.as_ptr(),
             pg_sys::CURSOR_OPT_PARALLEL_OK as i32,
             null_mut(),
diff --git a/tests/tests/scan.rs b/tests/tests/scan.rs
index 2835fdf6..8e3453fa 100644
--- a/tests/tests/scan.rs
+++ b/tests/tests/scan.rs
@@ -672,11 +672,12 @@ async fn test_view_foreign_table(mut conn: PgConnection, tempdir: TempDir) -> Re
     // cannot fully pushdown to the DuckDB
     "CREATE TABLE t1 (a int);".execute(&mut conn);
     "INSERT INTO t1 VALUES (1);".execute(&mut conn);
+
     r#"
     CREATE VIEW primitive_join_view AS
     SELECT *
     FROM primitive
-    JOIN t1 ON t1.a = primitive.int32_col
+    JOIN t1 ON t1.a = primitive.int32_col;
     "#
     .execute(&mut conn);