Skip to content

Commit 1349c4e

Browse files
danolivodanolivo
committed
Add store and exec plan functions
1 parent a39b8c6 commit 1349c4e

File tree

5 files changed

+263
-1
lines changed

5 files changed

+263
-1
lines changed

‎contrib/pg_execplan/init.sql

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,14 @@
11
\echo Use "CREATE EXTENSION pg_execplan" to load this file. \quit
22

3-
CREATE OR REPLACE FUNCTION @[email protected]_store_query_plan(query TEXT)
3+
-- Store plan of a query into a text file.
4+
-- query - query string which will be parsed and planned.
5+
-- filename - path to the file on a disk.
6+
CREATE OR REPLACE FUNCTION @[email protected]_store_query_plan(
7+
query TEXT,
8+
filename TEXT)
9+
RETURNS VOID AS 'pg_execplan'
10+
LANGUAGE C;
11+
12+
CREATE OR REPLACE FUNCTION @[email protected]_exec_query_plan(filename TEXT)
413
RETURNS VOID AS 'pg_execplan'
514
LANGUAGE C;

‎contrib/pg_execplan/pg_execplan.c

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
/*
2+
* pg_execplan.c
3+
*
4+
*/
5+
6+
#include "postgres.h"
7+
8+
#include "access/printtup.h"
9+
#include "commands/extension.h"
10+
#include "commands/prepare.h"
11+
#include "executor/executor.h"
12+
#include "nodes/plannodes.h"
13+
#include "tcop/pquery.h"
14+
#include "tcop/utility.h"
15+
#include "utils/builtins.h"
16+
#include "utils/memutils.h"
17+
#include "utils/plancache.h"
18+
#include "utils/snapmgr.h"
19+
20+
21+
#define EXPLAN_DEBUG_LEVEL 0
22+
23+
PG_MODULE_MAGIC;
24+
25+
PG_FUNCTION_INFO_V1(pg_store_query_plan);
26+
PG_FUNCTION_INFO_V1(pg_exec_query_plan);
27+
28+
void _PG_init(void);
29+
30+
/*
31+
* Module load/unload callback
32+
*/
33+
void
34+
_PG_init(void)
35+
{
36+
return;
37+
}
38+
39+
Datum
40+
pg_store_query_plan(PG_FUNCTION_ARGS)
41+
{
42+
char *query_string = TextDatumGetCString(PG_GETARG_DATUM(1)),
43+
*filename = TextDatumGetCString(PG_GETARG_DATUM(0)),
44+
*plan_string;
45+
int nstmts;
46+
FILE *fout;
47+
MemoryContext oldcontext;
48+
List *parsetree_list;
49+
RawStmt *parsetree;
50+
List *querytree_list,
51+
*plantree_list;
52+
QueryDesc *queryDesc;
53+
size_t string_len;
54+
55+
if (EXPLAN_DEBUG_LEVEL > 0)
56+
elog(LOG, "Store into %s plan of the query %s.", filename, query_string);
57+
58+
oldcontext = MemoryContextSwitchTo(MessageContext);
59+
60+
parsetree_list = pg_parse_query(query_string);
61+
nstmts = list_length(parsetree_list);
62+
if (nstmts != 1)
63+
elog(ERROR, "Query contains %d elements, but must contain only one.", nstmts);
64+
65+
parsetree = (RawStmt *) linitial(parsetree_list);
66+
querytree_list = pg_analyze_and_rewrite(parsetree, query_string, NULL, 0);
67+
plantree_list = pg_plan_queries(querytree_list, CURSOR_OPT_PARALLEL_OK, NULL);
68+
69+
queryDesc = CreateQueryDesc((PlannedStmt *) linitial(plantree_list),
70+
query_string,
71+
InvalidSnapshot,
72+
InvalidSnapshot,
73+
None_Receiver,
74+
0,
75+
0);
76+
77+
if (EXPLAN_DEBUG_LEVEL > 0)
78+
elog(INFO, "BEFORE writing %s ...", filename);
79+
80+
fout = fopen(filename, "wb");
81+
Assert(fout != NULL);
82+
string_len = strlen(query_string);
83+
fwrite(&string_len, sizeof(size_t), 1, fout);
84+
fwrite(query_string, sizeof(char), string_len, fout);
85+
86+
plan_string = nodeToString(queryDesc->plannedstmt);
87+
string_len = strlen(plan_string);
88+
fwrite(&string_len, sizeof(size_t), 1, fout);
89+
fwrite(plan_string, sizeof(char), string_len, fout);
90+
91+
fclose(fout);
92+
MemoryContextSwitchTo(oldcontext);
93+
PG_RETURN_VOID();
94+
}
95+
96+
static void
97+
LoadPlanFromFile(const char *filename, char **query_string, char **plan_string)
98+
{
99+
FILE *fin;
100+
size_t string_len;
101+
int nelems;
102+
103+
fin = fopen(filename, "rb");
104+
Assert(fin != NULL);
105+
106+
nelems = fread(&string_len, sizeof(size_t), 1, fin);
107+
Assert(nelems == 1);
108+
*query_string = palloc0(string_len + 1);
109+
nelems = fread(*query_string, sizeof(char), string_len, fin);
110+
Assert(nelems == string_len);
111+
112+
nelems = fread(&string_len, sizeof(size_t), 1, fin);
113+
Assert(nelems == 1);
114+
*plan_string = palloc0(string_len + 1);
115+
nelems = fread(*plan_string, sizeof(char), string_len, fin);
116+
Assert(nelems == string_len);
117+
118+
fclose(fin);
119+
120+
}
121+
122+
Datum
123+
pg_exec_query_plan(PG_FUNCTION_ARGS)
124+
{
125+
char *filename = TextDatumGetCString(PG_GETARG_DATUM(0)),
126+
*query_string = NULL,
127+
*plan_string = NULL;
128+
PlannedStmt *pstmt;
129+
ParamListInfo paramLI = NULL;
130+
CachedPlanSource *psrc;
131+
CachedPlan *cplan;
132+
Portal portal;
133+
DestReceiver *receiver;
134+
int16 format = 0;
135+
int eflags = 0;
136+
137+
LoadPlanFromFile(filename, &query_string, &plan_string);
138+
pstmt = (PlannedStmt *) stringToNode(plan_string);
139+
140+
psrc = CreateCachedPlan(NULL, query_string, query_string);
141+
CompleteCachedPlan(psrc, NIL, NULL, NULL, 0, NULL, NULL,
142+
CURSOR_OPT_GENERIC_PLAN, false);
143+
StorePreparedStatement(query_string, psrc, false);
144+
SetRemoteSubplan(psrc, pstmt);
145+
cplan = GetCachedPlan(psrc, paramLI, false);
146+
147+
if (EXPLAN_DEBUG_LEVEL > 0)
148+
elog(INFO, "query: %s\n", query_string);
149+
if (EXPLAN_DEBUG_LEVEL > 1)
150+
elog(INFO, "\nplan: %s\n", plan_string);
151+
152+
receiver = CreateDestReceiver(DestDebug);
153+
portal = CreateNewPortal();
154+
portal->visible = false;
155+
PortalDefineQuery(portal,
156+
NULL,
157+
query_string,
158+
query_string,
159+
cplan->stmt_list,
160+
cplan);
161+
PortalStart(portal, paramLI, eflags, InvalidSnapshot);
162+
PortalSetResultFormat(portal, 0, &format);
163+
(void) PortalRun(portal,
164+
FETCH_ALL,
165+
true,
166+
receiver,
167+
receiver,
168+
query_string);
169+
receiver->rDestroy(receiver);
170+
PortalDrop(portal, false);
171+
DropPreparedStatement(query_string, false);
172+
173+
if (EXPLAN_DEBUG_LEVEL > 0)
174+
elog(INFO, "query execution finished.\n");
175+
176+
PG_RETURN_VOID();
177+
}

