Writable Iceberg tables with REST catalog support (#68)

In this PR, we add support for writable REST catalog tables. The API may still change, so I refrain documenting it in the PR description. Once we finalise the APIs, the best way would be to put under https://github.com/Snowflake-Labs/pg_lake/blob/main/docs/iceberg-tables.md, so stay tuned for that.

In the earlier set of PRs (#47, #49, #51, #52 and #56) we prototyped adding support for writable rest catalog in different stages. However, it turns out that a single PR is better to tackle these very much related commits. It seemed overkill to maintain these set of PRs.

This new table type shares almost the same architecture with iceberg tables `catalog=postgres`.

#### Similarities 
- Tables are tracked the `pg_lake` catalogs such as `lake_table.files`, `lake_table.data_file_column_stats` and `lake_iceberg.tables_internal`.
- All metadata handling follows `ApplyIcebergMetadataChanges()` logic. Instead of generating a new `metadata.json` as we do for `catalog=postgres`, for these tables we collect the changes happened in the transaction, and apply to the REST catalog right after it is committed in Postgres.

#### Differences
- The `metadata_location` column in `lake_iceberg.tables_internal` is always `NULL` 
- Does not support RENAME TABLE / SET SCHEMA etc.


#### Some other notes on the implementation & design:
- We first `COMMIT` in Postgres, then in `post-commit` hook, send a `POST` request to REST catalog. So, it is possible that the changes are committed in Postgres, but not in REST catalog. This is a known limitation, and we'll have follow-up PRs to make sure we can recover from this situation.  
- Creating a table and modifying it in the same Postgres transaction cannot be committed atomically in REST catalog. There is no such API in REST catalog. So, there are some additional error scenarios where table creation committed in REST catalog, say not the full CTAS. This is an unfortunate limitation that we inherit from REST catalog APIs.
- Our implementation currently assumes that the Postgres is the single-writer to this table in the REST catalog. So, a concurrent modification breaks the table from Postgres side. For now, this is the current state. We plan to improve it in the future.



#### TODO:
- [x] `DROP partition_by` is not working (fixed by #79)
- [x] Concurrency
- [x] Certain DDLs do not work (e.g., ADD COLUMN with defaults), prevent much earlier
- [x] VACUUM regression tests
- [x] VACUUM failures (e.g., do we clean up properly?)
- [x] VACUUM (ICEBERG)
- [x] auto-vacuum test
- [x] Truncate test
- [x] savepoint
- [x] Complex TX test
- [x] Column names with quotes
- [x] Add column + add partition by + drop column in the same tx
- [x] Tests for read from postgres / iceberg, modify REST (or the other way around)
- [x] Zero column table?
- [x] DROP TABLE implemented, but needs tests (e.g., create - drop in the same tx, drop table removes the metadata from rest catalog etc).
- [x] `SET partition_by` to an already existing partition by is not supported in Polaris. We should skip sending such requests, instead only send `set partition_spec` alone. (fixed by #79)
- [ ] Recovery after failures (e.g., re-sync the previous snapshot/DDL) [Follow-up PR needed]
- [x] Cache access token, currently we fetch on every REST request interaction [Follow-up PR needed]
- [x] Cancel query
- [x] sequences / serial / generated columns etc.
- [x] All data types
- [ ] Docs [Follow-up PR needed]

Author: sfc-gh-okalaci

pg_lake_iceberg/include/pg_lake/iceberg/api/table_metadata.h

@@ -35,7 +35,7 @@ extern PGDLLEXPORT char *GenerateRemoteMetadataFilePath(int version, const char
 extern PGDLLEXPORT void UploadTableMetadataToURI(IcebergTableMetadata * tableMetadata, char *metadataURI);
 extern PGDLLEXPORT void AdjustAndRetainMetadataLogs(IcebergTableMetadata * metadata, char *prevMetadataPath, size_t snapshotLogLength, int64_t prev_last_updated_ms);
 extern PGDLLEXPORT void UpdateLatestSnapshot(IcebergTableMetadata * tableMetadata, IcebergSnapshot * newSnapshot);
-extern PGDLLEXPORT bool RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata, bool isVerbose);
+extern PGDLLEXPORT List *RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata, bool isVerbose);
 extern PGDLLEXPORT void GenerateSnapshotLogEntries(IcebergTableMetadata * metadata);
 extern PGDLLEXPORT int FindLargestPartitionFieldId(IcebergPartitionSpec * newSpec);
 extern PGDLLEXPORT void AppendCurrentPostgresSchema(Oid relationId, IcebergTableMetadata * metadata,

pg_lake_iceberg/include/pg_lake/iceberg/catalog.h

@@ -55,6 +55,8 @@ extern PGDLLEXPORT void DeleteInternalIcebergCatalogTable(Oid relationId);
 
 extern PGDLLEXPORT List *GetAllInternalIcebergRelationIds(void);
 extern PGDLLEXPORT char *GetIcebergMetadataLocation(Oid relationId, bool forUpdate);
+extern PGDLLEXPORT void LockIcebergPgLakeCatalogForUpdate(Oid relationId);
+extern PGDLLEXPORT char *GetIcebergCatalogMetadataLocation(Oid relationId, bool forUpdate);
 extern PGDLLEXPORT char *GetIcebergCatalogPreviousMetadataLocation(Oid relationId, bool forUpdate);
 extern PGDLLEXPORT void UpdateExternalCatalogMetadataLocation(char *catalogName, char *schemaName, char *tableName, const char *metadataLocation,
 															  const char *previousMetadataLocation);

pg_lake_iceberg/include/pg_lake/iceberg/metadata_operations.h

@@ -19,6 +19,6 @@
 
 #include "nodes/pg_list.h"
 
-extern PGDLLEXPORT void ApplyIcebergMetadataChanges(Oid relationId, List *metadataOperations, List *allTransforms, bool isVerbose);
+extern PGDLLEXPORT List *ApplyIcebergMetadataChanges(Oid relationId, List *metadataOperations, List *allTransforms, bool isVerbose);
 extern PGDLLEXPORT bool ShouldSkipMetadataChangeToIceberg(List *metadataOperationTypes);
 extern PGDLLEXPORT List *GetMetadataOperationTypes(List *metadataOperations);

pg_lake_iceberg/include/pg_lake/iceberg/metadata_spec.h

@@ -290,3 +290,5 @@ typedef struct IcebergTableMetadata
 
 extern PGDLLEXPORT IcebergTableMetadata * ReadIcebergTableMetadata(const char *tableMetadataPath);
 extern PGDLLEXPORT char *WriteIcebergTableMetadataToJson(IcebergTableMetadata * metadata);
+extern PGDLLEXPORT void AppendIcebergTableSchemaForRestCatalog(StringInfo command, IcebergTableSchema * schemas, size_t schemas_length);
+extern PGDLLEXPORT void AppendIcebergPartitionSpecFields(StringInfo command, IcebergPartitionSpecField * fields, size_t fields_length);

pg_lake_iceberg/include/pg_lake/rest_catalog/rest_catalog.h

@@ -20,8 +20,10 @@
 #include "postgres.h"
 #include "pg_lake/http/http_client.h"
 #include "pg_lake/util/rel_utils.h"
+#include "pg_lake/parquet/field.h"
+#include "pg_lake/iceberg/api/snapshot.h"
 
-extern char *RestCatalogHost;
+extern PGDLLEXPORT char *RestCatalogHost;
 extern char *RestCatalogClientId;
 extern char *RestCatalogClientSecret;
 
@@ -30,11 +32,43 @@ extern char *RestCatalogClientSecret;
 #define REST_CATALOG_NAMESPACE_NAME "%s/api/catalog/v1/%s/namespaces/%s"
 #define REST_CATALOG_NAMESPACE "%s/api/catalog/v1/%s/namespaces"
 
+#define REST_CATALOG_TABLE "%s/api/catalog/v1/%s/namespaces/%s/tables/%s"
+#define REST_CATALOG_TABLES "%s/api/catalog/v1/%s/namespaces/%s/tables"
+
 #define REST_CATALOG_AUTH_TOKEN_PATH "%s/api/catalog/v1/oauth/tokens"
-#define GET_REST_CATALOG_METADATA_LOCATION "%s/api/catalog/v1/%s/namespaces/%s/tables/%s"
 
-extern PGDLLEXPORT char *RestCatalogFetchAccessToken(void);
+#define REST_CATALOG_TRANSACTION_COMMIT "%s/api/catalog/v1/%s/transactions/commit"
+
+typedef enum RestCatalogOperationType
+{
+	REST_CATALOG_CREATE_TABLE = 0,
+	REST_CATALOG_ADD_SNAPSHOT = 1,
+	REST_CATALOG_ADD_SCHEMA = 2,
+	REST_CATALOG_SET_CURRENT_SCHEMA = 3,
+	REST_CATALOG_ADD_PARTITION = 4,
+	REST_CATALOG_REMOVE_SNAPSHOT = 5,
+	REST_CATALOG_DROP_TABLE = 6,
+	REST_CATALOG_SET_DEFAULT_PARTITION_ID = 7,
+}			RestCatalogOperationType;
+
+
+typedef struct RestCatalogRequest
+{
+	Oid			relationId;
+	RestCatalogOperationType operationType;
+
+	/*
+	 * For each request, holds the "action" part of the request body. We
+	 * concatenate all requests from multiple tables into a single transaction
+	 * commit request. The only exception is CREATE/DROP table, where body
+	 * holds the full request body.
+	 */
+	char	   *body;
+}			RestCatalogRequest;
+
 extern PGDLLEXPORT void RegisterNamespaceToRestCatalog(const char *catalogName, const char *namespaceName);
+extern PGDLLEXPORT void StartStageRestCatalogIcebergTableCreate(Oid relationId);
+extern PGDLLEXPORT char *FinishStageRestCatalogIcebergTableCreateRestRequest(Oid relationId, DataFileSchema * dataFileSchema, List *partitionSpecs);
 extern PGDLLEXPORT void ErrorIfRestNamespaceDoesNotExist(const char *catalogName, const char *namespaceName);
 extern PGDLLEXPORT char *GetRestCatalogName(Oid relationId);
 extern PGDLLEXPORT char *GetRestCatalogNamespace(Oid relationId);
@@ -43,3 +77,12 @@ extern PGDLLEXPORT bool IsReadOnlyRestCatalogIcebergTable(Oid relationId);
 extern PGDLLEXPORT char *GetMetadataLocationFromRestCatalog(const char *restCatalogName, const char *namespaceName,
 															const char *relationName);
 extern PGDLLEXPORT char *GetMetadataLocationForRestCatalogForIcebergTable(Oid relationId);
+extern PGDLLEXPORT void ReportHTTPError(HttpResult httpResult, int level);
+extern PGDLLEXPORT List *PostHeadersWithAuth(void);
+extern PGDLLEXPORT List *DeleteHeadersWithAuth(void);
+extern PGDLLEXPORT RestCatalogRequest * GetAddSnapshotCatalogRequest(IcebergSnapshot * newSnapshot, Oid relationId);
+extern PGDLLEXPORT RestCatalogRequest * GetAddSchemaCatalogRequest(Oid relationId, DataFileSchema * dataFileSchema);
+extern PGDLLEXPORT RestCatalogRequest * GetSetCurrentSchemaCatalogRequest(Oid relationId, int32_t schemaId);
+extern PGDLLEXPORT RestCatalogRequest * GetAddPartitionCatalogRequest(Oid relationId, List *partitionSpec);
+extern PGDLLEXPORT RestCatalogRequest * GetSetPartitionDefaultIdCatalogRequest(Oid relationId, int specId);
+extern PGDLLEXPORT RestCatalogRequest * GetRemoveSnapshotCatalogRequest(List *removedSnapshotIds, Oid relationId);

pg_lake_iceberg/src/iceberg/api/table_metadata.c

@@ -215,13 +215,13 @@ AddIcebergSnapshotToMetadata(IcebergTableMetadata * metadata, IcebergSnapshot *
  * It only does the in-memory operation, the caller is responsible for
  * persisting the changes.
  */
-bool
+List *
 RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata, bool isVerbose)
 {
 	if (metadata->snapshots_length == 0)
 	{
 		/* no snapshots yet, not possible to trigger, but let's be defensive */
-		return false;
+		return NIL;
 	}
 
 	/*
@@ -242,7 +242,7 @@ RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata,
 
 	if (expiredSnapshotCount == 0)
 		/* no snapshots to expire */
-		return false;
+		return NIL;
 
 	/* we might expire all snapshots, always retain at least 1 snapshot */
 	if (nonExpiredSnapshotCount == 0)
@@ -261,8 +261,15 @@ RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata,
 		nonExpiredSnapshotCount = 1;
 	}
 
+	List	   *expiredSnapshotIds = NIL;
+
 	for (int snapshotIndex = 0; snapshotIndex < expiredSnapshotCount; snapshotIndex++)
 	{
+		int64_t    *expiredSnapshotIdPtr = palloc(sizeof(int64_t));
+
+		*expiredSnapshotIdPtr = expiredSnapshots[snapshotIndex].snapshot_id;
+		expiredSnapshotIds = lappend(expiredSnapshotIds, expiredSnapshotIdPtr);
+
 		ereport(isVerbose ? INFO : DEBUG1,
 				(errmsg("expiring snapshot %" PRId64 " from %s",
 						expiredSnapshots[snapshotIndex].snapshot_id,
@@ -271,10 +278,16 @@ RemoveOldSnapshotsFromMetadata(Oid relationId, IcebergTableMetadata * metadata,
 
 	DeleteUnreferencedFiles(relationId, metadata, expiredSnapshots, expiredSnapshotCount, nonExpiredSnapshots, nonExpiredSnapshotCount);
 
-	metadata->snapshots = nonExpiredSnapshots;
-	metadata->snapshots_length = nonExpiredSnapshotCount;
+	IcebergCatalogType catalogType = GetIcebergCatalogType(relationId);
+	bool		writableRestCatalogTable = catalogType == REST_CATALOG_READ_WRITE;
+
+	if (!writableRestCatalogTable)
+	{
+		metadata->snapshots = nonExpiredSnapshots;
+		metadata->snapshots_length = nonExpiredSnapshotCount;
+	}
 
-	return true;
+	return expiredSnapshotIds;
 }
 
 

pg_lake_iceberg/src/iceberg/catalog.c

@@ -33,7 +33,6 @@
 
 char	   *IcebergDefaultLocationPrefix = NULL;
 
-static char *GetIcebergCatalogMetadataLocation(Oid relationId, bool forUpdate);
 static char *GetIcebergExternalMetadataLocation(Oid relationId);
 static char *GetIcebergCatalogMetadataLocationInternal(Oid relationId, bool isPrevMetadata, bool forUpdate);
 static char *GetIcebergCatalogColumnInternal(Oid relationId, char *columnName, bool forUpdate, bool errorIfNotFound);
@@ -66,7 +65,7 @@ InsertInternalIcebergCatalogTable(Oid relationId, const char *metadataLocation,
 
 	DECLARE_SPI_ARGS(3);
 	SPI_ARG_VALUE(1, OIDOID, relationId, false);
-	SPI_ARG_VALUE(2, TEXTOID, metadataLocation, false);
+	SPI_ARG_VALUE(2, TEXTOID, metadataLocation, metadataLocation == NULL);
 	SPI_ARG_VALUE(3, BOOLOID, hasCustomLocation, false);
 
 	SPI_START();
@@ -358,19 +357,41 @@ GetIcebergMetadataLocation(Oid relationId, bool forUpdate)
 {
 	Assert(IsIcebergTable(relationId));
 
-	if (IsInternalIcebergTable(relationId))
-	{
-		return GetIcebergCatalogMetadataLocation(relationId, forUpdate);
-	}
-	else
+	IcebergCatalogType catalogType = GetIcebergCatalogType(relationId);
+
+	if (IsExternalIcebergTable(relationId) || catalogType == REST_CATALOG_READ_WRITE)
 	{
-		if (forUpdate)
+		/*
+		 * We get the metadata location from the iceberg catalog table for
+		 * internal iceberg tables, except for REST_CATALOG_READ_WRITE.
+		 *
+		 * Writable rest catalog iceberg tables are a bit different. They are
+		 * internal iceberg tables, all their metadata is stored in our
+		 * catalogs, such as data files, column stats, etc. However, their
+		 * metadata location is stored in the rest catalog itself. So, we need
+		 * to fetch the metadata location from the rest catalog for those
+		 * tables, see below.
+		 */
+		if (forUpdate && catalogType != REST_CATALOG_READ_WRITE)
 			ereport(ERROR,
 					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 					 errmsg("Updating iceberg metadata is not supported for external Iceberg tables")));
+		else if (forUpdate && catalogType == REST_CATALOG_READ_WRITE)
+		{
+			/*
+			 * But we still need to acquire locks on the iceberg catalog table
+			 * row for internal iceberg tables to prevent concurrent updates.
+			 * So, we call GetIcebergCatalogMetadataLocation().
+			 */
+			LockIcebergPgLakeCatalogForUpdate(relationId);
+		}
 
 		return GetIcebergExternalMetadataLocation(relationId);
 	}
+	else
+	{
+		return GetIcebergCatalogMetadataLocation(relationId, forUpdate);
+	}
 }
 
 
@@ -381,7 +402,7 @@ GetIcebergMetadataLocation(Oid relationId, bool forUpdate)
 * If the metadata row for the table is going to be updated, the caller should
 * pass forUpdate as true.
 */
-static char *
+char *
 GetIcebergCatalogMetadataLocation(Oid relationId, bool forUpdate)
 {
 	Assert(IsInternalIcebergTable(relationId));
@@ -390,6 +411,32 @@ GetIcebergCatalogMetadataLocation(Oid relationId, bool forUpdate)
 }
 
 
+/*
+* LockIcebergPgLakeCatalogForUpdate acquires necessary locks on the
+* iceberg catalog table row for the given relation to serialize concurrent
+* updates. Note that normally GetIcebergCatalogMetadataLocation() already does
+* the trick for us. This function is only needed in cases where we need to
+* acquire the locks but don't have the metadata location itself, such as
+* writable rest catalog iceberg tables.
+*/
+void
+LockIcebergPgLakeCatalogForUpdate(Oid relationId)
+{
+	/*
+	 * We don't actually need the metadata location here. We just need to
+	 * acquire the necessary locks on the iceberg catalog table row to
+	 * serialize concurrent updates.
+	 *
+	 * Note that we cannot skip reading the actual metadata location and
+	 * returning it here as other callers might depend on the returned value.
+	 * So, we call GetIcebergCatalogMetadataLocation() as usual.
+	 */
+	bool		forUpdate = true;
+
+	GetIcebergCatalogMetadataLocationInternal(relationId, false, forUpdate);
+}
+
+
 /*
  * GetIcebergExternalMetadataLocation returns the metadata location for an external iceberg table.
  */
@@ -402,7 +449,7 @@ GetIcebergExternalMetadataLocation(Oid relationId)
 
 	char	   *currentMetadataPath = NULL;
 
-	if (icebergCatalogType == REST_CATALOG_READ_ONLY)
+	if (icebergCatalogType == REST_CATALOG_READ_ONLY || icebergCatalogType == REST_CATALOG_READ_WRITE)
 	{
 		currentMetadataPath = GetMetadataLocationForRestCatalogForIcebergTable(relationId);
 	}
@@ -549,13 +596,13 @@ bool
 RelationExistsInTheIcebergCatalog(Oid relationId)
 {
 	bool		forUpdate = false;
-	char	   *columnName = "metadata_location";
+	char	   *columnName = "table_name";
 	bool		errorIfNotFound = false;
 
-	char	   *metadataLocation =
+	char	   *tableName =
 		GetIcebergCatalogColumnInternal(relationId, columnName, forUpdate, errorIfNotFound);
 
-	return metadataLocation != NULL;
+	return tableName != NULL;
 }