Add DuckDB temporary tables

include/pgduckdb/pgduckdb_ddl.hpp

@@ -0,0 +1,5 @@
+extern "C" {
+#include "postgres.h"
+#include "nodes/nodes.h"
+}
+void DuckdbHandleDDL(Node *ParseTree, const char *queryString);

include/pgduckdb/pgduckdb_metadata_cache.hpp

@@ -5,4 +5,5 @@ extern "C" {
 namespace pgduckdb {
 bool IsExtensionRegistered();
 bool IsDuckdbOnlyFunction(Oid function_oid);
+Oid DuckdbTableAmOid();
 } // namespace pgduckdb

include/pgduckdb/pgduckdb_ruleutils.h

@@ -1,3 +1,5 @@
 #include "postgres.h"
 
+char *pgduckdb_relation_name(Oid relid);
 char *pgduckdb_function_name(Oid function_oid);
+char *pgduckdb_get_tabledef(Oid relation_id);

include/pgduckdb/pgduckdb_table_am.hpp

@@ -0,0 +1,8 @@
+extern "C" {
+#include "postgres.h"
+#include "access/tableam.h"
+}
+
+namespace pgduckdb {
+bool IsDuckdbTableAm(const TableAmRoutine *am);
+}

include/pgduckdb/utility/rename_ruleutils.h

@@ -20,3 +20,10 @@
 #define get_range_partbound_string       pgduckdb_get_range_partbound_string
 #define pg_get_statisticsobjdef_string   pgduckdb_pg_get_statisticsobjdef_string
 #define get_list_partvalue_string        pgduckdb_get_list_partvalue_string
+
+/*
+ * The following replaces all usages of generate_qualified_relation_name and
+ * generate_relation_name with calls to the pgduckdb_relation_name function
+ */
+#define generate_qualified_relation_name          pgduckdb_relation_name
+#define generate_relation_name(relid, namespaces) pgduckdb_relation_name(relid)

sql/pg_duckdb--0.0.1.sql

@@ -139,12 +139,38 @@ CREATE TABLE extensions (
     enabled BOOL DEFAULT TRUE
 );
 
+CREATE TABLE tables (
+    relid regclass PRIMARY KEY
+);
+
 CREATE OR REPLACE FUNCTION install_extension(extension_name TEXT) RETURNS bool
     LANGUAGE C AS 'MODULE_PATHNAME', 'install_extension';
 
 CREATE OR REPLACE FUNCTION raw_query(query TEXT) RETURNS void
     LANGUAGE C AS 'MODULE_PATHNAME', 'pgduckdb_raw_query';
 
+CREATE FUNCTION duckdb_am_handler(internal)
+RETURNS table_am_handler
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
+
+CREATE ACCESS METHOD duckdb
+    TYPE TABLE
+    HANDLER duckdb_am_handler;
+
+CREATE FUNCTION duckdb_drop_table_trigger() RETURNS event_trigger
+    AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE EVENT TRIGGER duckdb_drop_table_trigger ON sql_drop
+    EXECUTE FUNCTION duckdb_drop_table_trigger();
+
+CREATE FUNCTION duckdb_create_table_trigger() RETURNS event_trigger
+    AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE EVENT TRIGGER duckdb_create_table_trigger ON ddl_command_end
+    WHEN tag IN ('CREATE TABLE', 'CREATE TABLE AS')
+    EXECUTE FUNCTION duckdb_create_table_trigger();
+
 DO $$
 BEGIN
     RAISE WARNING 'To actually execute queries using DuckDB you need to run "SET duckdb.execution TO true;"';

src/pgduckdb_ddl.cpp

@@ -0,0 +1,179 @@
+#include "duckdb.hpp"
+#include <regex>
+
+extern "C" {
+#include "postgres.h"
+#include "access/tableam.h"
+#include "catalog/pg_type.h"
+#include "commands/event_trigger.h"
+#include "fmgr.h"
+#include "catalog/pg_authid_d.h"
+#include "funcapi.h"
+#include "nodes/print.h"
+#include "nodes/makefuncs.h"
+#include "optimizer/optimizer.h"
+#include "tcop/utility.h"
+#include "utils/syscache.h"
+#include "commands/event_trigger.h"
+#include "executor/spi.h"
+#include "miscadmin.h"
+
+#include "pgduckdb/pgduckdb_ruleutils.h"
+}
+
+#include "pgduckdb/pgduckdb_duckdb.hpp"
+#include "pgduckdb/vendor/pg_list.hpp"
+#include <inttypes.h>
+
+void
+DuckdbHandleDDL(Node *parsetree, const char *queryString) {
+	if (IsA(parsetree, CreateTableAsStmt)) {
+		auto stmt = castNode(CreateTableAsStmt, parsetree);
+		char *access_method = stmt->into->accessMethod ? stmt->into->accessMethod : default_table_access_method;
+		if (strcmp(access_method, "duckdb") != 0) {
+			/* not a duckdb table, so don't mess with the query */
+			return;
+		}
+
+		elog(ERROR, "DuckDB does not support CREATE TABLE AS yet");
+	}
+}
+
+extern "C" {
+
+PG_FUNCTION_INFO_V1(duckdb_create_table_trigger);
+
+Datum
+duckdb_create_table_trigger(PG_FUNCTION_ARGS) {
+	if (!CALLED_AS_EVENT_TRIGGER(fcinfo)) /* internal error */
+		elog(ERROR, "not fired by event trigger manager");
+
+	SPI_connect();
+
+	/* Temporarily escalate privileges to superuser so we can insert into duckdb.tables */
+	Oid saved_userid;
+	int sec_context;
+	GetUserIdAndSecContext(&saved_userid, &sec_context);
+	SetUserIdAndSecContext(BOOTSTRAP_SUPERUSERID, sec_context | SECURITY_LOCAL_USERID_CHANGE);
+	/*
+	 * We track the table oid in duckdb.tables so we can later check in our
+	 * delete event trigger if the table was created using the duckdb access
+	 * method. See the code comment in that function for more details.
+	 */
+	int ret = SPI_exec(R"(
+                INSERT INTO duckdb.tables(relid)
+                SELECT objid
+                FROM pg_catalog.pg_event_trigger_ddl_commands() cmds
+                JOIN pg_catalog.pg_class
+                ON cmds.objid = pg_class.oid
+                WHERE cmds.object_type = 'table'
+                AND pg_class.relam = (SELECT oid FROM pg_am WHERE amname = 'duckdb')
+				RETURNING relid)",
+	                   0);
+
+	/* Revert back to original privileges */
+	SetUserIdAndSecContext(saved_userid, sec_context);
+
+	if (ret != SPI_OK_INSERT_RETURNING)
+		elog(ERROR, "SPI_exec failed: error code %s", SPI_result_code_string(ret));
+
+	/* if we inserted a row it was a duckdb table */
+	auto isDuckdbTable = SPI_processed > 0;
+
+	if (!isDuckdbTable) {
+		SPI_finish();
+		PG_RETURN_NULL();
+	}
+	if (SPI_processed != 1) {
+		elog(ERROR, "Expected single table to be created, but found %" PRIu64, SPI_processed);
+	}
+	HeapTuple tuple = SPI_tuptable->vals[0];
+	bool isnull;
+	Datum relid_datum = SPI_getbinval(tuple, SPI_tuptable->tupdesc, 1, &isnull);
+	if (isnull) {
+		elog(ERROR, "Expected relid to be returned, but found NULL");
+	}
+	Oid relid = DatumGetObjectId(relid_datum);
+	SPI_finish();
+
+	const char *query_cstr = pgduckdb_get_tabledef(relid);
+	std::string query_string(query_cstr);
+
+	auto db = pgduckdb::DuckDBManager::Get().GetDatabase();
+	auto connection = duckdb::make_uniq<duckdb::Connection>(db);
+	auto &context = *connection->context;
+	auto result = context.Query(query_string, false);
+	if (result->HasError()) {
+		elog(ERROR, "(PGDuckDB/duckdb_create_table_trigger) Could not create table: %s", result->GetError().c_str());
+	}
+
+	PG_RETURN_NULL();
+}
+PG_FUNCTION_INFO_V1(duckdb_drop_table_trigger);
+
+Datum
+duckdb_drop_table_trigger(PG_FUNCTION_ARGS) {
+	if (!CALLED_AS_EVENT_TRIGGER(fcinfo)) /* internal error */
+		elog(ERROR, "not fired by event trigger manager");
+
+	SPI_connect();
+
+	/* Temporarily escalate privileges to superuser so we can delete from duckdb.tables */
+	Oid saved_userid;
+	int sec_context;
+	GetUserIdAndSecContext(&saved_userid, &sec_context);
+	SetUserIdAndSecContext(BOOTSTRAP_SUPERUSERID, sec_context | SECURITY_LOCAL_USERID_CHANGE);
+
+	// Because the table metadata is deleted from the postgres catalogs we
+	// cannot find out if the table was using the duckdb access method. So
+	// instead we keep our own metadata table that also tracks which tables are
+	// duckdb tables.
+	int ret = SPI_exec(R"(
+                DELETE FROM duckdb.tables
+                USING (
+                    SELECT objid, object_name, object_identity
+                    FROM pg_catalog.pg_event_trigger_dropped_objects()
+                    WHERE object_type = 'table'
+                ) objs
+                WHERE relid = objid
+                RETURNING objs.object_identity
+                )",
+	                   0);
+
+	/* Revert back to original privileges */
+	SetUserIdAndSecContext(saved_userid, sec_context);
+
+	if (ret != SPI_OK_DELETE_RETURNING)
+		elog(ERROR, "SPI_exec failed: error code %s", SPI_result_code_string(ret));
+
+	auto db = pgduckdb::DuckDBManager::Get().GetDatabase();
+	auto connection = duckdb::make_uniq<duckdb::Connection>(db);
+	auto &context = *connection->context;
+
+	auto result = context.Query("BEGIN TRANSACTION", false);
+	if (result->HasError()) {
+		elog(ERROR, "(PGDuckDB/duckdb_drop_table_trigger) Could not start transaction");
+	}
+
+	for (auto proc = 0; proc < SPI_processed; proc++) {
+		HeapTuple tuple = SPI_tuptable->vals[proc];
+
+		char *object_identity = SPI_getvalue(tuple, SPI_tuptable->tupdesc, 1);
+		// TODO: Handle std::string creation in a safe way for allocation failures
+		auto result = context.Query("DROP TABLE IF EXISTS " + std::string(object_identity), false);
+		if (result->HasError()) {
+			elog(ERROR, "(PGDuckDB/duckdb_drop_table_trigger) Could not drop table %s: %s", object_identity,
+			     result->GetError().c_str());
+		}
+	}
+
+	SPI_finish();
+
+	result = context.Query("COMMIT", false);
+	if (result->HasError()) {
+		elog(ERROR, "(PGDuckDB/duckdb_drop_table_trigger) Could not commit transaction");
+	}
+
+	PG_RETURN_NULL();
+}
+}

src/pgduckdb_duckdb.cpp

@@ -13,6 +13,11 @@
 #include "pgduckdb/pgduckdb_utils.hpp"
 #include "pgduckdb/catalog/pgduckdb_storage.hpp"
 
+extern "C" {
+#include "postgres.h"
+
+#include "utils/elog.h"
+}
 #include <string>
 
 #include <sys/types.h>
@@ -173,10 +178,6 @@ DuckdbCreateConnection(List *rtables, PlannerInfo *planner_info, List *needed_co
 
 	context.registered_state->Insert(
 	    "postgres_state", duckdb::make_shared_ptr<PostgresContextState>(rtables, planner_info, needed_columns, query));
-	auto res = context.Query("set search_path='pgduckdb.main'", false);
-	if (res->HasError()) {
-		elog(WARNING, "(DuckDB) %s", res->GetError().c_str());
-	}
 	return con;
 }