11use std:: { ffi:: CStr , str:: FromStr } ;
22
33use pgrx:: {
4- is_a,
4+ ereport , is_a,
55 pg_sys:: {
66 addRangeTableEntryForRelation, defGetInt32, defGetInt64, defGetString, get_namespace_name,
77 get_rel_namespace, makeDefElem, makeString, make_parsestate, quote_qualified_identifier,
88 AccessShareLock , AsPgCStr , CopyStmt , CreateTemplateTupleDesc , DefElem , List , NoLock , Node ,
99 NodeTag :: T_CopyStmt , Oid , ParseNamespaceItem , ParseState , PlannedStmt , QueryEnvironment ,
1010 RangeVar , RangeVarGetRelidExtended , RowExclusiveLock , TupleDescInitEntry ,
1111 } ,
12- PgBox , PgList , PgRelation , PgTupleDesc ,
12+ PgBox , PgList , PgLogLevel , PgRelation , PgSqlErrorCode , PgTupleDesc ,
1313} ;
1414use url:: Url ;
1515
@@ -300,7 +300,7 @@ pub(crate) fn copy_stmt_get_option(
300300 PgBox :: null ( )
301301}
302302
303- pub ( crate ) fn is_copy_to_parquet_stmt ( p_stmt : & PgBox < PlannedStmt > ) -> bool {
303+ fn is_copy_parquet_stmt ( p_stmt : & PgBox < PlannedStmt > , copy_from : bool ) -> bool {
304304 // the GUC pg_parquet.enable_copy_hook must be set to true
305305 if !ENABLE_PARQUET_COPY_HOOK . get ( ) {
306306 return false ;
@@ -314,7 +314,7 @@ pub(crate) fn is_copy_to_parquet_stmt(p_stmt: &PgBox<PlannedStmt>) -> bool {
314314
315315 let copy_stmt = unsafe { PgBox :: < CopyStmt > :: from_pg ( p_stmt. utilityStmt as _ ) } ;
316316
317- if copy_stmt. is_from {
317+ if copy_from != copy_stmt. is_from {
318318 return false ;
319319 }
320320
@@ -334,44 +334,35 @@ pub(crate) fn is_copy_to_parquet_stmt(p_stmt: &PgBox<PlannedStmt>) -> bool {
334334
335335 // extension checks are done via catalog (not yet searched via cache by postgres till pg18)
336336 // this is why we check them after the uri checks
337- extension_exists ( "pg_parquet" ) && !extension_exists ( "crunchy_query_engine" )
338- }
339-
340- pub ( crate ) fn is_copy_from_parquet_stmt ( p_stmt : & PgBox < PlannedStmt > ) -> bool {
341- // the GUC pg_parquet.enable_copy_hook must be set to true
342- if !ENABLE_PARQUET_COPY_HOOK . get ( ) {
343- return false ;
344- }
345337
346- let is_copy_stmt = unsafe { is_a ( p_stmt. utilityStmt , T_CopyStmt ) } ;
338+ // pg_parquet should be created
339+ if !extension_exists ( "pg_parquet" ) {
340+ ereport ! (
341+ PgLogLevel :: WARNING ,
342+ PgSqlErrorCode :: ERRCODE_WARNING ,
343+ "pg_parquet can handle this COPY command but is not enabled" ,
344+ "Run CREATE EXTENSION pg_parquet; to enable the pg_parquet extension." ,
345+ ) ;
347346
348- if !is_copy_stmt {
349347 return false ;
350348 }
351349
352- let copy_stmt = unsafe { PgBox :: < CopyStmt > :: from_pg ( p_stmt. utilityStmt as _ ) } ;
353-
354- if !copy_stmt. is_from {
350+ // crunchy_query_engine should not be created
351+ if extension_exists ( "crunchy_query_engine" ) {
355352 return false ;
356353 }
357354
358- if copy_stmt. is_program {
359- return false ;
360- }
361-
362- if copy_stmt. filename . is_null ( ) {
363- return false ;
364- }
365-
366- let uri = copy_stmt_uri ( p_stmt) . expect ( "uri is None" ) ;
355+ true
356+ }
367357
368- if !is_parquet_format_option ( p_stmt) && !is_parquet_uri ( uri) {
369- return false ;
370- }
358+ pub ( crate ) fn is_copy_to_parquet_stmt ( p_stmt : & PgBox < PlannedStmt > ) -> bool {
359+ let copy_from = false ;
360+ is_copy_parquet_stmt ( p_stmt, copy_from)
361+ }
371362
372- // extension checks are done via catalog (not yet searched via cache by postgres till pg18)
373- // this is why we check them after the uri checks
374- extension_exists ( "pg_parquet" ) && ! extension_exists ( "crunchy_query_engine" )
363+ pub ( crate ) fn is_copy_from_parquet_stmt ( p_stmt : & PgBox < PlannedStmt > ) -> bool {
364+ let copy_from = true ;
365+ is_copy_parquet_stmt ( p_stmt , copy_from )
375366}
376367
377368fn is_parquet_format_option ( p_stmt : & PgBox < PlannedStmt > ) -> bool {
0 commit comments