Some cleanup from last review addressing since I left everything in Cursor's hands

Signed-off-by: sfc-gh-npuka <naisila.puka@snowflake.com>

Author: sfc-gh-npuka

pg_lake_engine/src/utils/catalog_type.c

@@ -110,8 +110,7 @@ HasReadOnlyOption(List *options)
 /*
  * IsCatalogOwnedByExtension returns true if the catalog name is one of
  * the reserved built-in names: 'rest', 'object_store', or 'postgres'.
- * Comparison is case-insensitive so that "Postgres", "REST", etc. are
- * also recognized as reserved.
+ * Comparison is case-insensitive.
  */
 bool
 IsCatalogOwnedByExtension(const char *catalog)

pg_lake_iceberg/include/pg_lake/rest_catalog/rest_catalog.h

@@ -119,7 +119,7 @@ extern PGDLLEXPORT char *GetMetadataLocationForRestCatalogForIcebergTable(Oid re
 extern PGDLLEXPORT void ReportHTTPError(HttpResult httpResult, int level);
 extern PGDLLEXPORT List *PostHeadersWithAuth(RestCatalogOptions * opts);
 extern PGDLLEXPORT List *DeleteHeadersWithAuth(RestCatalogOptions * opts);
-extern PGDLLEXPORT HttpResult SendRequestToRestCatalog(HttpMethod method, const char *url, const char *body, List *headers, RestCatalogOptions * opts);
+extern PGDLLEXPORT HttpResult SendRequestToRestCatalog(RestCatalogOptions * opts, HttpMethod method, const char *url, const char *body, List *headers);
 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);

pg_lake_iceberg/src/rest_catalog/rest_catalog.c

@@ -68,7 +68,8 @@ int			RestCatalogAuthType = REST_CATALOG_AUTH_TYPE_OAUTH2;
 bool		RestCatalogEnableVendedCredentials = true;
 
 /*
- * Per-catalog token cache. Keyed by catalog.
+ * Per-rest-catalog token cache. Keyed by catalog.
+ * Should always be accessed via GetRestCatalogAccessToken()
  */
 #define TOKEN_CACHE_KEY_LEN NAMEDATALEN
 
@@ -82,6 +83,8 @@ typedef struct RestCatalogTokenCacheEntry
 static HTAB *RestCatalogTokenCache = NULL;
 static MemoryContext RestTokenCacheCtx = NULL;
 
+/* end of per-catalog token cache variables */
+
 static char *GetRestCatalogAccessToken(RestCatalogOptions * opts, bool forceRefreshToken);
 static void FetchRestCatalogAccessToken(RestCatalogOptions * opts, char **accessToken, int *expiresIn);
 static void CreateNamespaceOnRestCatalog(RestCatalogOptions * opts, const char *catalogName, const char *namespaceName);
@@ -723,8 +726,8 @@ StartStageRestCatalogIcebergTableCreate(Oid relationId)
 		headers = lappend(headers, vendedCreds);
 	}
 
