Add routine to get running worker pid (#295)

* Add extension_base.get_worker_pid() to look up a worker's PID by job ID

Exposes GetBaseWorkerPid() as an exportable C function (declared in
base_workers.h) that reads the shared-memory BaseWorkerHash under a
shared lock and returns the PID for the given worker ID in the current
database, or 0 if not found. The SQL wrapper get_worker_pid(int)
delegates to it. Tests cover both the running and not-found cases.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: David Christensen <david.christensen@snowflake.com>

* Remove get_worker_pid() from base SQL; keep only in upgrade path

The function is new in 3.3 so it belongs only in the 3.2->3.3 upgrade
script, not in the 1.6 base install SQL.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: David Christensen <david.christensen@snowflake.com>

* black format test_base_worker_launcher.py

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: David Christensen <david.christensen@snowflake.com>

---------

Signed-off-by: David Christensen <david.christensen@snowflake.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: sfc-gh-dachristensen

pg_extension_base/include/pg_extension_base/base_workers.h

@@ -115,4 +115,6 @@ extern PGDLLEXPORT void DeregisterBaseWorkerSelf(void);
 
 extern PGDLLEXPORT int32 MyBaseWorkerId;
 
+extern PGDLLEXPORT pid_t GetBaseWorkerPid(int32 workerId);
+
 #endif

pg_extension_base/pg_extension_base--3.2--3.3.sql

@@ -32,3 +32,15 @@ COMMENT ON FUNCTION extension_base.deregister_worker(text, bool)
  IS 'deregister a base worker';
 
 REVOKE ALL ON FUNCTION extension_base.deregister_worker(text, bool) FROM public;
+
+
+/* return the PID of a running base worker given its worker ID, or 0 if not found */
+CREATE FUNCTION extension_base.get_worker_pid(worker_id int)
+ RETURNS int
+ LANGUAGE c STRICT
+AS 'MODULE_PATHNAME', $function$pg_extension_base_get_worker_pid$function$;
+
+COMMENT ON FUNCTION extension_base.get_worker_pid(int)
+ IS 'return the PID of a running base worker, or 0 if not found';
+
+REVOKE ALL ON FUNCTION extension_base.get_worker_pid(int) FROM public;

pg_extension_base/src/base_worker_launcher.c

@@ -367,6 +367,7 @@ static void BaseWorkerTransactionCallback(XactEvent event, void *arg);
 /* SQL-callable functions */
 PG_FUNCTION_INFO_V1(pg_extension_base_register_worker);
 PG_FUNCTION_INFO_V1(pg_extension_base_deregister_worker);
+PG_FUNCTION_INFO_V1(pg_extension_base_get_worker_pid);
 PG_FUNCTION_INFO_V1(pg_extension_base_list_base_workers);
 PG_FUNCTION_INFO_V1(pg_extension_base_list_database_starters);
 
@@ -2925,6 +2926,42 @@ BaseWorkerTransactionCallback(XactEvent event, void *arg)
 }
 
 
+/*
+ * GetBaseWorkerPid returns the PID of a running base worker given its worker
+ * ID, or 0 if the worker is not found in the current database.
+ */
+pid_t
+GetBaseWorkerPid(int32 workerId)
+{
+	pid_t		workerPid = 0;
+
+	LWLockAcquire(&BaseWorkerControl->lock, LW_SHARED);
+
+	bool		isFound = false;
+	BaseWorkerEntry *workerEntry = GetBaseWorkerEntry(MyDatabaseId, workerId, &isFound);
+
+	if (isFound)
+		workerPid = workerEntry->workerPid;
+
+	LWLockRelease(&BaseWorkerControl->lock);
+
+	return workerPid;
+}
+
+
+/*
+ * pg_extension_base_get_worker_pid implements the extension_base.get_worker_pid
+ * SQL function.
+ */
+Datum
+pg_extension_base_get_worker_pid(PG_FUNCTION_ARGS)
+{
+	int32		workerId = PG_GETARG_INT32(0);
+
+	PG_RETURN_INT32(GetBaseWorkerPid(workerId));
+}
+
+
 /*
  * pg_extension_base_list_base_workers returns the state of the base workers in shared
  * memory.

pg_extension_base/tests/pytests/test_base_worker_launcher.py

@@ -431,6 +431,41 @@ def test_oneshot_worker_drop_database(superuser_conn):
     superuser_conn.autocommit = False
 
 
+def test_get_worker_pid_running(superuser_conn):
+    run_command(
+        "CREATE EXTENSION pg_extension_base_test_scheduler CASCADE", superuser_conn
+    )
+    superuser_conn.commit()
+    time.sleep(0.1)
+
+    worker_id = run_query(
+        "SELECT worker_id FROM extension_base.workers WHERE worker_name = 'pg_extension_base_test_scheduler_main_worker'",
+        superuser_conn,
+    )[0][0]
+
+    pid = run_query(
+        f"SELECT extension_base.get_worker_pid({worker_id})",
+        superuser_conn,
+    )[0][0]
+    assert pid > 0
+
+    run_command(
+        "DROP EXTENSION pg_extension_base_test_scheduler CASCADE", superuser_conn
+    )
+    superuser_conn.commit()
+
+
+def test_get_worker_pid_not_found(superuser_conn):
+    pid = run_query(
+        "SELECT extension_base.get_worker_pid(99999)",
+        superuser_conn,
+    )[
+        0
+    ][0]
+    assert pid == 0
+    superuser_conn.rollback()
+
+
 def get_pg_extension_workers(conn):
     query = "SELECT datname, application_name FROM pg_stat_activity WHERE backend_type LIKE 'pg_base_extension server %' OR backend_type LIKE 'pg base extension%' ORDER BY application_name"
     result = run_query(query, conn)