Add ALTER TABLE support for DuckDB tables (duckdb#652)

Adds all ALTER TABLE commands that DuckDB can parse except for COLLATION
changes and foreign keys. Not all of these ALTER commands are actually
supported by DuckDB (e.g. adding CHECK constraints). But this starts to
least forward those queries correctly to DuckDB, so once DuckDB supports
it pg_duckdb will start to automatically support it as well.

Fixes duckdb#536
Fixes duckdb#537
Fixes duckdb#538
Fixes duckdb#539
Fixes duckdb#540
Fixes duckdb#541
Fixes duckdb#542

---------

Co-authored-by: mkaruza <mkaruza@users.noreply.github.com>

Author: JelteF

include/pgduckdb/pgduckdb_ddl.hpp

@@ -2,5 +2,11 @@
 
 #include "pgduckdb/pg/declarations.hpp"
 
+namespace pgduckdb {
+enum class DDLType { NONE, CREATE_TABLE, ALTER_TABLE };
+/* Tracks the type of DDL statement that is currently being executed */
+extern DDLType top_level_duckdb_ddl_type;
+} // namespace pgduckdb
+
 void DuckdbTruncateTable(Oid relation_oid);
 void DuckdbInitUtilityHook();

include/pgduckdb/pgduckdb_ruleutils.h

@@ -12,6 +12,8 @@ char *pgduckdb_relation_name(Oid relid);
 char *pgduckdb_function_name(Oid function_oid, bool *use_variadic_p);
 char *pgduckdb_get_querydef(Query *);
 char *pgduckdb_get_tabledef(Oid relation_id);
+char *pgduckdb_get_alter_tabledef(Oid relation_oid, AlterTableStmt *alter_stmt);
+char *pgduckdb_get_rename_tabledef(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);

src/pgduckdb.cpp

@@ -10,6 +10,7 @@ extern "C" {
 #include "pgduckdb/pgduckdb.h"
 #include "pgduckdb/pgduckdb_node.hpp"
 #include "pgduckdb/pgduckdb_background_worker.hpp"
+#include "pgduckdb/pgduckdb_metadata_cache.hpp"
 #include "pgduckdb/pgduckdb_xact.hpp"
 
 static void DuckdbInitGUC(void);

src/pgduckdb_ddl.cpp

@@ -2,6 +2,7 @@
 #include "pgduckdb/pgduckdb_utils.hpp"
 #include "pgduckdb/pgduckdb_xact.hpp"
 #include "pgduckdb/pgduckdb_guc.h"
+#include "pgduckdb/pgduckdb_ddl.hpp"
 
 extern "C" {
 #include "postgres.h"
@@ -32,6 +33,7 @@ extern "C" {
 #include "pgduckdb/pgduckdb_ruleutils.h"
 }
 
+#include "pgduckdb/pgduckdb_guc.h"
 #include "pgduckdb/utility/cpp_wrapper.hpp"
 #include "pgduckdb/pgduckdb_duckdb.hpp"
 #include "pgduckdb/pgduckdb_background_worker.hpp"
@@ -40,6 +42,10 @@ extern "C" {
 #include "pgduckdb/vendor/pg_list.hpp"
 #include <inttypes.h>
 
+namespace pgduckdb {
+DDLType top_level_duckdb_ddl_type = DDLType::NONE;
+} // namespace pgduckdb
+
 /*
  * ctas_skip_data stores the original value of the skipData field of the
  * CreateTableAsStmt of the query that's currently being executed. For duckdb
@@ -181,6 +187,11 @@ DuckdbHandleDDL(PlannedStmt *pstmt, const char *query_string, ParamListInfo para
 		char *access_method = stmt->accessMethod ? stmt->accessMethod : default_table_access_method;
 		bool is_duckdb_table = strcmp(access_method, "duckdb") == 0;
 		if (is_duckdb_table) {
+			if (pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::NONE) {
+				ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+				                errmsg("Only one DuckDB table can be created in a single statement")));
+			}
+			pgduckdb::top_level_duckdb_ddl_type = pgduckdb::DDLType::CREATE_TABLE;
 			pgduckdb::ClaimCurrentCommandId();
 		}
 
@@ -202,6 +213,11 @@ DuckdbHandleDDL(PlannedStmt *pstmt, const char *query_string, ParamListInfo para
 		char *access_method = stmt->into->accessMethod ? stmt->into->accessMethod : default_table_access_method;
 		bool is_duckdb_table = strcmp(access_method, "duckdb") == 0;
 		if (is_duckdb_table) {
+			if (pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::NONE) {
+				ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+				                errmsg("Only one DuckDB table can be created in a single statement")));
+			}
+			pgduckdb::top_level_duckdb_ddl_type = pgduckdb::DDLType::CREATE_TABLE;
 			pgduckdb::ClaimCurrentCommandId();
 			/*
 			 * Force skipData to false for duckdb tables, so that Postgres does
@@ -313,6 +329,21 @@ DuckdbHandleDDL(PlannedStmt *pstmt, const char *query_string, ParamListInfo para
 			break;
 		}
 		return;
+	} else if (IsA(parsetree, AlterTableStmt)) {
+		auto stmt = castNode(AlterTableStmt, parsetree);
+		Oid relation_oid = RangeVarGetRelid(stmt->relation, AccessShareLock, false);
+		Relation relation = RelationIdGetRelation(relation_oid);
+		/*
+		 * Certain CREATE TABLE commands also trigger an ALTER TABLE command,
+		 * specifically if you use REFERENCES it will alter the table
+		 * afterwards. We currently only do this to get a better error message,
+		 * because we don't support REFERENCES anyway.
+		 */
+		if (pgduckdb::IsDuckdbTable(relation) && pgduckdb::top_level_duckdb_ddl_type == pgduckdb::DDLType::NONE) {
+			pgduckdb::top_level_duckdb_ddl_type = pgduckdb::DDLType::ALTER_TABLE;
+			pgduckdb::ClaimCurrentCommandId();
+		}
+		RelationClose(relation);
 	}
 }
 
@@ -416,6 +447,15 @@ DECLARE_PG_FUNCTION(duckdb_create_table_trigger) {
 		PG_RETURN_NULL();
 	}
 
+	if (pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::CREATE_TABLE &&
+	    pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::NONE) {
+		ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+		                errmsg("Cannot create a DuckDB table this way, use CREATE TABLE or CREATE TABLE ... AS")));
+		PG_RETURN_NULL();
+	}
+	/* Reset it back to NONE, for the remainder of the event trigger */
+	pgduckdb::top_level_duckdb_ddl_type = pgduckdb::DDLType::NONE;
+
 	EventTriggerData *trigger_data = (EventTriggerData *)fcinfo->context;
 	Node *parsetree = trigger_data->parsetree;
 