-	HttpResult	httpResult = SendRequestToRestCatalog(HTTP_POST, postUrl, body->data,
-													  headers, opts);
+	HttpResult	httpResult = SendRequestToRestCatalog(opts, HTTP_POST, postUrl, body->data,
+													  headers);
 
 	if (httpResult.status != 200)
 	{
@@ -859,9 +862,8 @@ RegisterNamespaceToRestCatalog(RestCatalogOptions * opts, const char *catalogNam
 		psprintf(REST_CATALOG_NAMESPACE_NAME,
 				 opts->host, URLEncodePath(catalogName),
 				 URLEncodePath(namespaceName));
-	HttpResult	httpResult = SendRequestToRestCatalog(HTTP_GET, getUrl, NULL,
-													  GetHeadersWithAuth(opts),
-													  opts);
+	HttpResult	httpResult = SendRequestToRestCatalog(opts, HTTP_GET, getUrl, NULL,
+													  GetHeadersWithAuth(opts));
 
 	switch (httpResult.status)
 	{
@@ -951,9 +953,8 @@ ErrorIfRestNamespaceDoesNotExist(RestCatalogOptions * opts, const char *catalogN
 		psprintf(REST_CATALOG_NAMESPACE_NAME,
 				 opts->host, URLEncodePath(catalogName),
 				 URLEncodePath(namespaceName));
-	HttpResult	httpResult = SendRequestToRestCatalog(HTTP_GET, getUrl, NULL,
-													  GetHeadersWithAuth(opts),
-													  opts);
+	HttpResult	httpResult = SendRequestToRestCatalog(opts, HTTP_GET, getUrl, NULL,
+													  GetHeadersWithAuth(opts));
 
 	/* namespace not found */
 	if (httpResult.status == 404)
@@ -1002,8 +1003,7 @@ GetMetadataLocationFromRestCatalog(RestCatalogOptions * opts, const char *restCa
 				 opts->host, URLEncodePath(restCatalogName), URLEncodePath(namespaceName), URLEncodePath(relationName));
 
 	List	   *headers = GetHeadersWithAuth(opts);
-	HttpResult	hr = SendRequestToRestCatalog(HTTP_GET, getUrl, NULL, headers,
-											  opts);
+	HttpResult	hr = SendRequestToRestCatalog(opts, HTTP_GET, getUrl, NULL, headers);
 
 	if (hr.status != 200)
 	{
@@ -1051,9 +1051,8 @@ CreateNamespaceOnRestCatalog(RestCatalogOptions * opts, const char *catalogName,
 		psprintf(REST_CATALOG_NAMESPACE, opts->host,
 				 URLEncodePath(catalogName));
 
-	HttpResult	httpResult = SendRequestToRestCatalog(HTTP_POST, postUrl, body.data,
-													  PostHeadersWithAuth(opts),
-													  opts);
+	HttpResult	httpResult = SendRequestToRestCatalog(opts, HTTP_POST, postUrl, body.data,
+													  PostHeadersWithAuth(opts));
 
 	if (httpResult.status != 200)
 	{
@@ -1155,6 +1154,12 @@ BuildTokenCacheKey(char *key, const RestCatalogOptions * opts)
  * Any ALTER/DROP SERVER blows away the entire token cache so stale
  * credentials are never reused.  The cache is rebuilt lazily on the
  * next token lookup.
+ *
+ * We ignore hashvalue and reset the whole cache rather than selectively
+ * invalidating a single server's entry (as postgres_fdw does).  With a
+ * handful of servers and infrequent ALTER SERVER, the cost of a few
+ * extra OAuth round-trips is negligible compared to the complexity of
+ * keying the cache by OID for selective invalidation.
  */
 static void
 InvalidateRestTokenCache(Datum arg, int cacheid, uint32 hashvalue)
@@ -1169,6 +1174,11 @@ InvalidateRestTokenCache(Datum arg, int cacheid, uint32 hashvalue)
 
 /*
  * Initialize the per-catalog token cache hash table if needed.
+ *
+ * TokenCacheCallbackRegistered is separate from RestCatalogTokenCache because
+ * the callback must be registered exactly once per backend lifetime
+ * (CacheRegisterSyscacheCallback appends to a fixed-size array), while
+ * RestCatalogTokenCache is reset to NULL on every invalidation.
  */
 static bool TokenCacheCallbackRegistered = false;
 
@@ -1318,10 +1328,14 @@ FetchRestCatalogAccessToken(RestCatalogOptions * opts, char **accessToken, int *
 
 	headers = lappend(headers, "Content-Type: application/x-www-form-urlencoded");
 
-	/* POST — pass NULL opts to skip 419 token refresh (avoids recursion) */
-	HttpResult	httpResponse = SendRequestToRestCatalog(HTTP_POST, accessTokenUrl,
-														body.data, headers,
-														NULL);
+	/*
+	 * Pass NULL opts so SendRequestToRestCatalog skips the 419 token-refresh
+	 * retry branch.  Otherwise a 419 here would call
+	 * GetRestCatalogAccessToken -> FetchRestCatalogAccessToken ->
+	 * SendRequestToRestCatalog in an infinite loop.
+	 */
+	HttpResult	httpResponse = SendRequestToRestCatalog(NULL, HTTP_POST, accessTokenUrl,
+														body.data, headers);
 
 	if (httpResponse.status != 200)
 		ereport(ERROR,

pg_lake_table/src/transaction/track_iceberg_metadata_changes.c

@@ -292,10 +292,9 @@ PostAllRestCatalogRequests(void)
 			if (createTableRequest != NULL)
 			{
 				HttpResult	httpResult =
-					SendRequestToRestCatalog(HTTP_POST, requestPerTable->tableRestUrl,
-											 createTableRequest->body,
-											 PostHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts),
-											 PgLakeXactRestCatalog->catalogOpts);
+					SendRequestToRestCatalog(PgLakeXactRestCatalog->catalogOpts, HTTP_POST,
+											 requestPerTable->tableRestUrl, createTableRequest->body,
+											 PostHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts));
 
 				if (httpResult.status != 200)
 				{
@@ -310,10 +309,9 @@ PostAllRestCatalogRequests(void)
 			else if (dropTableRequest != NULL)
 			{
 				HttpResult	httpResult =
-					SendRequestToRestCatalog(HTTP_DELETE, requestPerTable->tableRestUrl,
-											 NULL,
-											 DeleteHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts),
-											 PgLakeXactRestCatalog->catalogOpts);
+					SendRequestToRestCatalog(PgLakeXactRestCatalog->catalogOpts, HTTP_DELETE,
+											 requestPerTable->tableRestUrl, NULL,
+											 DeleteHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts));
 
 				if (httpResult.status != 204)
 				{
@@ -433,9 +431,9 @@ PostAllRestCatalogRequests(void)
 
 		char	   *url = psprintf(REST_CATALOG_TRANSACTION_COMMIT,
 								   PgLakeXactRestCatalog->catalogOpts->host, catalogName);
-		HttpResult	httpResult = SendRequestToRestCatalog(HTTP_POST, url, batchRequestBody->data,
-														  PostHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts),
-														  PgLakeXactRestCatalog->catalogOpts);
+		HttpResult	httpResult = SendRequestToRestCatalog(PgLakeXactRestCatalog->catalogOpts, HTTP_POST,
+														  url, batchRequestBody->data,
+														  PostHeadersWithAuth(PgLakeXactRestCatalog->catalogOpts));
 
 		if (httpResult.status != 204)
 		{

pg_lake_table/tests/pytests/test_modify_iceberg_rest_table.py

@@ -923,7 +923,7 @@ def test_token_cache_reuses_token_across_catalog_ops(
 
     Uses pg_lake_iceberg.http_client_trace_traffic to observe actual
     HTTP traffic: each token fetch shows up as a POST to .../oauth/tokens
-    in the connection notices.
+    in the connection notices. Does 3 back-to-back inserts.
     """
     if installcheck:
         return
@@ -935,18 +935,17 @@ def test_token_cache_reuses_token_across_catalog_ops(
     pg_conn.commit()
 
     run_command(
-        f"CREATE TABLE {SCHEMA_NAME}.{TABLE_NAME} (id bigint, value text) "
-        f"USING iceberg WITH (catalog='rest')",
+        "SET pg_lake_iceberg.http_client_trace_traffic TO on",
         pg_conn,
     )
-    pg_conn.commit()
+    pg_conn.notices.clear()
 
     run_command(
-        "SET pg_lake_iceberg.http_client_trace_traffic TO on",
+        f"CREATE TABLE {SCHEMA_NAME}.{TABLE_NAME} (id bigint, value text) "
+        f"USING iceberg WITH (catalog='rest')",
         pg_conn,
     )
-
-    pg_conn.notices.clear()
+    pg_conn.commit()
 
     for i in range(3):
         run_command(
@@ -958,8 +957,8 @@ def test_token_cache_reuses_token_across_catalog_ops(
     token_fetches = sum(
         1 for n in pg_conn.notices if "oauth/tokens" in n and "POST" in n
     )
-    assert token_fetches <= 1, (
-        f"Expected at most 1 OAuth token fetch (cached), got {token_fetches}. "
+    assert token_fetches == 1, (
+        f"Expected exactly 1 OAuth token fetch (cached), got {token_fetches}. "
         f"Notices:\n" + "\n".join(pg_conn.notices)
     )
 
@@ -987,6 +986,8 @@ def test_alter_server_credentials_invalidates_token_cache(
     next catalog operation re-fetches it.  We verify this by enabling
     HTTP traffic tracing and checking that a POST to .../oauth/tokens
     appears after the ALTER SERVER (proving the cache was invalidated).
+    Test that cache is invalidated on bogus credentials (1 fetch, commit fails),
+    then cache is invalidated again on restored credentials (1 fetch, commit succeeds).
     """
     if installcheck:
         return
@@ -1078,8 +1079,8 @@ def test_alter_server_credentials_invalidates_token_cache(
         "Expected COMMIT to fail after ALTER SERVER set bogus client_id "
         "(cache should have been invalidated, forcing re-auth with bad creds)"
     )
-    assert post_alter_fetches >= 1, (
-        f"Expected token re-fetch after ALTER SERVER (cache invalidated), "
+    assert post_alter_fetches == 1, (
+        f"Expected exactly 1 token re-fetch after ALTER SERVER (cache invalidated), "
         f"got {post_alter_fetches}. Notices ({len(post_alter_notices)}):\n"
         + "\n".join(post_alter_notices)
     )
@@ -1090,12 +1091,22 @@ def test_alter_server_credentials_invalidates_token_cache(
     )
     superuser_conn.commit()
 
+    superuser_conn.notices.clear()
+
     run_command(
         f"INSERT INTO {SCHEMA_NAME}.{TABLE_NAME} VALUES (4)",
         superuser_conn,
     )
     superuser_conn.commit()
 
+    restore_fetches = sum(
+        1 for n in superuser_conn.notices if "oauth/tokens" in n and "POST" in n
+    )
+    assert restore_fetches == 1, (
+        f"Expected exactly 1 token re-fetch after restoring credentials, "
+        f"got {restore_fetches}. Notices:\n" + "\n".join(superuser_conn.notices)
+    )
+
     results = run_query(
         f"SELECT count(*) FROM {SCHEMA_NAME}.{TABLE_NAME}", superuser_conn
     )