Skip to content

Commit

Permalink
fix: worker_restart without security definer
Browse files Browse the repository at this point in the history
  • Loading branch information
steve-chavez committed Oct 17, 2024
1 parent 6649118 commit f26dadf
Show file tree
Hide file tree
Showing 5 changed files with 55 additions and 8 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# PG_NET
*A PostgreSQL extension that enables asynchronous (non-blocking) HTTP/HTTPS requests with SQL*.

Requires libcurl >= 7.83.
Requires libcurl >= 7.83. Compatible with PostgreSQL > = 12.

![PostgreSQL version](https://img.shields.io/badge/postgresql-12+-blue.svg)
[![License](https://img.shields.io/pypi/l/markdown-subtemplate.svg)](https://github.com/supabase/pg_net/blob/master/LICENSE)
Expand Down
5 changes: 5 additions & 0 deletions sql/pg_net--0.11.0--0.11.1.sql
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,11 @@ alter function net._http_collect_response ( bigint, boolean) security invoker;

alter function net.http_collect_response ( bigint, boolean) security invoker;

create or replace function net.worker_restart()
returns bool
language 'c'
as 'pg_net';

grant usage on schema net to PUBLIC;
grant all on all sequences in schema net to PUBLIC;
grant all on all tables in schema net to PUBLIC;
4 changes: 4 additions & 0 deletions sql/pg_net.sql
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,10 @@ create or replace function net._encode_url_with_params_array(url text, params_ar
immutable
as 'pg_net';

create or replace function net.worker_restart()
returns bool
language 'c'
as 'pg_net';

-- Interface to make an async request
-- API: Public
Expand Down
41 changes: 34 additions & 7 deletions src/worker.c
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
#include <string.h>
#include <inttypes.h>

#include <storage/shmem.h>

#include "util.h"
#include "core.h"

Expand All @@ -40,17 +42,25 @@ _Static_assert(LIBCURL_VERSION_NUM >= MIN_LIBCURL_VERSION_NUM, REQUIRED_LIBCURL_

PG_MODULE_MAGIC;

static char *guc_ttl;
static int guc_batch_size;
static char* guc_database_name;
static MemoryContext CurlMemContext = NULL;
static char* guc_ttl;
static int guc_batch_size;
static char* guc_database_name;
static MemoryContext CurlMemContext = NULL;
static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
static long latch_timeout = 1000;
static volatile sig_atomic_t got_sigterm = false;
static volatile sig_atomic_t got_sighup = false;
static bool* restart_worker = NULL;

void _PG_init(void);
PGDLLEXPORT void pg_net_worker(Datum main_arg) pg_attribute_noreturn();

static long latch_timeout = 1000;
static volatile sig_atomic_t got_sigterm = false;
static volatile sig_atomic_t got_sighup = false;
PG_FUNCTION_INFO_V1(worker_restart);
Datum worker_restart(PG_FUNCTION_ARGS) {
bool result = DatumGetBool(DirectFunctionCall1(pg_reload_conf, (Datum) NULL)); // reload the config
*restart_worker = true;
PG_RETURN_BOOL(result && *restart_worker); // TODO is not necessary to return a bool here, but we do it to maintain backward compatibility
}

static void
handle_sigterm(SIGNAL_ARGS)
Expand Down Expand Up @@ -141,6 +151,12 @@ void pg_net_worker(Datum main_arg) {
ProcessConfigFile(PGC_SIGHUP);
}

if (restart_worker && *restart_worker) {
*restart_worker = false;
elog(INFO, "Restarting pg_net worker");
break;
}

delete_expired_responses(guc_ttl, guc_batch_size);

consume_request_queue(lstate.curl_mhandle, guc_batch_size, CurlMemContext);
Expand Down Expand Up @@ -206,6 +222,14 @@ void pg_net_worker(Datum main_arg) {
proc_exit(EXIT_FAILURE);
}

static void net_shmem_startup(void) {
if (prev_shmem_startup_hook)
prev_shmem_startup_hook();

restart_worker = ShmemAlloc(sizeof(bool));
*restart_worker = false;
}

void _PG_init(void) {
if (IsBinaryUpgrade) {
return;
Expand All @@ -226,6 +250,9 @@ void _PG_init(void) {
.bgw_restart_time = 1,
});

prev_shmem_startup_hook = shmem_startup_hook;
shmem_startup_hook = net_shmem_startup;

CurlMemContext = AllocSetContextCreate(TopMemoryContext,
"pg_net curl context",
ALLOCSET_DEFAULT_MINSIZE,
Expand Down
11 changes: 11 additions & 0 deletions test/test_privileges.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,17 @@ def test_net_on_another_role(sess):
).fetchone()
assert response[0] == "SUCCESS"

## can use the net.worker_restart function
response = sess.execute(
text(
"""
set local role to another;
select net.worker_restart();
"""
)
).fetchone()
assert response[0] == True

sess.execute(text("""
set local role postgres;
drop role another;
Expand Down

0 comments on commit f26dadf

Please sign in to comment.