Skip to content

Commit 62adaf6

Browse files
steve-chavezsteve-chavez
committed
refactor: deduplicate logic in create extension
The `handle_create_extension` function repeated statements unnecessarily: - switched to superuser multiple times - downgraded to privileged role multiple times The switching only needed to be done one time at the start, and downgrading one time at the end. The same effect is maintained, no tests are broken. Also `run_process_utility_hook` is executed once for the `create extension` logic now. Partly addresses #79.
1 parent 6dd0411 commit 62adaf6

File tree

3 files changed

+113
-148
lines changed

3 files changed

+113
-148
lines changed

‎src/privileged_extensions.c‎

Lines changed: 82 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -67,94 +67,99 @@ static void run_custom_script(const char *filename, const char *extname,
6767
running_custom_script = false;
6868
}
6969

70-
void handle_create_extension(
71-
void (*process_utility_hook)(PROCESS_UTILITY_PARAMS),
72-
PROCESS_UTILITY_PARAMS, const char *privileged_extensions,
73-
const char *superuser,
74-
const char *privileged_extensions_custom_scripts_path,
75-
const extension_parameter_overrides *epos, const size_t total_epos) {
76-
CreateExtensionStmt *stmt = (CreateExtensionStmt *)pstmt->utilityStmt;
77-
char *filename = (char *)palloc(MAXPGPATH);
78-
79-
// Run global before-create script.
80-
{
81-
DefElem *d_schema = NULL;
82-
DefElem *d_new_version = NULL;
83-
DefElem *d_cascade = NULL;
84-
char *extschema = NULL;
85-
char *extversion = NULL;
86-
bool extcascade = false;
87-
ListCell *option_cell = NULL;
88-
bool already_switched_to_superuser = false;
89-
90-
foreach (option_cell, stmt->options) {
91-
DefElem *defel = (DefElem *)lfirst(option_cell);
92-
93-
if (strcmp(defel->defname, "schema") == 0) {
94-
d_schema = defel;
95-
extschema = defGetString(d_schema);
96-
} else if (strcmp(defel->defname, "new_version") == 0) {
97-
d_new_version = defel;
98-
extversion = defGetString(d_new_version);
99-
} else if (strcmp(defel->defname, "cascade") == 0) {
100-
d_cascade = defel;
101-
extcascade = defGetBoolean(d_cascade);
102-
}
70+
void run_global_before_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path){
71+
DefElem *d_schema = NULL, *d_new_version = NULL, *d_cascade = NULL;
72+
char *extschema = NULL, *extversion = NULL;
73+
bool extcascade = false;
74+
char filename[MAXPGPATH];
75+
76+
ListCell *option_cell = NULL;
77+
78+
foreach (option_cell, options) {
79+
DefElem *defel = (DefElem *)lfirst(option_cell);
80+
81+
if (strcmp(defel->defname, "schema") == 0) {
82+
d_schema = defel;
83+
extschema = defGetString(d_schema);
84+
} else if (strcmp(defel->defname, "new_version") == 0) {
85+
d_new_version = defel;
86+
extversion = defGetString(d_new_version);
87+
} else if (strcmp(defel->defname, "cascade") == 0) {
88+
d_cascade = defel;
89+
extcascade = defGetBoolean(d_cascade);
10390
}
91+
}
10492

105-
switch_to_superuser(superuser,
106-
&already_switched_to_superuser);
107-
108-
snprintf(filename, MAXPGPATH, "%s/before-create.sql",
109-
privileged_extensions_custom_scripts_path);
110-
run_custom_script(filename, stmt->extname, extschema, extversion,
111-
extcascade);
93+
snprintf(filename, MAXPGPATH, "%s/before-create.sql",
94+
privileged_extensions_custom_scripts_path);
95+
run_custom_script(filename, extname, extschema, extversion,
96+
extcascade);
97+
}
11298

113-
if (!already_switched_to_superuser) {
114-
switch_to_original_role();
99+
void run_ext_before_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path){
100+
DefElem *d_schema = NULL;
101+
DefElem *d_new_version = NULL;
102+
DefElem *d_cascade = NULL;
103+
char *extschema = NULL;
104+
char *extversion = NULL;
105+
bool extcascade = false;
106+
ListCell *option_cell = NULL;
107+
char filename[MAXPGPATH];
108+
109+
foreach (option_cell, options) {
110+
DefElem *defel = (DefElem *)lfirst(option_cell);
111+
112+
if (strcmp(defel->defname, "schema") == 0) {
113+
d_schema = defel;
114+
extschema = defGetString(d_schema);
115+
} else if (strcmp(defel->defname, "new_version") == 0) {
116+
d_new_version = defel;
117+
extversion = defGetString(d_new_version);
118+
} else if (strcmp(defel->defname, "cascade") == 0) {
119+
d_cascade = defel;
120+
extcascade = defGetBoolean(d_cascade);
115121
}
116122
}
117123

118-
// Run per-extension before-create script.
119-
{
120-
DefElem *d_schema = NULL;
121-
DefElem *d_new_version = NULL;
122-
DefElem *d_cascade = NULL;
123-
char *extschema = NULL;
124-
char *extversion = NULL;
125-
bool extcascade = false;
126-
ListCell *option_cell = NULL;
127-
bool already_switched_to_superuser = false;
128-
129-
foreach (option_cell, stmt->options) {
130-
DefElem *defel = (DefElem *)lfirst(option_cell);
131-
132-
if (strcmp(defel->defname, "schema") == 0) {
133-
d_schema = defel;
134-
extschema = defGetString(d_schema);
135-
} else if (strcmp(defel->defname, "new_version") == 0) {
136-
d_new_version = defel;
137-
extversion = defGetString(d_new_version);
138-
} else if (strcmp(defel->defname, "cascade") == 0) {
139-
d_cascade = defel;
140-
extcascade = defGetBoolean(d_cascade);
141-
}
142-
}
143-
144-
switch_to_superuser(superuser,
145-
&already_switched_to_superuser);
146124

147-
snprintf(filename, MAXPGPATH, "%s/%s/before-create.sql",
148-
privileged_extensions_custom_scripts_path, stmt->extname);
149-
run_custom_script(filename, stmt->extname, extschema, extversion,
150-
extcascade);
125+
snprintf(filename, MAXPGPATH, "%s/%s/before-create.sql",
126+
privileged_extensions_custom_scripts_path, extname);
127+
run_custom_script(filename, extname, extschema, extversion,
128+
extcascade);
129+
}
151130

152-
if (!already_switched_to_superuser) {
153-
switch_to_original_role();
131+
void run_ext_after_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path){
132+
DefElem *d_schema = NULL;
133+
DefElem *d_new_version = NULL;
134+
DefElem *d_cascade = NULL;
135+
char *extschema = NULL;
136+
char *extversion = NULL;
137+
bool extcascade = false;
138+
ListCell *option_cell = NULL;
139+
char filename[MAXPGPATH];
140+
141+
foreach (option_cell, options) {
142+
DefElem *defel = (DefElem *)lfirst(option_cell);
143+
144+
if (strcmp(defel->defname, "schema") == 0) {
145+
d_schema = defel;
146+
extschema = defGetString(d_schema);
147+
} else if (strcmp(defel->defname, "new_version") == 0) {
148+
d_new_version = defel;
149+
extversion = defGetString(d_new_version);
150+
} else if (strcmp(defel->defname, "cascade") == 0) {
151+
d_cascade = defel;
152+
extcascade = defGetBoolean(d_cascade);
154153
}
155154
}
156155

157-
// Apply overrides.
156+
snprintf(filename, MAXPGPATH, "%s/%s/after-create.sql",
157+
privileged_extensions_custom_scripts_path, extname);
158+
run_custom_script(filename, extname, extschema, extversion,
159+
extcascade);
160+
}
161+
162+
void override_create_ext_statement(CreateExtensionStmt *stmt, const size_t total_epos, const extension_parameter_overrides *epos){
158163
for (size_t i = 0; i < total_epos; i++) {
159164
if (strcmp(epos[i].name, stmt->extname) == 0) {
160165
const extension_parameter_overrides *epo = &epos[i];
@@ -189,63 +194,6 @@ void handle_create_extension(
189194
}
190195
}
191196
}
192-
193-
// Run `CREATE EXTENSION`.
194-
if (is_string_in_comma_delimited_string(stmt->extname,
195-
privileged_extensions)) {
196-
bool already_switched_to_superuser = false;
197-
switch_to_superuser(superuser,
198-
&already_switched_to_superuser);
199-
200-
run_process_utility_hook(process_utility_hook);
201-
202-
if (!already_switched_to_superuser) {
203-
switch_to_original_role();
204-
}
205-
} else {
206-
run_process_utility_hook(process_utility_hook);
207-
}
208-
209-
// Run per-extension after-create script.
210-
{
211-
DefElem *d_schema = NULL;
212-
DefElem *d_new_version = NULL;
213-
DefElem *d_cascade = NULL;
214-
char *extschema = NULL;
215-
char *extversion = NULL;
216-
bool extcascade = false;
217-
ListCell *option_cell = NULL;
218-
bool already_switched_to_superuser = false;
219-
220-
foreach (option_cell, stmt->options) {
221-
DefElem *defel = (DefElem *)lfirst(option_cell);
222-
223-
if (strcmp(defel->defname, "schema") == 0) {
224-
d_schema = defel;
225-
extschema = defGetString(d_schema);
226-
} else if (strcmp(defel->defname, "new_version") == 0) {
227-
d_new_version = defel;
228-
extversion = defGetString(d_new_version);
229-
} else if (strcmp(defel->defname, "cascade") == 0) {
230-
d_cascade = defel;
231-
extcascade = defGetBoolean(d_cascade);
232-
}
233-
}
234-
235-
switch_to_superuser(superuser,
236-
&already_switched_to_superuser);
237-
238-
snprintf(filename, MAXPGPATH, "%s/%s/after-create.sql",
239-
privileged_extensions_custom_scripts_path, stmt->extname);
240-
run_custom_script(filename, stmt->extname, extschema, extversion,
241-
extcascade);
242-
243-
if (!already_switched_to_superuser) {
244-
switch_to_original_role();
245-
}
246-
}
247-
248-
pfree(filename);
249197
}
250198

251199
bool all_extensions_are_privileged(List *objects, const char *privileged_extensions){

‎src/privileged_extensions.h‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,16 @@
44
#include "extensions_parameter_overrides.h"
55
#include "utils.h"
66

7-
extern void handle_create_extension(
8-
void (*process_utility_hook)(PROCESS_UTILITY_PARAMS),
9-
PROCESS_UTILITY_PARAMS, const char *privileged_extensions,
10-
const char *superuser,
11-
const char *privileged_extensions_custom_scripts_path,
12-
const extension_parameter_overrides *epos, const size_t total_epos);
13-
147
bool all_extensions_are_privileged(List *objects, const char *privileged_extensions);
158

169
bool is_extension_privileged(const char *extname, const char *privileged_extensions);
1710

11+
void run_global_before_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path);
12+
13+
void run_ext_before_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path);
14+
15+
void run_ext_after_create_script(char *extname, List *options, const char *privileged_extensions_custom_scripts_path);
16+
17+
void override_create_ext_statement(CreateExtensionStmt *stmt, const size_t total_epos, const extension_parameter_overrides *epos);
18+
1819
#endif

‎src/supautils.c‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -448,13 +448,29 @@ static void supautils_hook(PROCESS_UTILITY_PARAMS) {
448448

449449
constrain_extension(stmt->extname, cexts, total_cexts);
450450

451-
handle_create_extension(prev_hook,
452-
PROCESS_UTILITY_ARGS,
453-
privileged_extensions,
454-
supautils_superuser,
455-
privileged_extensions_custom_scripts_path,
456-
epos, total_epos);
457-
return;
451+
if (is_extension_privileged(stmt->extname, privileged_extensions)) {
452+
bool already_switched_to_superuser = false;
453+
454+
switch_to_superuser(supautils_superuser, &already_switched_to_superuser);
455+
456+
run_global_before_create_script(stmt->extname, stmt->options, privileged_extensions_custom_scripts_path);
457+
458+
run_ext_before_create_script(stmt->extname, stmt->options, privileged_extensions_custom_scripts_path);
459+
460+
override_create_ext_statement(stmt, total_epos, epos);
461+
462+
run_process_utility_hook(prev_hook);
463+
464+
run_ext_after_create_script(stmt->extname, stmt->options, privileged_extensions_custom_scripts_path);
465+
466+
if (!already_switched_to_superuser) {
467+
switch_to_original_role();
468+
}
469+
470+
return;
471+
}
472+
473+
break;
458474
}
459475

460476
/*

0 commit comments

Comments
 (0)