Navigation Menu

Skip to content

Instantly share code, notes, and snippets.

@steve-chavez
Forked from ArtemGr/makefile
Created June 22, 2022 17:40
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save steve-chavez/7e84d85a80c23554a0ef639fdf69e5db to your computer and use it in GitHub Desktop.
Save steve-chavez/7e84d85a80c23554a0ef639fdf69e5db to your computer and use it in GitHub Desktop.
SPI helper for a PostgreSQL CHECK when using one-to-many with a jsonb.
all: spi.so
spi.so: spi.o makefile
g++ -shared -o spi.so spi.o
cp --remove-destination spi.so /var/lib/postgresql/spi.so
spi.o: spi.cc makefile
g++ -g -O2 -Wall -std=c++11 -fpic -c -o spi.o -I/usr/include/postgresql -I/usr/include/postgresql/9.4/server spi.cc
clean:
rm -f spi.o spi.so
#include <iostream>
#include <glim/gstring.hpp> // https://code.google.com/p/libglim/
extern "C" {
#include <postgres.h>
#include <executor/spi.h> // postgresql-server-dev-9.4
#include <utils/builtins.h>
#include <utils/jsonb.h>
}
static inline int32_t numericToInt32 (Numeric numeric) {
return DatumGetInt32 (DirectFunctionCall1 (numeric_int4, NumericGetDatum (numeric)));
}
// http://www.postgresql.org/docs/9.4/interactive/xfunc-c.html
// http://www.postgresql.org/docs/9.4/interactive/spi.html
extern "C" {
#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif
void _PG_init() {
//std::cout << "have_ids spi init." << std::endl; // Goes into /var/log/postgresql/postgresql-9.4-main.log
}
PG_FUNCTION_INFO_V1 (have_ids); // http://www.postgresql.org/docs/9.4/interactive/xfunc-c.html
// Register with (su - postgres; psql db):
// CREATE FUNCTION have_ids (cstring, jsonb) RETURNS BOOLEAN AS '/var/lib/postgresql/spi', 'have_ids' LANGUAGE C STABLE STRICT COST 100;
// cf. http://www.postgresql.org/docs/9.4/interactive/sql-createfunction.html
// Check with:
// echo "SELECT have_ids ('foo', '[1, 2, 3]');" | sudo -u postgres psql db
/// Checks if all the ids from the given array are present in the given table.
/// SELECT have_ids ('foo', '[1, 2, 3]');
Datum have_ids (PG_FUNCTION_ARGS) try {
bool missingIds {false};
const char* table = PG_GETARG_CSTRING (0);
Jsonb* categories = PG_GETARG_JSONB (1);
JsonbIterator* it = JsonbIteratorInit (&categories->root); if (it) {
if (SPI_connect() != SPI_OK_CONNECT) GTHROW ("!SPI_connect");
JsonbValue val; JsonbIteratorToken tok;
while ((tok = JsonbIteratorNext (&it, &val, true)) != WJB_DONE) {
//std::cout << "have_ids] tok " << tok << std::endl; // WJB_BEGIN_ARRAY, WJB_ELEM, WJB_END_ARRAY.
if (tok == WJB_ELEM) {
if (val.type != JsonbValue::jbvNumeric) GTHROW ("!jbvNumeric: " + std::to_string (val.type));
int32_t id = numericToInt32 (val.val.numeric);
//std::cout << "have_ids] " << id << std::endl;
GSTRING_ON_STACK (sql, 256) << "SELECT COUNT(*) FROM " << table << " WHERE id = " << id;
//std::cout << "have_ids] " << sql << std::endl;
if (SPI_execute (sql.c_str(), true, 1) != SPI_OK_SELECT) GTHROW ("!SPI_execute");
if (SPI_processed != 1) GTHROW ("!SPI_processed"); // Number of rows processed.
if (SPI_tuptable == nullptr) GTHROW ("!SPI_tuptable");
bool isNull {false};
Datum count = SPI_getbinval (SPI_tuptable->vals[0], SPI_tuptable->tupdesc, 1, &isNull);
if (isNull) GTHROW ("isNull");
if (DatumGetInt32 (count) != 1) {missingIds = true; break;}
}
}
if (SPI_finish() != SPI_OK_FINISH) GTHROW ("!SPI_finish");
}
PG_RETURN_BOOL (!missingIds);
} catch (const std::exception& ex) {
GSTRING_ON_STACK (error, 256) << "have_ids] " << ex.what();
elog (ERROR, error.c_str());
}
void _PG_fini() {
//std::cout << "have_ids spi fini." << std::endl;
}
}
@steve-chavez
Copy link
Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment