pg_advice_comment.c

/*-------------------------------------------------------------------------
 *
 * pg_advice_comment.c
 *	  supply pg_plan_advice advice strings from query comments
 *
 * This module registers an advisor with pg_plan_advice. For each query that
 * is planned, it scans the raw query text for a hint-style comment of the form
 *
 *		/STAR+ ADVICE STAR/
 *
 * (replace STAR with *)
 *
 * and returns the text found between the opening and closing delimiters
 * verbatim as the advice string. pg_plan_advice parses that string with the
 * exact same parser it uses for the pg_plan_advice.advice GUC, so the full
 * advice grammar (SEQ_SCAN, INDEX_SCAN, JOIN_ORDER, HASH_JOIN, GATHER, ...) is
 * available without this module needing to understand any of it.
 *
 * The comment is located with a full SQL lexical scanner (query_scan.l, ported
 * from pg_hint_plan), so occurrences of the marker inside string literals,
 * quoted identifiers, dollar-quoted bodies, or ordinary comments are not
 * mistaken for advice. Only the first advice comment in the query is used.
 *
 * There are no GUCs and no SQL-callable functions: load the module via
 * shared_preload_libraries (alongside pg_plan_advice) and it just works.
 *
 * Copyright (c) 2016-2026, PostgreSQL Global Development Group
 *
 *	  pg_advice_comment.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"

#include "commands/explain_state.h"
#include "fmgr.h"
#include "lib/stringinfo.h"
#include "miscadmin.h"
#include "nodes/parsenodes.h"
#include "nodes/pathnodes.h"
#include "parser/scansup.h"

#include "query_scan.h"

PG_MODULE_MAGIC;

/*
 * Hook type used by pg_plan_advice to collect advice strings. This mirrors
 * the declaration in contrib/pg_plan_advice/pg_plan_advice.h. We declare it
 * locally rather than including that header, because the header is not part of
 * the standard server installation. If the upstream signature ever changes,
 * this will need to be kept in sync.
 */
typedef char *(*pg_plan_advice_advisor_hook) (PlannerGlobal *glob,
											   Query *parse,
											   const char *query_string,
											   int cursorOptions,
											   ExplainState *es);

void		_PG_init(void);

static char *pgac_advisor(PlannerGlobal *glob,
						  Query *parse,
						  const char *query_string,
						  int cursorOptions,
						  ExplainState *es);
static char *pgac_get_advice_from_comment(const char *query_string);

/*
 * Module initialization: register our advisor with pg_plan_advice.
 *
 * We resolve pg_plan_advice_add_advisor() at runtime so that we don't need to
 * link against pg_plan_advice. pg_plan_advice must be loadable as "$libdir/
 * pg_plan_advice"; normally it will already be loaded via
 * shared_preload_libraries.
 */
void
_PG_init(void)
{
	void		(*add_advisor_fn) (pg_plan_advice_advisor_hook hook);

	add_advisor_fn = (void (*) (pg_plan_advice_advisor_hook))
		load_external_function("pg_plan_advice", "pg_plan_advice_add_advisor",
							   true, NULL);
	(*add_advisor_fn) (pgac_advisor);
}

/*
 * Advisor callback.
 *
 * Return the advice string found in the first advice comment of the query, or
 * NULL if there is none. When we return NULL, pg_plan_advice falls through to
 * any other advisor or to the pg_plan_advice.advice GUC.
 *
 * The returned string is allocated in the current (short-lived) memory
 * context; pg_plan_advice only needs it to remain valid long enough to be
 * parsed.
 */
static char *
pgac_advisor(PlannerGlobal *glob, Query *parse,
			 const char *query_string, int cursorOptions,
			 ExplainState *es)
{
	return pgac_get_advice_from_comment(query_string);
}

/*
 * Extract the advice string from the first advice comment in the query text.
 *
 * This drives the SQL lexical scanner in query_scan.l, which copies the
 * contents of the first "/STAR+ ... STAR/" comment (and nothing else) into an
 * output buffer. Returns a palloc'd copy of that text with surrounding
 * whitespace trimmed, or NULL if there is no advice comment (or it is empty).
 */
static char *
pgac_get_advice_from_comment(const char *query_string)
{
	QueryScanState sstate;
	StringInfo	advice_buf;
	char	   *result = NULL;
	char	   *start;
	char	   *end;

	if (query_string == NULL)
		return NULL;

	sstate = query_scan_create();
	advice_buf = makeStringInfo();

	query_scan_setup(sstate, query_string, strlen(query_string), WARNING);
	for (;;)
	{
		QueryScanResult sr = query_scan(sstate, advice_buf);

		if (sr == QUERY_SCAN_EOL)
			break;

		/*
		 * Make sure we can break out of the loop if stuck. This should not be
		 * necessary in practice, but it gives an escape should the scanner
		 * contain a bug.
		 */
		CHECK_FOR_INTERRUPTS();
	}

	query_scan_finish(sstate);

	/* Trim surrounding whitespace; an empty comment supplies no advice. */
	start = advice_buf->data;
	end = advice_buf->data + advice_buf->len;
	while (start < end && scanner_isspace(*start))
		start++;
	while (end > start && scanner_isspace(*(end - 1)))
		end--;
	if (start < end)
		result = pnstrdup(start, end - start);

	pfree(advice_buf->data);
	pfree(advice_buf);
	return result;
}