Branch data Line data Source code
1 : : /*-------------------------------------------------------------------------
2 : : *
3 : : * test_extensible.c
4 : : * Tests for extensible nodes and custom scans
5 : : *
6 : : * Copyright (c) 2026, PostgreSQL Global Development Group
7 : : *
8 : : * IDENTIFICATION
9 : : * src/test/modules/test_extensible/test_extensible.c
10 : : *
11 : : *-------------------------------------------------------------------------
12 : : */
13 : :
14 : : #include "postgres.h"
15 : :
16 : : #include "access/table.h"
17 : : #include "access/tableam.h"
18 : : #include "catalog/namespace.h"
19 : : #include "executor/executor.h"
20 : : #include "fmgr.h"
21 : : #include "miscadmin.h"
22 : : #include "nodes/extensible.h"
23 : : #include "nodes/nodes.h"
24 : : #include "nodes/plannodes.h"
25 : : #include "nodes/readfuncs.h"
26 : : #include "optimizer/pathnode.h"
27 : : #include "optimizer/paths.h"
28 : : #include "optimizer/restrictinfo.h"
29 : : #include "utils/builtins.h"
30 : : #include "utils/guc.h"
31 : : #include "utils/lsyscache.h"
32 : :
33 : 2 : PG_MODULE_MAGIC;
34 : :
35 : : /* Name of the test table that triggers our CustomScan injection */
36 : : #define TEST_TABLE_NAME "test_extensible_tbl"
37 : :
38 : : /*
39 : : * TestExtNode - an ExtensibleNode subtype carrying our planning data.
40 : : */
41 : : typedef struct TestExtNode
42 : : {
43 : : ExtensibleNode base; /* must be first */
44 : : Oid relid; /* OID of the relation being scanned */
45 : : int repeat_count; /* how many times to return each scanned row */
46 : : } TestExtNode;
47 : :
48 : : #define TEST_EXT_NODE_NAME "TestExtNode"
49 : : #define TEST_CUSTOM_SCAN_NAME "TestCustomScan"
50 : :
51 : : /* GUC: how many times the custom scan returns each row */
52 : : static int test_repeat_count = 2;
53 : :
54 : : static TestExtNode *text_to_test_ext_node(text *txt);
55 : :
56 : : /* Encode a TestExtNode into its serialized representation */
57 : : #define TEST_EXT_NODE_TO_TEXT(node) cstring_to_text(nodeToString(node))
58 : :
59 : : /* Decode a function argument back into a TestExtNode */
60 : : #define PG_GETARG_TEST_EXT_NODE(n) text_to_test_ext_node(PG_GETARG_TEXT_PP(n))
61 : :
62 : 2 : PG_FUNCTION_INFO_V1(test_get_extensible_node_methods);
63 : 2 : PG_FUNCTION_INFO_V1(test_get_custom_scan_methods);
64 : 2 : PG_FUNCTION_INFO_V1(test_ext_node_make);
65 : 2 : PG_FUNCTION_INFO_V1(test_ext_node_copy);
66 : 2 : PG_FUNCTION_INFO_V1(test_ext_node_equal);
67 : 2 : PG_FUNCTION_INFO_V1(test_ext_node_get_relid);
68 : 2 : PG_FUNCTION_INFO_V1(test_ext_node_get_repeat_count);
69 : :
70 : : /*
71 : : * ExtensibleNodeMethods callbacks.
72 : : *
73 : : * Note that nodeOut() and nodeRead() must agree on the set and order of
74 : : * serialized fields.
75 : : */
76 : : static void
77 : 13 : test_ext_node_copy_cb(ExtensibleNode *newnode, const ExtensibleNode *oldnode)
78 : : {
79 : 13 : ((TestExtNode *) newnode)->relid = ((const TestExtNode *) oldnode)->relid;
80 : 13 : ((TestExtNode *) newnode)->repeat_count =
81 : 13 : ((const TestExtNode *) oldnode)->repeat_count;
82 : 13 : }
83 : :
84 : : static bool
85 : 3 : test_ext_node_equal_cb(const ExtensibleNode *a, const ExtensibleNode *b)
86 : : {
87 : 3 : return ((const TestExtNode *) a)->relid ==
88 [ + + ]: 5 : ((const TestExtNode *) b)->relid &&
89 : 2 : ((const TestExtNode *) a)->repeat_count ==
90 [ + + ]: 2 : ((const TestExtNode *) b)->repeat_count;
91 : : }
92 : :
93 : : static void
94 : 20 : test_ext_node_out_cb(StringInfo str, const ExtensibleNode *node)
95 : : {
96 : 20 : appendStringInfo(str, " :relid %u", ((const TestExtNode *) node)->relid);
97 : 20 : appendStringInfo(str, " :repeat_count %d",
98 : 20 : ((const TestExtNode *) node)->repeat_count);
99 : 20 : }
100 : :
101 : : /*
102 : : * Fetch the next token, erroring out instead of returning NULL.
103 : : *
104 : : * Unlike anything in readfuncs.c, this callback is reachable with arbitrary
105 : : * strings through SQL function calls, so we need this check.
106 : : */
107 : : static const char *
108 : 86 : test_ext_node_next_token(ReadNodeContext *ctx)
109 : : {
110 : : int length;
111 : 86 : const char *token = pg_strtok(ctx, &length);
112 : :
113 [ + + ]: 86 : if (token == NULL)
114 [ + - ]: 1 : ereport(ERROR,
115 : : (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
116 : : errmsg("unexpected end of \"%s\"", TEST_EXT_NODE_NAME)));
117 : :
118 : 85 : return token;
119 : : }
120 : :
121 : : static void
122 : 22 : test_ext_node_read_cb(ReadNodeContext *ctx, ExtensibleNode *node)
123 : : {
124 : 22 : TestExtNode *tnode = (TestExtNode *) node;
125 : :
126 : 22 : (void) test_ext_node_next_token(ctx); /* skip :relid */
127 : 22 : tnode->relid = atooid(test_ext_node_next_token(ctx));
128 : :
129 : 21 : (void) test_ext_node_next_token(ctx); /* skip :repeat_count */
130 : 21 : tnode->repeat_count = atoi(test_ext_node_next_token(ctx));
131 : 21 : }
132 : :
133 : : static const ExtensibleNodeMethods test_ext_node_methods =
134 : : {
135 : : .extnodename = TEST_EXT_NODE_NAME,
136 : : .node_size = sizeof(TestExtNode),
137 : : .nodeCopy = test_ext_node_copy_cb,
138 : : .nodeEqual = test_ext_node_equal_cb,
139 : : .nodeOut = test_ext_node_out_cb,
140 : : .nodeRead = test_ext_node_read_cb,
141 : : };
142 : :
143 : : /*
144 : : * TestCustomScanState - execution state for the custom scan
145 : : */
146 : : typedef struct TestCustomScanState
147 : : {
148 : : CustomScanState css; /* must be first */
149 : : TableScanDesc scandesc;
150 : : int repeat_count; /* repeat_count from TestExtNode */
151 : : int repeats_left; /* how many more times to return current row */
152 : : } TestCustomScanState;
153 : :
154 : : /*
155 : : * Executor callbacks
156 : : */
157 : :
158 : : /*
159 : : * Retrieve our private planning data from a CustomScan node.
160 : : */
161 : : static TestExtNode *
162 : 22 : test_get_ext_node(CustomScan *cscan)
163 : : {
164 : : TestExtNode *tnode;
165 : :
166 : : Assert(list_length(cscan->custom_private) == 1);
167 : 22 : tnode = (TestExtNode *) linitial(cscan->custom_private);
168 : :
169 : : Assert(IsA(tnode, ExtensibleNode));
170 : : Assert(strcmp(tnode->base.extnodename, TEST_EXT_NODE_NAME) == 0);
171 : :
172 : 22 : return tnode;
173 : : }
174 : :
175 : : /*
176 : : * BeginCustomScan.
177 : : */
178 : : static void
179 : 11 : test_begin_custom_scan(CustomScanState *node, EState *estate, int eflags)
180 : : {
181 : 11 : TestCustomScanState *tstate = (TestCustomScanState *) node;
182 : 11 : TestExtNode *tnode = test_get_ext_node((CustomScan *) node->ss.ps.plan);
183 : 11 : Relation rel = node->ss.ss_currentRelation;
184 : :
185 : : Assert(tnode->repeat_count > 0);
186 : 11 : tstate->repeat_count = tnode->repeat_count;
187 : 11 : tstate->repeats_left = 0;
188 : 11 : tstate->scandesc = NULL;
189 : :
190 : : /* A plain EXPLAIN never executes the plan, so skip */
191 [ + + ]: 11 : if (eflags & EXEC_FLAG_EXPLAIN_ONLY)
192 : 5 : return;
193 : :
194 : 6 : tstate->scandesc = table_beginscan(rel, estate->es_snapshot, 0, NULL,
195 : : SO_NONE);
196 : : }
197 : :
198 : : /*
199 : : * Access method for ExecScan(): return the next tuple to be considered, or
200 : : * NULL when the scan is done.
201 : : */
202 : : static TupleTableSlot *
203 : 38 : test_scan_next(ScanState *node)
204 : : {
205 : 38 : TestCustomScanState *tstate = (TestCustomScanState *) node;
206 : 38 : TupleTableSlot *slot = node->ss_ScanTupleSlot;
207 : :
208 : : /* Return the current tuple again if it still has repeats left */
209 [ + + ]: 38 : if (tstate->repeats_left > 0)
210 : : {
211 : 18 : tstate->repeats_left--;
212 : 18 : return slot;
213 : : }
214 : :
215 [ + + ]: 20 : if (!table_scan_getnextslot(tstate->scandesc, ForwardScanDirection, slot))
216 : 5 : return NULL;
217 : :
218 : 15 : tstate->repeats_left = tstate->repeat_count - 1;
219 : 15 : return slot;
220 : : }
221 : :
222 : : /*
223 : : * Recheck method for ExecScan(), used only during EvalPlanQual rechecks.
224 : : * We evaluate no quals of our own, so there is nothing to recheck.
225 : : */
226 : : static bool
227 : 0 : test_scan_recheck(ScanState *node, TupleTableSlot *slot)
228 : : {
229 : 0 : return true;
230 : : }
231 : :
232 : : static TupleTableSlot *
233 : 36 : test_exec_custom_scan(CustomScanState *node)
234 : : {
235 : 36 : return ExecScan(&node->ss, test_scan_next, test_scan_recheck);
236 : : }
237 : :
238 : : static void
239 : 11 : test_end_custom_scan(CustomScanState *node)
240 : : {
241 : 11 : TestCustomScanState *tstate = (TestCustomScanState *) node;
242 : :
243 : : /* No scan started under EXEC_FLAG_EXPLAIN_ONLY */
244 [ + + ]: 11 : if (tstate->scandesc != NULL)
245 : 6 : table_endscan(tstate->scandesc);
246 : 11 : }
247 : :
248 : : /*
249 : : * ReScanCustomScan: reset our own state as well as ExecScan()'s.
250 : : */
251 : : static void
252 : 0 : test_rescan_custom_scan(CustomScanState *node)
253 : : {
254 : 0 : TestCustomScanState *tstate = (TestCustomScanState *) node;
255 : :
256 : 0 : tstate->repeats_left = 0;
257 : 0 : table_rescan(tstate->scandesc, NULL);
258 : 0 : ExecScanReScan(&node->ss);
259 : 0 : }
260 : :
261 : : static const CustomExecMethods test_custom_exec_methods =
262 : : {
263 : : .CustomName = TEST_CUSTOM_SCAN_NAME,
264 : : .BeginCustomScan = test_begin_custom_scan,
265 : : .ExecCustomScan = test_exec_custom_scan,
266 : : .EndCustomScan = test_end_custom_scan,
267 : : .ReScanCustomScan = test_rescan_custom_scan,
268 : : };
269 : :
270 : : /*
271 : : * CreateCustomScanState() allocates the CustomScanState and fills in its node
272 : : * tag and its methods; everything else is left to ExecInitCustomScan().
273 : : *
274 : : * slotOps is set here because ExecInitCustomScan() reads it before creating
275 : : * the scan tuple slot, and before it opens the scan relation itself. So we
276 : : * open the relation ourselves (its OID travels via our ExtensibleNode) just
277 : : * to learn its slot type from table_slot_callbacks().
278 : : *
279 : : * A lock is taken rather than passing NoLock, since in a parallel worker
280 : : * nothing has locked the relation yet; the worker takes its own lock later
281 : : * in ExecGetRangeTableRelation(). AccessShareLock matches what a plain scan
282 : : * uses.
283 : : */
284 : : static Node *
285 : 11 : test_create_custom_scan_state(CustomScan *cscan)
286 : : {
287 : : TestCustomScanState *tstate;
288 : 11 : TestExtNode *tnode = test_get_ext_node(cscan);
289 : : Relation rel;
290 : :
291 : : tstate = (TestCustomScanState *)
292 : 11 : newNode(sizeof(TestCustomScanState), T_CustomScanState);
293 : 11 : tstate->css.methods = &test_custom_exec_methods;
294 : :
295 : 11 : rel = table_open(tnode->relid, AccessShareLock);
296 : 11 : tstate->css.slotOps = table_slot_callbacks(rel);
297 : 11 : table_close(rel, NoLock);
298 : :
299 : 11 : return (Node *) tstate;
300 : : }
301 : :
302 : : static const CustomScanMethods test_custom_scan_methods =
303 : : {
304 : : .CustomName = TEST_CUSTOM_SCAN_NAME,
305 : : .CreateCustomScanState = test_create_custom_scan_state,
306 : : };
307 : :
308 : : /*
309 : : * Planner callbacks
310 : : */
311 : :
312 : : /*
313 : : * PlanCustomPath turns our CustomPath into the CustomScan plan node that the
314 : : * executor runs.
315 : : */
316 : : static Plan *
317 : 10 : test_plan_custom_path(PlannerInfo *root,
318 : : RelOptInfo *rel,
319 : : struct CustomPath *best_path,
320 : : List *tlist,
321 : : List *clauses,
322 : : List *custom_plans)
323 : : {
324 : 10 : CustomScan *cscan = makeNode(CustomScan);
325 : :
326 : 10 : cscan->scan.plan.targetlist = tlist;
327 : :
328 : : /*
329 : : * Restriction clauses arrive as RestrictInfos; reduce them to bare
330 : : * expressions for the plan's qual. Pseudoconstants are dropped since the
331 : : * core evaluates them in a gating Result node above us.
332 : : */
333 : 10 : cscan->scan.plan.qual = extract_actual_clauses(clauses, false);
334 : 10 : cscan->scan.scanrelid = rel->relid;
335 : 10 : cscan->flags = best_path->flags;
336 : :
337 : : /* Our CustomPath has no child paths */
338 : 10 : cscan->custom_plans = custom_plans;
339 : 10 : cscan->custom_exprs = NIL;
340 : :
341 : : /* Pass the ExtensibleNode from the path to the plan */
342 : 10 : cscan->custom_private = best_path->custom_private;
343 : 10 : cscan->custom_scan_tlist = NIL;
344 : 10 : cscan->custom_relids = NULL;
345 : 10 : cscan->methods = &test_custom_scan_methods;
346 : :
347 : 10 : return (Plan *) cscan;
348 : : }
349 : :
350 : : static const CustomPathMethods test_custom_path_methods =
351 : : {
352 : : .CustomName = TEST_CUSTOM_SCAN_NAME,
353 : : .PlanCustomPath = test_plan_custom_path,
354 : : };
355 : :
356 : : static set_rel_pathlist_hook_type prev_set_rel_pathlist_hook = NULL;
357 : :
358 : : static void
359 : 21 : test_set_rel_pathlist(PlannerInfo *root, RelOptInfo *rel,
360 : : Index rti, RangeTblEntry *rte)
361 : : {
362 : : CustomPath *cpath;
363 : : TestExtNode *tnode;
364 : : char *relname;
365 : :
366 [ - + ]: 21 : if (prev_set_rel_pathlist_hook)
367 : 0 : prev_set_rel_pathlist_hook(root, rel, rti, rte);
368 : :
369 : : /*
370 : : * Only consider plain base relations with a table AM, no inheritance and
371 : : * no tablesample, keeping the logic simple.
372 : : */
373 [ + + + + ]: 21 : if (rel->reloptkind != RELOPT_BASEREL || rte->rtekind != RTE_RELATION)
374 : 7 : return;
375 [ - + - - ]: 14 : if (rte->relkind != RELKIND_RELATION && rte->relkind != RELKIND_MATVIEW)
376 : 0 : return;
377 [ + + + + ]: 14 : if (rte->tablesample != NULL || rte->inh)
378 : 4 : return;
379 : :
380 : : /*
381 : : * Only inject our CustomPath for the specific marker table, free of
382 : : * namespace.
383 : : */
384 : 10 : relname = get_rel_name(rte->relid);
385 [ + - - + ]: 10 : if (relname == NULL || strcmp(relname, TEST_TABLE_NAME) != 0)
386 : 0 : return;
387 : :
388 : : /* TestExtNode for the executor callbacks */
389 : 10 : tnode = (TestExtNode *) newNode(sizeof(TestExtNode), T_ExtensibleNode);
390 : 10 : tnode->base.extnodename = TEST_EXT_NODE_NAME;
391 : 10 : tnode->relid = rte->relid;
392 : :
393 : : /*
394 : : * Read once. Changes to the GUC do not affect already-planned queries.
395 : : * Each row is returned repeat_count times.
396 : : */
397 : 10 : tnode->repeat_count = test_repeat_count;
398 : :
399 : : /*
400 : : * Use a cost of zero to force our custom path; a real provider should
401 : : * estimate the cost honestly, but it does not matter for this module.
402 : : */
403 : 10 : cpath = makeNode(CustomPath);
404 : 10 : cpath->path.pathtype = T_CustomScan;
405 : 10 : cpath->path.parent = rel;
406 : 10 : cpath->path.pathtarget = rel->reltarget;
407 : 10 : cpath->path.rows = rel->rows * tnode->repeat_count;
408 : 10 : cpath->path.startup_cost = 0;
409 : 10 : cpath->path.total_cost = 0;
410 : :
411 : : /*
412 : : * Consider it as parallel safe, since our scan touches nothing but its
413 : : * own relation and keeps no state outside the CustomScanState.
414 : : */
415 : 10 : cpath->path.parallel_safe = rel->consider_parallel;
416 : :
417 : 10 : cpath->flags = 0;
418 : 10 : cpath->custom_paths = NIL;
419 : 10 : cpath->custom_private = list_make1(tnode);
420 : 10 : cpath->methods = &test_custom_path_methods;
421 : :
422 : 10 : add_path(rel, (Path *) cpath);
423 : : }
424 : :
425 : : /*
426 : : * test_get_extensible_node_methods
427 : : *
428 : : * Thin wrapper around GetExtensibleNodeMethods(). Returns the registered
429 : : * extnodename, or NULL when missing_ok = true and the name is not found.
430 : : */
431 : : Datum
432 : 3 : test_get_extensible_node_methods(PG_FUNCTION_ARGS)
433 : : {
434 : 3 : char *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
435 : 3 : bool missing_ok = PG_GETARG_BOOL(1);
436 : : const ExtensibleNodeMethods *methods;
437 : :
438 : 3 : methods = GetExtensibleNodeMethods(name, missing_ok);
439 [ + + ]: 2 : if (methods == NULL)
440 : 1 : PG_RETURN_NULL();
441 : :
442 : 1 : PG_RETURN_TEXT_P(cstring_to_text(methods->extnodename));
443 : : }
444 : :
445 : : /*
446 : : * test_get_custom_scan_methods
447 : : *
448 : : * Thin wrapper around GetCustomScanMethods(). Returns the registered
449 : : * CustomName, or NULL when missing_ok = true and the name is not found.
450 : : */
451 : : Datum
452 : 3 : test_get_custom_scan_methods(PG_FUNCTION_ARGS)
453 : : {
454 : 3 : char *name = text_to_cstring(PG_GETARG_TEXT_PP(0));
455 : 3 : bool missing_ok = PG_GETARG_BOOL(1);
456 : : const CustomScanMethods *methods;
457 : :
458 : 3 : methods = GetCustomScanMethods(name, missing_ok);
459 [ + + ]: 2 : if (methods == NULL)
460 : 1 : PG_RETURN_NULL();
461 : :
462 : 1 : PG_RETURN_TEXT_P(cstring_to_text(methods->CustomName));
463 : : }
464 : :
465 : : /*
466 : : * Decode a TestExtNode via stringToNode(), rejecting a string describing
467 : : * some other kind of node instead of misinterpreting it as one of ours.
468 : : */
469 : : static TestExtNode *
470 : 11 : text_to_test_ext_node(text *txt)
471 : : {
472 : 11 : Node *node = stringToNode(text_to_cstring(txt));
473 : :
474 [ + - + - ]: 10 : if (node == NULL || !IsA(node, ExtensibleNode) ||
475 [ - + ]: 10 : strcmp(((ExtensibleNode *) node)->extnodename, TEST_EXT_NODE_NAME) != 0)
476 [ # # ]: 0 : ereport(ERROR,
477 : : (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
478 : : errmsg("argument is not a serialized \"%s\"",
479 : : TEST_EXT_NODE_NAME)));
480 : :
481 : 10 : return (TestExtNode *) node;
482 : : }
483 : :
484 : : /*
485 : : * test_ext_node_make
486 : : *
487 : : * Builds a TestExtNode and returns it.
488 : : */
489 : : Datum
490 : 7 : test_ext_node_make(PG_FUNCTION_ARGS)
491 : : {
492 : : TestExtNode *tnode;
493 : :
494 : 7 : tnode = (TestExtNode *) newNode(sizeof(TestExtNode), T_ExtensibleNode);
495 : 7 : tnode->base.extnodename = TEST_EXT_NODE_NAME;
496 : 7 : tnode->relid = PG_GETARG_OID(0);
497 : 7 : tnode->repeat_count = PG_GETARG_INT32(1);
498 : :
499 : 7 : PG_RETURN_TEXT_P(TEST_EXT_NODE_TO_TEXT(tnode));
500 : : }
501 : :
502 : : /*
503 : : * test_ext_node_copy
504 : : */
505 : : Datum
506 : 2 : test_ext_node_copy(PG_FUNCTION_ARGS)
507 : : {
508 : 2 : TestExtNode *tnode = PG_GETARG_TEST_EXT_NODE(0);
509 : :
510 : 2 : PG_RETURN_TEXT_P(TEST_EXT_NODE_TO_TEXT(copyObject(tnode)));
511 : : }
512 : :
513 : : /*
514 : : * test_ext_node_equal
515 : : */
516 : : Datum
517 : 3 : test_ext_node_equal(PG_FUNCTION_ARGS)
518 : : {
519 : 3 : TestExtNode *a = PG_GETARG_TEST_EXT_NODE(0);
520 : 3 : TestExtNode *b = PG_GETARG_TEST_EXT_NODE(1);
521 : :
522 : 3 : PG_RETURN_BOOL(equal(a, b));
523 : : }
524 : :
525 : : /*
526 : : * test_ext_node_get_relid
527 : : * test_ext_node_get_repeat_count
528 : : *
529 : : * Field accessors for the custom node contents.
530 : : */
531 : : Datum
532 : 2 : test_ext_node_get_relid(PG_FUNCTION_ARGS)
533 : : {
534 : 2 : PG_RETURN_OID(PG_GETARG_TEST_EXT_NODE(0)->relid);
535 : : }
536 : :
537 : : Datum
538 : 1 : test_ext_node_get_repeat_count(PG_FUNCTION_ARGS)
539 : : {
540 : 1 : PG_RETURN_INT32(PG_GETARG_TEST_EXT_NODE(0)->repeat_count);
541 : : }
542 : :
543 : : /*
544 : : * Module initialization
545 : : */
546 : : void
547 : 2 : _PG_init(void)
548 : : {
549 : 2 : RegisterCustomScanMethods(&test_custom_scan_methods);
550 : 2 : RegisterExtensibleNodeMethods(&test_ext_node_methods);
551 : :
552 : 2 : DefineCustomIntVariable("test_extensible.repeat_count",
553 : : "Number of times the custom scan returns each row.",
554 : : NULL,
555 : : &test_repeat_count,
556 : : 2,
557 : : 1,
558 : : 100,
559 : : PGC_USERSET,
560 : : 0,
561 : : NULL, NULL, NULL);
562 : :
563 : 2 : MarkGUCPrefixReserved("test_extensible");
564 : :
565 : : /* Install the path-list hook to inject CustomPaths for the test table */
566 : 2 : prev_set_rel_pathlist_hook = set_rel_pathlist_hook;
567 : 2 : set_rel_pathlist_hook = test_set_rel_pathlist;
568 : 2 : }
|