‎contrib/pg_execplan/pg_execplan.control

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# pg_execplan extension
2+
comment = 'Execute raw query plan on remote node'
3+
default_version = '0.1'
4+
module_pathname = '$libdir/pg_execplan'
5+
relocatable = false
6+
requires = 'postgres_fdw'

‎src/backend/utils/cache/plancache.c

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1906,3 +1906,71 @@ ResetPlanCache(void)
19061906
}
19071907
}
19081908
}
1909+
1910+
void
1911+
SetRemoteSubplan(CachedPlanSource *plansource, PlannedStmt *rstmt)
1912+
{
1913+
CachedPlan *plan;
1914+
MemoryContext plan_context;
1915+
MemoryContext oldcxt;
1916+
PlannedStmt *stmt;
1917+
1918+
Assert(plansource->raw_parse_tree == NULL);
1919+
Assert(plansource->query_list == NIL);
1920+
1921+
/*
1922+
* Make dedicated query context to store cached plan. It is in current
1923+
* memory context for now, later it will be reparented to
1924+
* CachedMemoryContext. If it is in CachedMemoryContext initially we would
1925+
* have to destroy it in case of error.
1926+
*/
1927+
plan_context = AllocSetContextCreate(CurrentMemoryContext,
1928+
"CachedPlan",
1929+
ALLOCSET_DEFAULT_SIZES);
1930+
oldcxt = MemoryContextSwitchTo(plan_context);
1931+
1932+
stmt = makeNode(PlannedStmt);
1933+
1934+
stmt->commandType = rstmt->commandType;
1935+
stmt->hasReturning = rstmt->hasReturning;
1936+
stmt->resultRelations = rstmt->resultRelations;
1937+
stmt->subplans = rstmt->subplans;
1938+
stmt->rowMarks = rstmt->rowMarks;
1939+
stmt->planTree = rstmt->planTree;
1940+
stmt->rtable = rstmt->rtable;
1941+
1942+
stmt->canSetTag = true;
1943+
stmt->transientPlan = false;
1944+
stmt->utilityStmt = NULL;
1945+
stmt->rewindPlanIDs = NULL;
1946+
stmt->relationOids = rstmt->relationOids;
1947+
stmt->invalItems = rstmt->invalItems;
1948+
1949+
/*
1950+
* Create and fill the CachedPlan struct within the new context.
1951+
*/
1952+
plan = (CachedPlan *) palloc(sizeof(CachedPlan));
1953+
plan->magic = CACHEDPLAN_MAGIC;
1954+
plan->stmt_list = list_make1(stmt);
1955+
plan->saved_xmin = InvalidTransactionId;
1956+
plan->refcount = 1; /* will be referenced by plansource */
1957+
plan->context = plan_context;
1958+
plan->dependsOnRole = false;
1959+
if (plansource->is_saved)
1960+
{
1961+
MemoryContextSetParent(plan_context, CacheMemoryContext);
1962+
plan->is_saved = true;
1963+
}
1964+
else
1965+
{
1966+
MemoryContextSetParent(plan_context,
1967+
MemoryContextGetParent(plansource->context));
1968+
plan->is_saved = false;
1969+
}
1970+
plan->is_valid = true;
1971+
plan->is_oneshot = false;
1972+
1973+
plansource->gplan = plan;
1974+
1975+
MemoryContextSwitchTo(oldcxt);
1976+
}

‎src/include/utils/plancache.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
#include "access/tupdesc.h"
1919
#include "nodes/params.h"
20+
#include "nodes/plannodes.h"
2021

2122
/* Forward declaration, to avoid including parsenodes.h here */
2223
struct RawStmt;
@@ -178,5 +179,6 @@ extern CachedPlan *GetCachedPlan(CachedPlanSource *plansource,
178179
ParamListInfo boundParams,
179180
bool useResOwner);
180181
extern void ReleaseCachedPlan(CachedPlan *plan, bool useResOwner);
182+
extern void SetRemoteSubplan(CachedPlanSource *plansource, PlannedStmt *rstmt);
181183

182184
#endif /* PLANCACHE_H */

0 commit comments

Comments
 (0)