Skip to content

Commit

Permalink
feat: before-create script
Browse files Browse the repository at this point in the history
  • Loading branch information
@soedirgo
soedirgo committed Aug 14, 2023
1 parent 0068d70 commit da58877
Show file tree
Hide file tree
Showing 5 changed files with 216 additions and 26 deletions.
26 changes: 26 additions & 0 deletions nix/pg_tle.nix
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
{ stdenv, fetchFromGitHub, postgresql, flex }:

stdenv.mkDerivation rec {
pname = "pg_tle";
version = "1.0.4";

nativeBuildInputs = [ flex ];
buildInputs = [ postgresql ];

src = fetchFromGitHub {
owner = "aws";
repo = pname;
rev = "refs/tags/v${version}";
hash = "sha256-W/7pLy/27VatCdzUh1NZ4K2FRMD1erfHiFV2eY2x2W0=";
};

makeFlags = [ "FLEX=flex" ];

installPhase = ''
mkdir -p $out/{lib,share/postgresql/extension}
cp *.so $out/lib
cp *.sql $out/share/postgresql/extension
cp *.control $out/share/postgresql/extension
'';
}
18 changes: 15 additions & 3 deletions shell.nix
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@ let
};
pgWithExt = { postgresql } :
let
pg = postgresql.withPackages (p: [ (supautils {inherit postgresql;}) ]);
pg = postgresql.withPackages (p: [
(supautils { inherit postgresql; })
(callPackage ./nix/pg_tle.nix { inherit postgresql; })
]);
ver = builtins.head (builtins.splitVersion postgresql.version);
script = ''
export PATH=${pg}/bin:"$PATH"
Expand All @@ -34,11 +37,11 @@ let
PGTZ=UTC initdb --no-locale --encoding=UTF8 --nosync -U "$PGUSER"
options="-F -c listen_addresses=\"\" -k $PGDATA -c shared_preload_libraries=\"supautils\""
options="-F -c listen_addresses=\"\" -k $PGDATA -c shared_preload_libraries=\"pg_tle, supautils\""
reserved_roles="supabase_storage_admin, anon, reserved_but_not_yet_created, authenticator*"
reserved_memberships="pg_read_server_files, pg_write_server_files, pg_execute_server_program, role_with_reserved_membership"
privileged_extensions="hstore, postgres_fdw"
privileged_extensions="hstore, postgres_fdw, pg_tle"
privileged_extensions_custom_scripts_path="$tmpdir/privileged_extensions_custom_scripts"
privileged_role="privileged_role"
privileged_role_allowed_configs="session_replication_role, pgrst.*, other.nested.*"
Expand All @@ -49,6 +52,15 @@ let
pg_ctl start -o "$options" -o "$reserved_stuff_options" -o "$placeholder_stuff_options"
mkdir -p "$tmpdir/privileged_extensions_custom_scripts/hstore"
echo "do \$\$
begin
if not exists (select from pg_extension where extname = 'pg_tle') then
return;
end if;
if exists (select from pgtle.available_extensions() where name = @extname@) then
raise notice 'extname: %, extschema: %, extversion: %, extcascade: %', @extname@, @extschema@, @extversion@, @extcascade@;
end if;
end \$\$;" > "$tmpdir/privileged_extensions_custom_scripts/before-create.sql"
echo 'create table t1();' > "$tmpdir/privileged_extensions_custom_scripts/hstore/before-create.sql"
echo 'drop table t1; create table t2 as values (1);' > "$tmpdir/privileged_extensions_custom_scripts/hstore/after-create.sql"
Expand Down
161 changes: 140 additions & 21 deletions src/privileged_extensions.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include <catalog/pg_authid.h>
#include <catalog/pg_collation.h>
#include <catalog/pg_type.h>
#include <commands/defrem.h>
#include <executor/spi.h>
#include <miscadmin.h>
#include <nodes/pg_list.h>
Expand All @@ -18,37 +19,72 @@
#include "privileged_extensions.h"
#include "utils.h"

