Skip to content

Commit

Permalink
chore
Browse files Browse the repository at this point in the history
  • Loading branch information
Weijun-H committed Oct 27, 2024
1 parent 5935635 commit 795dc23
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 26 deletions.
30 changes: 16 additions & 14 deletions src/hooks/utility.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"),
};

Expand All @@ -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
Expand Down
23 changes: 12 additions & 11 deletions src/hooks/utility/view.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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(),
Expand Down
3 changes: 2 additions & 1 deletion tests/tests/scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down

0 comments on commit 795dc23

Please sign in to comment.