Allow subextensions to use DuckDB execution

Author: dpxcc

include/pgduckdb/pgduckdb_ruleutils.h

@@ -15,8 +15,8 @@ char *pgduckdb_get_tabledef(Oid relation_id);
 char *pgduckdb_get_alter_tabledef(Oid relation_oid, AlterTableStmt *alter_stmt);
 char *pgduckdb_get_rename_relationdef(Oid relation_oid, RenameStmt *rename_stmt);
 bool pgduckdb_is_not_default_expr(Node *node, void *context);
-List *pgduckdb_db_and_schema(const char *postgres_schema_name, bool is_duckdb_table);
-const char *pgduckdb_db_and_schema_string(const char *postgres_schema_name, bool is_duckdb_table);
+List *pgduckdb_db_and_schema(const char *postgres_schema_name, const char *duckdb_table_am_name);
+const char *pgduckdb_db_and_schema_string(const char *postgres_schema_name, const char *duckdb_table_am_name);
 bool pgduckdb_is_duckdb_row(Oid type_oid);
 bool pgduckdb_is_unresolved_type(Oid type_oid);
 bool pgduckdb_is_fake_type(Oid type_oid);

include/pgduckdb/pgduckdb_table_am.hpp

@@ -2,4 +2,8 @@
 
 namespace pgduckdb {
 bool IsDuckdbTableAm(const TableAmRoutine *am);
-}
+
+const char *DuckdbTableAmGetName(const TableAmRoutine *am);
+
+const char *DuckdbTableAmGetName(Oid relid);
+} // namespace pgduckdb

src/pgduckdb_ddl.cpp

