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 @extschema@.pg_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 @extschema@.pg_store_query_plan(
7+
query TEXT,
8+
filename TEXT)
9+
RETURNS VOID AS 'pg_execplan'
10+
LANGUAGE C;
11+
12+
CREATE OR REPLACE FUNCTION @extschema@.pg_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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy