@@ -24,12 +24,16 @@ afterEach(() => {
2424 Registry . reset ( )
2525} )
2626
27- function postgres ( relkind = "r" , reltuples = 5 , translate = ( sql : string ) => sql ) {
27+ function postgres (
28+ relkind = "r" ,
29+ reltuples = 5 ,
30+ translate = ( sql : string , _options ?: Parameters < Connector [ "execute" ] > [ 3 ] ) => sql ,
31+ ) {
2832 database . run ( "INSERT INTO pg_class VALUES (1, ?, ?)" , [ relkind , reltuples ] )
29- const execute = mock ( async ( ...[ sql ] : Parameters < Connector [ "execute" ] > ) => {
33+ const execute = mock ( async ( ...[ sql , , , options ] : Parameters < Connector [ "execute" ] > ) => {
3034 // Keep the generated WHERE and FILTER clauses executable: removing either
3135 // guard condition must change results, rather than leave a canned mock row.
32- const query = translate ( sql )
36+ const query = translate ( sql , options )
3337 . replace ( "pg_catalog.pg_class" , "pg_class" )
3438 . replace ( / p g _ c a t a l o g \. t o _ r e g c l a s s \( E ' (?: [ ^ ' ] | ' ' ) * ' \) / , "1" )
3539 const statement = database . query ( query )
@@ -41,6 +45,96 @@ function postgres(relkind = "r", reltuples = 5, translate = (sql: string) => sql
4145}
4246
4347describe ( "predicate_effect with mocked Postgres" , ( ) => {
48+ test . each ( [
49+ { name : "predicates" , base_sql : "orders" , predicates : [ "custom_write_udf() > 0" ] as string [ ] , calls : 2 } ,
50+ {
51+ name : "populations" ,
52+ base_sql : "SELECT * FROM orders WHERE custom_write_udf() > 0" ,
53+ predicates : [ "id > 0" ] as string [ ] ,
54+ calls : 1 ,
55+ } ,
56+ ] ) ( "uses read-only execution to reject writing UDFs in $name" , async ( { base_sql, predicates, calls } ) => {
57+ let writes = 0
58+ const execute = postgres ( "r" , 5 , ( sql , options ) => {
59+ if ( ! sql . includes ( "custom_write_udf()" ) ) return sql
60+ // Model the connector's transaction boundary for a function unknown to the denylist.
61+ if ( options ?. readOnly ) throw new Error ( "cannot execute INSERT in a read-only transaction" )
62+ writes ++
63+ return sql . replaceAll ( "custom_write_udf()" , "1" )
64+ } )
65+ await rejects (
66+ checkSql ( { warehouse, check : "predicate_effect" , base_sql, predicates, query_timeout_ms : 4321 } ) ,
67+ / r e a d - o n l y t r a n s a c t i o n / ,
68+ )
69+ expect ( writes ) . toBe ( 0 )
70+ expect ( execute ) . toHaveBeenCalledTimes ( calls )
71+ for ( const call of execute . mock . calls ) {
72+ expect ( call [ 3 ] ) . toMatchObject ( { noLimit : true , queryTimeoutMs : 4321 , readOnly : true } )
73+ }
74+ } )
75+
76+ for ( const [ kind , calls ] of [
77+ [
78+ "server-control" ,
79+ [
80+ "pg_promote()" ,
81+ 'pg_catalog."pg_promote" /* promote */ ()' ,
82+ "PG_RELOAD_CONF()" ,
83+ 'pg_catalog."pg_reload_conf" -- reload\n()' ,
84+ ] ,
85+ ] ,
86+ [
87+ "dblink-write" ,
88+ [
89+ "dblink_send_query('remote', 'DELETE FROM important_table')" ,
90+ `public."dblink_send_query" /* send */ ('remote', 'DELETE FROM important_table')` ,
91+ "DBLINK_CONNECT('remote', 'dbname=remote')" ,
92+ `public."dblink_connect" -- connect\n('remote', 'dbname=remote')` ,
93+ "dblink_connect_u('remote', 'dbname=remote')" ,
94+ `public."dblink_connect_u" /* connect */ ('remote', 'dbname=remote')` ,
95+ ] ,
96+ ] ,
97+ ] satisfies [ string , string [ ] ] [ ] ) {
98+ test . each ( calls ) ( `rejects ${ kind } predicate %s before query execution` , async ( call ) => {
99+ const predicate = `${ call } IS NOT NULL`
100+ const execute = postgres ( "r" , 5 , ( sql ) => sql . replaceAll ( call , "1" ) )
101+ await rejects (
102+ checkSql ( { warehouse, check : "predicate_effect" , base_sql : "orders" , predicates : [ predicate ] } ) ,
103+ / s q l _ c h e c k o n l y a c c e p t s a s i n g l e r e a d - o n l y S E L E C T / ,
104+ )
105+ expect ( execute ) . not . toHaveBeenCalled ( )
106+ } )
107+
108+ test . each ( calls ) ( `rejects ${ kind } population %s before query execution` , async ( call ) => {
109+ const execute = postgres ( "r" , 5 , ( sql ) => sql . replaceAll ( call , "1" ) )
110+ await rejects (
111+ checkSql ( {
112+ warehouse,
113+ check : "predicate_effect" ,
114+ base_sql : `SELECT ${ call } AS value FROM orders` ,
115+ predicates : [ "value IS NOT NULL" ] ,
116+ } ) ,
117+ / s q l _ c h e c k o n l y a c c e p t s a s i n g l e r e a d - o n l y S E L E C T / ,
118+ )
119+ expect ( execute ) . not . toHaveBeenCalled ( )
120+ } )
121+ }
122+
123+ test . each ( [ "pg_promote" , "pg_reload_conf" , "dblink_send_query" , "dblink_connect" , "dblink_connect_u" ] ) (
124+ "accepts %s in literals, comments and column identifiers" ,
125+ async ( name ) => {
126+ const execute = postgres ( )
127+ const result = await checkSql ( {
128+ warehouse,
129+ check : "predicate_effect" ,
130+ base_sql : `SELECT id AS "${ name } " FROM orders` ,
131+ predicates : [ `"${ name } " > 2 -- ${ name } ()` , `length('${ name } ()') > 0 /* ${ name } () */` ] ,
132+ } )
133+ expect ( result . metadata ) . toMatchObject ( { total_rows : 5 , combined_kept_rows : 3 } )
134+ expect ( execute ) . toHaveBeenCalledTimes ( 1 )
135+ } ,
136+ )
137+
44138 test . each ( [
45139 "pg_terminate_backend(id)" ,
46140 "PG_TERMINATE_BACKEND(id)" ,
0 commit comments