Skip to content

Commit 2f3c6e8

Browse files
sfc-gh-okalacisfc-gh-okalaci
committed
Move CachedExtensionIds and SPI helpers to pg_extension_base
pg_extension_base is the foundational extension in the dependency chain, making it the natural home for shared infrastructure used by all extensions. This eliminates code duplication and allows pg_extension_base itself to use proper SPI helpers with extension owner context switching. Specific changes: - Move extension_ids.h/c (CachedExtensionIds infrastructure) from pg_lake_engine - Move spi_helpers.h from pg_lake_engine - Add PgExtensionBase cache in separate pg_extension_base_ids.h/c files - Update SPI_START_EXTENSION_OWNER macro to work with any CachedExtensionIds - Use SPI_START_EXTENSION_OWNER(PgExtensionBase) in base_worker_launcher.c for register/deregister functions to support non-superuser callers - Update all extension Makefiles to include pg_extension_base headers
1 parent 2bc8648 commit 2f3c6e8

59 files changed

Lines changed: 239 additions & 244 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

pg_lake_engine/include/pg_lake/extensions/extension_ids.h renamed to pg_extension_base/include/pg_extension_base/extension_ids.h

File renamed without changes.

‎pg_extension_base/include/pg_extension_base/pg_extension_base_ids.h‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright 2025 Snowflake Inc.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
#pragma once
19+
20+
#include "pg_extension_base/extension_ids.h"
21+
22+
#define PG_EXTENSION_BASE_NAME "pg_extension_base"
23+
24+
/* cached extension IDs for pg_extension_base */
25+
extern PGDLLEXPORT CachedExtensionIds *PgExtensionBase;
26+
27+
extern PGDLLEXPORT void InitializePgExtensionBaseCache(void);

‎pg_extension_base/include/pg_extension_base/spi_helpers.h‎

Lines changed: 125 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,87 @@
1818
#ifndef SPI_UTILITIES_H
1919
#define SPI_UTILITIES_H
2020