@@ -859,6 +899,16 @@ DECLARE_PG_FUNCTION(duckdb_alter_table_trigger) {
 		PG_RETURN_NULL();
 	}
 
+	/* Reset since we don't need it anymore */
+	if (pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::ALTER_TABLE &&
+	    pgduckdb::top_level_duckdb_ddl_type != pgduckdb::DDLType::NONE) {
+		ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+		                errmsg("Cannot ALTER a DuckDB table this way, please use ALTER TABLE")));
+		PG_RETURN_NULL();
+	}
+	/* Reset it back to NONE, for the remainder of the event trigger */
+	pgduckdb::top_level_duckdb_ddl_type = pgduckdb::DDLType::NONE;
+
 	SPI_connect();
 
 	/*
@@ -889,20 +939,20 @@ DECLARE_PG_FUNCTION(duckdb_alter_table_trigger) {
 		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'
+		WHERE cmds.object_type in ('table', 'table column')
 		AND pg_class.relam = (SELECT oid FROM pg_am WHERE amname = 'duckdb')
 		UNION ALL
 		SELECT objid as relid, false AS needs_to_check_temporary_set
 		FROM pg_catalog.pg_event_trigger_ddl_commands() cmds
 		JOIN duckdb.tables AS ddbtables
 		ON cmds.objid = ddbtables.relid
-		WHERE cmds.object_type = 'table'
+		WHERE cmds.object_type in ('table', 'table column')
 		UNION ALL
 		SELECT objid as relid, true AS needs_to_check_temporary_set
 		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'
+		WHERE cmds.object_type in ('table', 'table column')
 		AND pg_class.relam != (SELECT oid FROM pg_am WHERE amname = 'duckdb')
 		AND pg_class.relpersistence = 't'
 		)",
@@ -946,7 +996,29 @@ DECLARE_PG_FUNCTION(duckdb_alter_table_trigger) {
 		}
 	}
 
-	elog(ERROR, "DuckDB does not support ALTER TABLE yet");
+	/* Forcibly allow whatever writes Postgres did for this command */
+	pgduckdb::ClaimCurrentCommandId(true);
+
+	/* We're going to run multiple queries in DuckDB, so we need to start a
+	 * transaction to ensure ACID guarantees hold. */
+	auto connection = pgduckdb::DuckDBManager::GetConnection(true);
+
+	EventTriggerData *trigdata = (EventTriggerData *)fcinfo->context;
+	char *alter_table_stmt_string;
+	if (IsA(trigdata->parsetree, AlterTableStmt)) {
+		AlterTableStmt *alter_table_stmt = (AlterTableStmt *)trigdata->parsetree;
+		alter_table_stmt_string = pgduckdb_get_alter_tabledef(relid, alter_table_stmt);
+	} else if (IsA(trigdata->parsetree, RenameStmt)) {
+		RenameStmt *rename_stmt = (RenameStmt *)trigdata->parsetree;
+		alter_table_stmt_string = pgduckdb_get_rename_tabledef(relid, rename_stmt);
+	} else {
+		elog(ERROR, "Unexpected parsetree type: %d", nodeTag(trigdata->parsetree));
+	}
+
+	elog(DEBUG1, "Executing: %s", alter_table_stmt_string);
+	auto res = pgduckdb::DuckDBQueryOrThrow(*connection, alter_table_stmt_string);
+
+	PG_RETURN_NULL();
 }
 
 /*