@@ -968,7 +968,7 @@ DuckdbHandleViewStmtPre(Node *parsetree, PlannedStmt *pstmt, const char *query_s
 	/* END OF COPIED LOGIC */
 
 	char *duckdb_query_string = pgduckdb_get_querydef((Query *)copyObjectImpl(viewParse));
-	List *db_and_schema = pgduckdb_db_and_schema(schema_name, true);
+	List *db_and_schema = pgduckdb_db_and_schema(schema_name, "duckdb");
 
 	char *duckdb_db = (char *)linitial(db_and_schema);
 	char *duckdb_schema = (char *)lsecond(db_and_schema);
@@ -1002,7 +1002,7 @@ DuckdbHandleViewStmtPost(Node *parsetree) {
 	auto default_db = pgduckdb::DuckDBManager::Get().GetDefaultDBName();
 	char *postgres_schema_name = get_namespace_name(rel->rd_rel->relnamespace);
 
-	const char *duckdb_db = (const char *)linitial(pgduckdb_db_and_schema(postgres_schema_name, true));
+	const char *duckdb_db = (const char *)linitial(pgduckdb_db_and_schema(postgres_schema_name, "duckdb"));
 	relation_close(rel, NoLock);
 
 	Oid arg_types[] = {OIDOID, TEXTOID, TEXTOID, TEXTOID};
@@ -1294,7 +1294,7 @@ DECLARE_PG_FUNCTION(duckdb_create_table_trigger) {
 		Oid saved_userid;
 		int sec_context;
 		const char *postgres_schema_name = get_namespace_name_or_temp(get_rel_namespace(relid));
-		const char *duckdb_db = (const char *)linitial(pgduckdb_db_and_schema(postgres_schema_name, true));
+		const char *duckdb_db = (const char *)linitial(pgduckdb_db_and_schema(postgres_schema_name, "duckdb"));
 		auto default_db = pgduckdb::DuckDBManager::Get().GetDefaultDBName();
 
 		Oid arg_types[] = {OIDOID, TEXTOID, TEXTOID, TEXTOID};
@@ -1588,7 +1588,7 @@ DECLARE_PG_FUNCTION(duckdb_drop_trigger) {
 			char *object_type = SPI_getvalue(tuple, SPI_tuptable->tupdesc, 3);
 			char *drop_query =
 			    psprintf("DROP %s IF EXISTS %s.%s", object_type,
-			             pgduckdb_db_and_schema_string(postgres_schema_name, true), quote_identifier(table_name));
+			             pgduckdb_db_and_schema_string(postgres_schema_name, "duckdb"), quote_identifier(table_name));
 			pgduckdb::DuckDBQueryOrThrow(*connection, drop_query);
 
 			deleted_duckdb_relations++;

src/pgduckdb_hooks.cpp

@@ -66,14 +66,7 @@ ContainsCatalogTable(List *rtes) {
 
 static bool
 IsDuckdbTable(Oid relid) {
-	if (relid == InvalidOid) {
-		return false;
-	}
-
-	auto rel = RelationIdGetRelation(relid);
-	bool result = pgduckdb::IsDuckdbTableAm(rel->rd_tableam);
-	RelationClose(rel);
-	return result;
+	return pgduckdb::DuckdbTableAmGetName(relid) != nullptr;
 }
 
 static bool

src/pgduckdb_ruleutils.cpp

@@ -472,11 +472,15 @@ pgduckdb_write_row_refname(StringInfo buf, char *refname, bool is_top_level) {
  * are not escaped yet.
  */
 List *
-pgduckdb_db_and_schema(const char *postgres_schema_name, bool is_duckdb_table) {
-	if (!is_duckdb_table) {
+pgduckdb_db_and_schema(const char *postgres_schema_name, const char *duckdb_table_am_name) {
+	if (duckdb_table_am_name == nullptr) {
 		return list_make2((void *)"pgduckdb", (void *)postgres_schema_name);
 	}
 
+	if (strcmp("duckdb", duckdb_table_am_name) != 0) {
+		return list_make2((void *)duckdb_table_am_name, (void *)postgres_schema_name);
+	}
+
 	if (strcmp("pg_temp", postgres_schema_name) == 0) {
 		return list_make2((void *)"pg_temp", (void *)"main");
 	}
@@ -538,8 +542,8 @@ pgduckdb_db_and_schema(const char *postgres_schema_name, bool is_duckdb_table) {
  * database are quoted if necessary.
  */
 const char *
-pgduckdb_db_and_schema_string(const char *postgres_schema_name, bool is_duckdb_table) {
-	List *db_and_schema = pgduckdb_db_and_schema(postgres_schema_name, is_duckdb_table);
+pgduckdb_db_and_schema_string(const char *postgres_schema_name, const char *duckdb_table_am_name) {
+	List *db_and_schema = pgduckdb_db_and_schema(postgres_schema_name, duckdb_table_am_name);
 	const char *db_name = (const char *)linitial(db_and_schema);
 	const char *schema_name = (const char *)lsecond(db_and_schema);
 	return psprintf("%s.%s", quote_identifier(db_name), quote_identifier(schema_name));
@@ -558,9 +562,9 @@ pgduckdb_relation_name(Oid relation_oid) {
 	Form_pg_class relation = (Form_pg_class)GETSTRUCT(tp);
 	const char *relname = NameStr(relation->relname);
 	const char *postgres_schema_name = get_namespace_name_or_temp(relation->relnamespace);
-	bool is_duckdb_table = pgduckdb::IsDuckdbTable(relation);
+	const char *duckdb_table_am_name = pgduckdb::DuckdbTableAmGetName(relation_oid);
 
-	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, is_duckdb_table);
+	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, duckdb_table_am_name);
 
 	char *result = psprintf("%s.%s", db_and_schema, quote_identifier(relname));
 
@@ -611,7 +615,8 @@ pgduckdb_get_tabledef(Oid relation_oid) {
 	Relation relation = relation_open(relation_oid, AccessShareLock);
 	const char *relation_name = pgduckdb_relation_name(relation_oid);
 	const char *postgres_schema_name = get_namespace_name_or_temp(relation->rd_rel->relnamespace);
-	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, pgduckdb::IsDuckdbTable(relation));
+	const char *duckdb_table_am_name = pgduckdb::DuckdbTableAmGetName(relation->rd_tableam);
+	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, duckdb_table_am_name);
 
 	StringInfoData buffer;
 	initStringInfo(&buffer);
@@ -799,7 +804,7 @@ pgduckdb_get_viewdef(const ViewStmt *stmt, const char *postgres_schema_name, con
 	StringInfoData buffer;
 	initStringInfo(&buffer);
 
-	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, true);
+	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, "duckdb");
 	appendStringInfo(&buffer, "CREATE SCHEMA IF NOT EXISTS %s; ", db_and_schema);
 
 	appendStringInfoString(&buffer, "CREATE ");
@@ -879,7 +884,7 @@ pgduckdb_get_rename_relationdef(Oid relation_oid, RenameStmt *rename_stmt) {
 	Assert(pgduckdb::IsDuckdbTable(relation) || pgduckdb::IsMotherDuckView(relation));
 
 	const char *postgres_schema_name = get_namespace_name_or_temp(relation->rd_rel->relnamespace);
-	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, true);
+	const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, "duckdb");
 	const char *old_table_name = psprintf("%s.%s", db_and_schema, quote_identifier(rename_stmt->relation->relname));
 
 	const char *relation_type = "TABLE";

src/pgduckdb_table_am.cpp

@@ -10,6 +10,9 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  */
 
+#include "duckdb/common/string.hpp"
+#include "duckdb/common/unordered_map.hpp"
+
 #include "pgduckdb/pgduckdb_ddl.hpp"
 
 extern "C" {
@@ -492,9 +495,35 @@ duckdb_am_handler(FunctionCallInfo /*funcinfo*/) {
 }
 }
 
+static duckdb::unordered_map<const TableAmRoutine * /*am*/, duckdb::string /*name*/> duckdb_table_ams = {
+    {&duckdb_methods, "duckdb"}};
+
+extern "C" __attribute__((visibility("default"))) void
+RegisterDuckdbTableAm(const char *name, const TableAmRoutine *am) {
+	duckdb_table_ams[am] = name;
+}
+
 namespace pgduckdb {
 bool
 IsDuckdbTableAm(const TableAmRoutine *am) {
 	return am == &duckdb_methods;
 }
+
+const char *
+DuckdbTableAmGetName(const TableAmRoutine *am) {
+	auto it = duckdb_table_ams.find(am);
+	return it == duckdb_table_ams.end() ? nullptr : it->second.c_str();
+}
+
+const char *
+DuckdbTableAmGetName(Oid relid) {
+	if (relid == InvalidOid) {
+		return nullptr;
+	}
+
+	auto rel = RelationIdGetRelation(relid);
+	const char *name = DuckdbTableAmGetName(rel->rd_tableam);
+	RelationClose(rel);
+	return name;
+}
 } // namespace pgduckdb