21+
#include "miscadmin.h"
22+
2123
#include "executor/spi.h"
24+
#include "utils/builtins.h"
25+
#include "utils/guc.h"
26+
#include "utils/jsonb.h"
27+
#include "utils/pg_lsn.h"
28+
29+
/* SPI macros for setting parameters */
30+
#define DATUMIZE_TEXTOID(Value) CStringGetTextDatum(Value)
31+
#define DATUMIZE_TEXTARRAYOID(Value) PointerGetDatum(Value)
32+
#define DATUMIZE_CHAROID(Value) CharGetDatum(Value)
33+
#define DATUMIZE_BOOLOID(Value) BoolGetDatum(Value)
34+
#define DATUMIZE_OIDOID(Value) ObjectIdGetDatum(Value)
35+
#define DATUMIZE_INT2OID(Value) Int16GetDatum(Value)
36+
#define DATUMIZE_INT2ARRAYOID(Value) PointerGetDatum(Value)
37+
#define DATUMIZE_INT4OID(Value) Int32GetDatum(Value)
38+
#define DATUMIZE_INT4ARRAYOID(Value) PointerGetDatum(Value)
39+
#define DATUMIZE_INT8OID(Value) Int64GetDatum(Value)
40+
#define DATUMIZE_INT8ARRAYOID(Value) PointerGetDatum(Value)
41+
#define DATUMIZE_FLOAT4OID(Value) Float4GetDatum(Value)
42+
#define DATUMIZE_FLOAT4ARRAYOID(Value) PointerGetDatum(Value)
43+
#define DATUMIZE_FLOAT8OID(Value) Float8GetDatum(Value)
44+
#define DATUMIZE_FLOAT8ARRAYOID(Value) PointerGetDatum(Value)
45+
#define DATUMIZE_TIMESTAMPTZOID(Value) TimestampTzGetDatum(Value)
46+
#define DATUMIZE_JSONOID(Value) CStringGetTextDatum(Value)
47+
#define DATUMIZE_JSONBOID(Value) JsonbPGetDatum(Value)
48+
#define DATUMIZE_BYTEAOID(Value) PointerGetDatum(Value)
49+
#define DATUMIZE_BYTEAARRAYOID(Value) PointerGetDatum(Value)
50+
#define DATUMIZE_LSNOID(Value) LSNGetDatum(Value)
51+
#define DATUMIZE_NAMEOID(Value) NameGetDatum(Value)
52+
#define DATUMIZE(TypeId,Value) DATUMIZE_##TypeId(Value)
53+
54+
#define DECLARE_SPI_ARGS(Count) \
55+
const int spiArgCount = Count; \
56+
Oid spiArgTypes[Count]; \
57+
Datum spiArgValues[Count]; \
58+
char spiArgNulls[Count]; \
59+
memset(spiArgNulls, ' ', Count);
60+
61+
#define SPI_ARG_DATUM(Index,TypeId,Datum) \
62+
{ \
63+
spiArgTypes[Index-1] = TypeId; \
64+
spiArgValues[Index-1] = Datum; \
65+
spiArgNulls[Index-1] = ' '; \
66+
}
67+
68+
#define SPI_ARG_NULL(Index,TypeId) \
69+
{ \
70+
spiArgTypes[Index-1] = TypeId; \
71+
spiArgNulls[Index-1] = 'n'; \
72+
}
73+
74+
#define SPI_ARG_VALUE(Index,TypeId,Value,IsNull) \
75+
{ \
76+
if (IsNull) \
77+
{ \
78+
SPI_ARG_NULL(Index,TypeId); \
79+
} \
80+
else \
81+
{ \
82+
SPI_ARG_DATUM(Index,TypeId,DATUMIZE_##TypeId(Value)); \
83+
} \
84+
}
85+
86+
#define SPI_EXECUTE(Query,ReadOnly) \
87+
{ \
88+
int spiStatus = SPI_execute_with_args(Query, spiArgCount, spiArgTypes, \
89+
spiArgValues, spiArgNulls, ReadOnly, 0); \
90+
if (spiStatus < 0) \
91+
{ \
92+
ereport(ERROR, (errmsg("failed to execute SPI query"))); \
93+
} \
94+
}
2295

2396
/* SPI macros for reading results */
2497
#define DEDATUMIZE_TEXTOID(Value) TextDatumGetCString(Value)
25-
#define DEDATUMIZE_NAMEOID(Value) DatumGetName(Value)
98+
#define DEDATUMIZE_TEXTARRAYOID(Value) DatumGetArrayTypeP(Value)
99+
#define DEDATUMIZE_CHAROID(Value) DatumGetChar(Value)
26100
#define DEDATUMIZE_BOOLOID(Value) DatumGetBool(Value)
101+
#define DEDATUMIZE_OIDOID(Value) DatumGetObjectId(Value)
27102
#define DEDATUMIZE_INT2OID(Value) DatumGetInt16(Value)
28103
#define DEDATUMIZE_INT2ARRAYOID(Value) DatumGetArrayTypeP(Value)
29104
#define DEDATUMIZE_INT4OID(Value) DatumGetInt32(Value)
@@ -34,9 +109,13 @@
34109
#define DEDATUMIZE_FLOAT4ARRAYOID(Value) DatumGetArrayTypeP(Value)
35110
#define DEDATUMIZE_FLOAT8OID(Value) DatumGetFloat8(Value)
36111
#define DEDATUMIZE_FLOAT8ARRAYOID(Value) DatumGetArrayTypeP(Value)
112+
#define DEDATUMIZE_TIMESTAMPTZOID(Value) DatumGetTimestampTz(Value)
113+
#define DEDATUMIZE_JSONOID(Value) TextDatumGetCString(Value)
37114
#define DEDATUMIZE_JSONBOID(Value) DatumGetJsonbP(Value)
38115
#define DEDATUMIZE_BYTEAOID(Value) DatumGetByteaP(Value)
39116
#define DEDATUMIZE_BYTEAARRAYOID(Value) DatumGetArrayTypeP(Value)
117+
#define DEDATUMIZE_LSNOID(Value) DatumGetLSN(Value)
118+
#define DEDATUMIZE_NAMEOID(Value) DatumGetName(Value)
40119
#define DEDATUMIZE(TypeId,Value) DATUMIZE_##TypeId(Value)
41120

42121
#define GET_SPI_DATUM(RowIndex,ColumnNumber,IsNull) \
@@ -46,4 +125,49 @@
46125
#define GET_SPI_VALUE(TypeId,RowIndex,ColumnNumber,IsNull) \
47126
DEDATUMIZE_##TypeId(GET_SPI_DATUM(RowIndex,ColumnNumber,IsNull))
48127

128+
#define SPI_START_VARS() \
129+
Oid _savedUserId = InvalidOid; \
130+
int _savedSecurityContext = 0;
131+
132+
/*
133+
* We do not want auto_explain to concern itself with our SPI queries,
134+
* since it might lead to failure, and unnecessary spilling of internals
135+
* into logs.
136+
*/
137+
#define DISABLE_QUERY_TRACKING() \
138+
int _spiGUCNestLevel = NewGUCNestLevel(); \
139+
(void) set_config_option("auto_explain.log_min_duration", "-1", PGC_SUSET, PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);\
140+
(void) set_config_option("pgaudit.log", "none", PGC_SUSET, PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);\
141+
(void) set_config_option("pg_stat_statements.track", "none", PGC_SUSET, PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
142+
143+
144+
#define RESET_QUERY_TRACKING() \
145+
AtEOXact_GUC(true, _spiGUCNestLevel);
146+
147+
#define SPI_START() \
148+
SPI_START_VARS() \
149+
DISABLE_QUERY_TRACKING() \
150+
SPI_connect();
151+
152+
/*
153+
* SPI_START_EXTENSION_OWNER runs SPI queries as the extension owner.
154+
*
155+
* This requires the caller to have ExtensionOwnerId() available
156+
* (from pg_extension_base/extension_ids.h).
157+
*
158+
* Usage: SPI_START_EXTENSION_OWNER(PgLakeTable)
159+
*/
160+
#define SPI_START_EXTENSION_OWNER(Extension) \
161+
SPI_START_VARS() \
162+
GetUserIdAndSecContext(&_savedUserId, &_savedSecurityContext); \
163+
SetUserIdAndSecContext(ExtensionOwnerId(Extension), SECURITY_LOCAL_USERID_CHANGE); \
164+
DISABLE_QUERY_TRACKING() \
165+
SPI_connect();
166+
167+
#define SPI_END() \
168+
if (_savedUserId != InvalidOid) \
169+
SetUserIdAndSecContext(_savedUserId, _savedSecurityContext); \
170+
RESET_QUERY_TRACKING(); \
171+
SPI_finish();
172+
49173
#endif

‎pg_extension_base/src/base_worker_launcher.c‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@
9797
#include "tcop/utility.h"
9898

9999
#include "pg_extension_base/base_workers.h"
100+
#include "pg_extension_base/pg_extension_base_ids.h"
100101
#include "pg_extension_base/pg_compat.h"
101102
#include "pg_extension_base/spi_helpers.h"
102103

@@ -403,6 +404,9 @@ InitializeBaseWorkerLauncher(void)
403404
return;
404405
}
405406

407+
/* initialize the extension IDs cache for pg_extension_base */
408+
InitializePgExtensionBaseCache();
409+
406410
/* set up DDL hooks */
407411
PreviousProcessUtility =
408412
ProcessUtility_hook != NULL ? ProcessUtility_hook : standard_ProcessUtility;
@@ -2199,11 +2203,14 @@ DeregisterBaseWorker_internal(int32 workerId)
21992203

22002204
/*
22012205
* InsertBaseWorkerRegistration inserts an entry into pg_extension_base.workers.
2206+
*
2207+
* We run as the extension owner to allow non-superusers who have been granted
2208+
* EXECUTE on register_worker to insert into the workers table.
22022209
*/
22032210
static int32
22042211
InsertBaseWorkerRegistration(char *workerName, Oid extensionId, Oid entryPointFunctionId)
22052212
{
2206-
SPI_connect();
2213+
SPI_START_EXTENSION_OWNER(PgExtensionBase);
22072214

22082215
char *extensionName = get_extension_name(extensionId);
22092216

@@ -2243,7 +2250,7 @@ InsertBaseWorkerRegistration(char *workerName, Oid extensionId, Oid entryPointFu
22432250

22442251
int32 workerId = DatumGetInt32(workerIdDatum);
22452252

2246-
SPI_finish();
2253+
SPI_END();
22472254

22482255
return workerId;
22492256
}
@@ -2274,11 +2281,14 @@ DatabaseIsTemplate(Oid databaseId)
22742281
/*
22752282
* DeleteBaseWorkerRegistrationByName deletes an entry from pg_extension_base.workers
22762283
* by name.
2284+
*
2285+
* We run as the extension owner to allow non-superusers who have been granted
2286+
* EXECUTE on deregister_worker to delete from the workers table.
22772287
*/
22782288
static int32
22792289
DeleteBaseWorkerRegistrationByName(char *workerName)
22802290
{
2281-
SPI_connect();
2291+
SPI_START_EXTENSION_OWNER(PgExtensionBase);
22822292

22832293
int argCount = 1;
22842294
Oid argTypes[] = {TEXTOID};
@@ -2306,23 +2316,26 @@ DeleteBaseWorkerRegistrationByName(char *workerName)
23062316

23072317
int32 workerId = DatumGetInt32(workerIdDatum);
23082318

2309-
SPI_finish();
2319+
SPI_END();
23102320

23112321
return workerId;
23122322
}
23132323

23142324
/*
23152325
* DeleteBaseWorkerRegistrationById deletes an entry from pg_extension_base.workers
23162326
* by id.
2327+
*
2328+
* We run as the extension owner to allow non-superusers who have been granted
2329+
* EXECUTE on deregister_worker to delete from the workers table.
23172330
*/
23182331
static void
23192332
DeleteBaseWorkerRegistrationById(int32 workerId)
23202333
{
2321-
SPI_connect();
2334+
SPI_START_EXTENSION_OWNER(PgExtensionBase);
23222335

23232336
int argCount = 1;
23242337
Oid argTypes[] = {INT4OID};
2325-
Datum argValues[] = {workerId};
2338+
Datum argValues[] = {Int32GetDatum(workerId)};
23262339

23272340
const char *argNulls = " ";
23282341
bool readOnly = false;
@@ -2336,7 +2349,7 @@ DeleteBaseWorkerRegistrationById(int32 workerId)
23362349
if (SPI_processed != 1)
23372350
ereport(ERROR, (errmsg("could not find worker id %d", workerId)));
23382351

2339-
SPI_finish();
2352+
SPI_END();
23402353
}
23412354

23422355
/*

pg_lake_engine/src/extensions/extension_ids.c renamed to pg_extension_base/src/extension_ids.c

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,10 @@
2424
#include "access/htup_details.h"
2525
#include "access/table.h"
2626
#include "catalog/objectaccess.h"
27-
#include "catalog/objectaccess.h"
2827
#include "catalog/pg_extension.h"
2928
#include "catalog/pg_type.h"
3029
#include "commands/extension.h"
31-
#include "pg_lake/extensions/extension_ids.h"
30+
#include "pg_extension_base/extension_ids.h"
3231
#include "parser/parse_func.h"
3332
#include "utils/catcache.h"
3433
#include "utils/fmgroids.h"
@@ -261,7 +260,7 @@ ExtensionSchemaId(CachedExtensionIds * extension)
261260

262261

263262
/*
264-
* ExtensionOwnerId returns the OID of the extension schema.
263+
* ExtensionOwnerId returns the OID of the extension owner.
265264
*/
266265
Oid
267266
ExtensionOwnerId(CachedExtensionIds * extension)

pg_lake_engine/src/extensions/pg_extension_base.c renamed to pg_extension_base/src/pg_extension_base_ids.c

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,24 +16,24 @@
1616
*/
1717

1818
/*
19-
* Utility functions for pg_lake_benchmark extension OIDs.
19+
* Cached extension IDs for pg_extension_base.
2020
*/
2121
#include "postgres.h"
2222

23-
#include "pg_lake/extensions/pg_extension_base.h"
24-
#include "pg_lake/extensions/extension_ids.h"
23+
#include "pg_extension_base/pg_extension_base_ids.h"
2524

2625

2726
/*
28-
* Generic extension state.
27+
* Cached extension state for pg_extension_base.
2928
*/
3029
CachedExtensionIds *PgExtensionBase;
3130

31+
3232
/*
3333
* Set up extension ID caching for pg_extension_base.
3434
*/
3535
void
36-
InitializePGExtensionBaseCache(void)
36+
InitializePgExtensionBaseCache(void)
3737
{
38-
PgExtensionBase = CreateExtensionIdsCache(PG_EXTENSION_BASE, NULL, NULL);
38+
PgExtensionBase = CreateExtensionIdsCache(PG_EXTENSION_BASE_NAME, NULL, NULL);
3939
}

‎pg_lake_benchmark/Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ OBJS := $(patsubst %.c,%.o,$(sort $(SOURCES)))
88

99
SHLIB_LINK_INTERNAL = $(libpq)
1010

11-
PG_CPPFLAGS = -I$(libpq_srcdir) -Iinclude -I../pg_lake_engine/include -std=gnu99 -g
11+
PG_CPPFLAGS = -I$(libpq_srcdir) -Iinclude -I../pg_extension_base/include -I../pg_lake_engine/include -std=gnu99 -g
1212

1313
UNAME_S := $(shell uname -s)
1414
ifeq ($(UNAME_S),Darwin)

‎pg_lake_benchmark/src/benchmark.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
#include "pg_lake/benchmark.h"
2323
#include "pg_lake/pgduck/client.h"
24-
#include "pg_lake/util/spi_helpers.h"
24+
#include "pg_extension_base/spi_helpers.h"
2525

2626

2727
/*

‎pg_lake_copy/Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ SOURCES := $(wildcard src/*.c) $(wildcard src/*/*.c)
77
OBJS := $(patsubst %.c,%.o,$(sort $(SOURCES)))
88

99
SHLIB_LINK = $(libpq)
10-
PG_CPPFLAGS = -I$(libpq_srcdir) -Iinclude -I../pg_lake_engine/include -I../pg_lake_iceberg/include -I../pg_lake_table/include -std=gnu99 -g -Wall -Wextra -Wno-unused-parameter -Werror
10+
PG_CPPFLAGS = -I$(libpq_srcdir) -Iinclude -I../pg_extension_base/include -I../pg_lake_engine/include -I../pg_lake_iceberg/include -I../pg_lake_table/include -std=gnu99 -g -Wall -Wextra -Wno-unused-parameter -Werror
1111

1212
UNAME_S := $(shell uname -s)
1313
ifeq ($(UNAME_S),Darwin)

‎pg_lake_engine/include/pg_lake/extensions/btree_gist.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
#pragma once
1919

20-
#include "pg_lake/extensions/extension_ids.h"
20+
#include "pg_extension_base/extension_ids.h"
2121

2222
#define BTREE_GIST "btree_gist"
2323

0 commit comments

Comments
 (0)