// TODO: interpolate extschema, current_role, current_database_owner
static void run_custom_script(const char *filename) {
// Prevent recursively running custom scripts
static bool running_custom_script = false;

static void run_custom_script(const char *filename, const char *extname,
const char *extschema, const char *extversion,
bool extcascade) {
if (running_custom_script) {
return;
}
running_custom_script = true;
PushActiveSnapshot(GetTransactionSnapshot());
SPI_connect();
{
char *sql_tmp1 = "do $$\n"
"begin\n"
" execute pg_read_file(";
char *sql_tmp2 = quote_literal_cstr(filename);
char *sql_tmp3 = ");\n"
"exception\n"
" when undefined_file then\n"
" -- skip\n"
"end\n"
"$$;";
size_t sql_len = strlen(sql_tmp1) + strlen(sql_tmp2) + strlen(sql_tmp3);
char *sql_tmp01 = "do $$\n"
"begin\n"
" execute\n"
" replace(\n"
" replace(\n"
" replace(\n"
" replace(\n"
" pg_read_file(\n";
char *sql_tmp02 = quote_literal_cstr(filename);
char *sql_tmp03 = " ),\n"
" '@extname@', ";
char *sql_tmp04 = quote_literal_cstr(quote_literal_cstr(extname));
char *sql_tmp05 = " ),\n"
" '@extschema@', ";
char *sql_tmp06 =
extschema == NULL
? "'null'"
: quote_literal_cstr(quote_literal_cstr(extschema));
char *sql_tmp07 = " ),\n"
" '@extversion@', ";
char *sql_tmp08 =
extversion == NULL
? "'null'"
: quote_literal_cstr(quote_literal_cstr(extversion));
char *sql_tmp09 = " ), "
" '@extcascade@', ";
char *sql_tmp10 = extcascade ? "'true'" : "'false'";
char *sql_tmp11 = " );\n"
"exception\n"
" when undefined_file then\n"
" -- skip\n"
"end\n"
"$$;";
size_t sql_len =
strlen(sql_tmp01) + strlen(sql_tmp02) + strlen(sql_tmp03) +
strlen(sql_tmp04) + strlen(sql_tmp05) + strlen(sql_tmp06) +
strlen(sql_tmp07) + strlen(sql_tmp08) + strlen(sql_tmp09) +
strlen(sql_tmp10) + strlen(sql_tmp11);
char *sql = (char *)palloc(sql_len);
int rc;

snprintf(sql, sql_len, "%s%s%s", sql_tmp1, sql_tmp2, sql_tmp3);
snprintf(sql, sql_len, "%s%s%s%s%s%s%s%s%s%s%s", sql_tmp01, sql_tmp02,
sql_tmp03, sql_tmp04, sql_tmp05, sql_tmp06, sql_tmp07,
sql_tmp08, sql_tmp09, sql_tmp10, sql_tmp11);

rc = SPI_execute(sql, false, 0);
if (rc != SPI_OK_UTILITY) {
elog(ERROR, "SPI_execute failed with error code %d", rc);
}

pfree(sql_tmp2);
pfree(sql);
}
SPI_finish();
PopActiveSnapshot();
running_custom_script = false;
}

void handle_create_extension(
Expand All @@ -59,13 +95,72 @@ void handle_create_extension(
CreateExtensionStmt *stmt = (CreateExtensionStmt *)pstmt->utilityStmt;
char *filename = (char *)palloc(MAXPGPATH);

// Run before-create script.
// Run global before-create script.
{
DefElem *d_schema = NULL;
DefElem *d_new_version = NULL;
DefElem *d_cascade = NULL;
char *extschema = NULL;
char *extversion = NULL;
bool extcascade = false;
ListCell *option_cell = NULL;

foreach (option_cell, stmt->options) {
DefElem *defel = (DefElem *)lfirst(option_cell);

if (strcmp(defel->defname, "schema") == 0) {
d_schema = defel;
extschema = defGetString(d_schema);
} else if (strcmp(defel->defname, "new_version") == 0) {
d_new_version = defel;
extversion = defGetString(d_new_version);
} else if (strcmp(defel->defname, "cascade") == 0) {
d_cascade = defel;
extcascade = defGetBoolean(d_cascade);
}
}

switch_to_superuser(privileged_extensions_superuser);

snprintf(filename, MAXPGPATH, "%s/before-create.sql",
privileged_extensions_custom_scripts_path);
run_custom_script(filename, stmt->extname, extschema, extversion,
extcascade);

switch_to_original_role();
}

// Run per-extension before-create script.
{
DefElem *d_schema = NULL;
DefElem *d_new_version = NULL;
DefElem *d_cascade = NULL;
char *extschema = NULL;
char *extversion = NULL;
bool extcascade = false;
ListCell *option_cell = NULL;

foreach (option_cell, stmt->options) {
DefElem *defel = (DefElem *)lfirst(option_cell);

if (strcmp(defel->defname, "schema") == 0) {
d_schema = defel;
extschema = defGetString(d_schema);
} else if (strcmp(defel->defname, "new_version") == 0) {
d_new_version = defel;
extversion = defGetString(d_new_version);
} else if (strcmp(defel->defname, "cascade") == 0) {
d_cascade = defel;
extcascade = defGetBoolean(d_cascade);
}
}

switch_to_superuser(privileged_extensions_superuser);

snprintf(filename, MAXPGPATH, "%s/%s/before-create.sql",
privileged_extensions_custom_scripts_path, stmt->extname);
run_custom_script(filename);
run_custom_script(filename, stmt->extname, extschema, extversion,
extcascade);

switch_to_original_role();
}
Expand All @@ -82,13 +177,37 @@ void handle_create_extension(
run_process_utility_hook(process_utility_hook);
}

// Run after-create script.
// Run per-extension after-create script.
{
DefElem *d_schema = NULL;
DefElem *d_new_version = NULL;
DefElem *d_cascade = NULL;
char *extschema = NULL;
char *extversion = NULL;
bool extcascade = false;
ListCell *option_cell = NULL;

foreach (option_cell, stmt->options) {
DefElem *defel = (DefElem *)lfirst(option_cell);

if (strcmp(defel->defname, "schema") == 0) {
d_schema = defel;
extschema = defGetString(d_schema);
} else if (strcmp(defel->defname, "new_version") == 0) {
d_new_version = defel;
extversion = defGetString(d_new_version);
} else if (strcmp(defel->defname, "cascade") == 0) {
d_cascade = defel;
extcascade = defGetBoolean(d_cascade);
}
}

switch_to_superuser(privileged_extensions_superuser);

snprintf(filename, MAXPGPATH, "%s/%s/after-create.sql",
privileged_extensions_custom_scripts_path, stmt->extname);
run_custom_script(filename);
run_custom_script(filename, stmt->extname, extschema, extversion,
extcascade);

switch_to_original_role();
}
Expand Down
23 changes: 22 additions & 1 deletion test/expected/privileged_extensions.out
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
-- non-superuser extensions role
create role extensions_role login;
grant all on database postgres to extensions_role;
alter default privileges for role postgres in schema public grant all on tables to extensions_role;
set role extensions_role;
\echo
Expand All @@ -15,7 +16,7 @@ select '1=>2'::hstore;
drop extension hstore;
\echo

-- custom scripts are run
-- per-extension custom scripts are run
select * from t2;
column1
---------
Expand All @@ -27,6 +28,26 @@ drop table t2;
set role extensions_role;
\echo

-- global extension custom scripts are run
create extension pg_tle;
reset role;
grant pgtle_admin to extensions_role;
set role extensions_role;
select pgtle.install_extension('foo', '1', '', 'select 1', '{}');
install_extension
-------------------
t
(1 row)

create extension foo cascade;
NOTICE: extname: foo, extschema: <NULL>, extversion: <NULL>, extcascade: t
drop extension pg_tle cascade;
NOTICE: drop cascades to 3 other objects
DETAIL: drop cascades to function pgtle."foo.control"()
drop cascades to function pgtle."foo--1.sql"()
drop cascades to extension foo
\echo

-- custom scripts are run even for superusers
reset role;
create extension hstore;
Expand Down
14 changes: 13 additions & 1 deletion test/sql/privileged_extensions.sql
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
-- non-superuser extensions role
create role extensions_role login;
grant all on database postgres to extensions_role;
alter default privileges for role postgres in schema public grant all on tables to extensions_role;
set role extensions_role;
\echo
Expand All @@ -11,14 +12,25 @@ select '1=>2'::hstore;
drop extension hstore;
\echo

-- custom scripts are run
-- per-extension custom scripts are run
select * from t2;

reset role;
drop table t2;
set role extensions_role;
\echo

-- global extension custom scripts are run
create extension pg_tle;
reset role;
grant pgtle_admin to extensions_role;
set role extensions_role;
select pgtle.install_extension('foo', '1', '', 'select 1', '{}');
create extension foo cascade;

drop extension pg_tle cascade;
\echo

-- custom scripts are run even for superusers
reset role;
create extension hstore;
Expand Down

0 comments on commit da58877

Please sign in to comment.