/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * src/metadata/collection.c
 *
 * Implementation of collection metadata cache.
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"
#include "fmgr.h"
#include "miscadmin.h"

#include "access/xact.h"
#include "catalog/pg_attribute.h"
#include "catalog/pg_class.h"
#include "catalog/namespace.h"
#include "commands/extension.h"
#include "executor/spi.h"
#include "lib/stringinfo.h"
#include "nodes/makefuncs.h"
#include "nodes/pg_list.h"
#include "parser/parse_func.h"
#include "storage/lmgr.h"
#include "funcapi.h"
#include "utils/builtins.h"
#include "utils/catcache.h"
#include "utils/guc.h"
#include "utils/inval.h"
#include "utils/lsyscache.h"
#include "utils/memutils.h"
#include "utils/syscache.h"
#include "utils/version_utils.h"

#include "metadata/collection.h"
#include "metadata/metadata_cache.h"
#include "utils/documentdb_errors.h"
#include "metadata/relation_utils.h"
#include "utils/query_utils.h"
#include "utils/guc_utils.h"
#include "metadata/metadata_guc.h"
#include "api_hooks.h"
#include "commands/parse_error.h"
#include "utils/feature_counter.h"
#include "jsonschema/bson_json_schema_tree.h"

#define CREATE_COLLECTION_FUNC_NARGS 2


/*
 * NameToCollectionCacheEntry is an entry in NameToCollectionHash that maps
 * a qualified collection name to a collection.
 */
typedef struct NameToCollectionCacheEntry
{
	/* Mongo qualified name of the collection */
	MongoCollectionName name;

	/* collection metadata */
	MongoCollection collection;

	/* set to false when an invalidation for the relation is received */
	bool isValid;
} NameToCollectionCacheEntry;

/*
 * RelationIdToCollectionCacheEntry is an entry in RelationIdToCollectionHash
 * that maps a relation OID to a collection.
 */
typedef struct RelationIdToCollectionCacheEntry
{
	/* OID of the Postgres relation underlying the collection */
	Oid relationId;

	/* collection metadata */
	MongoCollection collection;

	/* set to false when an invalidation for the relation is received */
	bool isValid;
} RelationIdToCollectionCacheEntry;


static const char CharactersNotAllowedInDatabaseNames[7] = {
	'/', '\\', '.', ' ', '"', '$', '\0'
};
static const int CharactersNotAllowedInDatabaseNamesLength =
	sizeof(CharactersNotAllowedInDatabaseNames);

static const char CharactersNotAllowedInCollectionNames[2] = { '$', '\0' };
static const int CharactersNotAllowedInCollectionNamesLength =
	sizeof(CharactersNotAllowedInCollectionNames);

static const char *ValidSystemCollectionNames[5] = {
	"system.users", "system.js", "system.views", "system.profile", "system.dbSentinel"
};
static const int ValidSystemCollectionNamesLength = 5;

/* Not allowing any writes to the below system namespaces */
static const char *NonWritableSystemCollectionNames[4] = {
	"system.users", "system.js", "system.views", "system.profile"
};
static const int NonWritableSystemCollectionNamesLength = 4;

static const uint32_t MaxDatabaseCollectionLength = 235;
static const StringView SystemPrefix = { .length = 7, .string = "system." };

extern bool UseLocalExecutionShardQueries;
extern bool ForceLocalExecutionShardQueries;
extern bool EnableSchemaValidation;
extern int MaxSchemaValidatorSize;
extern bool EnableOnlyCollectionCacheInvalidateOnCollectionChanges;
extern bool EnableNewNamespaceValidation;

/* user-defined functions */
PG_FUNCTION_INFO_V1(command_collection_table);
PG_FUNCTION_INFO_V1(command_invalidate_collection_cache);
PG_FUNCTION_INFO_V1(command_get_next_collection_id);
PG_FUNCTION_INFO_V1(command_ensure_valid_db_coll);
PG_FUNCTION_INFO_V1(validate_dbname);
PG_FUNCTION_INFO_V1(command_get_collection);
PG_FUNCTION_INFO_V1(command_get_collection_or_view);
PG_FUNCTION_INFO_V1(validate_collection_cache_entry);

/* forward declarations */
static void InitializeCollectionsHash(void);
static bool GetMongoCollectionFromCatalogById(uint64 collectionId, Oid relationId,
											  MongoCollection *collection);
static bool GetMongoCollectionFromCatalogByNameDatum(Datum databaseNameDatum,
													 Datum collectionNameDatum,
													 MongoCollection *collection);
static Oid GetRelationIdForCollectionTableName(char *collectionTableName,
											   LOCKMODE lockMode);
static AttrNumber GetMongoDataCreationTimeVarAttrNumber(Oid collectionOid);
static MongoCollection * GetMongoCollectionByNameDatumCore(Datum databaseNameDatum,
														   Datum collectionNameDatum,
														   LOCKMODE lockMode);
static Datum GetCollectionOrViewCore(PG_FUNCTION_ARGS, bool allowViews);
static void CopyCollectionOptions(const pgbson *options,
								  MongoCollection *collection);
static void UpdateCollectionOptions(MongoCollection *collection, const
									pgbson *updateSpec);
static void RemoveCollectionOptions(MongoCollection *collection, const
									pgbson *removeSpec);
static void ValidateSystemNamespace(const char *databaseName, const char *collectionName);


/*
 * CollectionCacheIsValid determines whether the collections hashes are
 * valid. It is set to false before initialization, if OOMs occurred while
 * in critical sections of cache construction, and after global invalidations.
 */
static bool CollectionCacheIsValid = false;

/* memory context in which we allocate collections hashes */
static MemoryContext CollectionsCacheContext = NULL;

/* (database name, collection name) -> collection hash */
static HTAB *NameToCollectionHash = NULL;

/* (relation OID) -> collection hash */
static HTAB *RelationIdToCollectionHash = NULL;


/*
 * InitializeCollectionsHash (re)creates the collections hashes if they are
 * not valid.
 *
 * At the end of this function, either CollectionCacheIsValid is true or
 * an OOM was thrown. In the latter case, we will try again on the next
 * call.
 */
static void
InitializeCollectionsHash(void)
{
	/* make sure the metadata cache is initalized */
	InitializeDocumentDBApiExtensionCache();

	/* should not call this function if extension does not exist */
	Assert(IsDocumentDBApiExtensionActive());

	if (CollectionCacheIsValid)
	{
		/* already built the hashes */
		return;
	}

	if (CollectionsCacheContext == NULL)
	{
		CollectionsCacheContext = AllocSetContextCreate(CacheMemoryContext,
														"Cache context for collection",
														ALLOCSET_DEFAULT_SIZES);
	}

	/* reset any previously allocated memory */
	MemoryContextReset(CollectionsCacheContext);

	int hashFlags = HASH_ELEM | HASH_BLOBS | HASH_CONTEXT;

	/* create the (database name, collection name) -> collection hash */
	HASHCTL info;
	memset(&info, 0, sizeof(info));
	info.keysize = sizeof(MongoCollectionName);
	info.entrysize = sizeof(NameToCollectionCacheEntry);
	info.hcxt = CollectionsCacheContext;

	NameToCollectionHash = hash_create("Name to Collection ID Hash", 32,
									   &info, hashFlags);

	/* create the (relation OID) -> collection hash */
	memset(&info, 0, sizeof(info));
	info.keysize = sizeof(Oid);
	info.entrysize = sizeof(RelationIdToCollectionCacheEntry);
	info.hcxt = CollectionsCacheContext;

	RelationIdToCollectionHash = hash_create("Relation ID to Collection ID Hash", 32,
											 &info, hashFlags);

	CollectionCacheIsValid = true;
}


/*
 * ResetCollectionsCache is called when we rebuild the cache from scratch.
 * We need not worry about freeing memory here, since DocumentDBApiMetadataCacheContext
 * is reset as part of the process. We only set ResetCollectionsCache such
 * that we rebuild the hashes in InitializeCollectionsHash.
 */
void
ResetCollectionsCache(void)
{
	CollectionCacheIsValid = false;
}


/*
 * InvalidateCollectionByRelationId is called when receiving invalidation of a specific
 * relation ID.
 *
 * This can happen any time postgres code calls AcceptInvalidationMessages(), e.g
 * after obtaining a relation lock. We remove entries from the cache. They will
 * still be temporarily usable until new entries are added to the cache.
 */
void
InvalidateCollectionByRelationId(Oid relationId)
{
	if (!CollectionCacheIsValid)
	{
		/* hashes are not valid, just wait for cache rebuild */
		return;
	}

	/* delete entry from the relation OID -> collection cache (if any) */
	bool foundInCache = false;
	RelationIdToCollectionCacheEntry *entryById =
		hash_search(RelationIdToCollectionHash, &relationId, HASH_REMOVE,
					&foundInCache);

	if (foundInCache)
	{
		MongoCollection *collection = &(entryById->collection);

		/* delete entry from the collection name -> collection cache */
		NameToCollectionCacheEntry *entryByName =
			hash_search(NameToCollectionHash, &(collection->name),
						HASH_REMOVE, &foundInCache);

		/*
		 * The name-hash entry may not exist if the collection was
		 * only loaded via GetMongoCollectionByColId (which only
		 * populates RelationIdToCollectionHash, not NameToCollectionHash).
		 */

		if (foundInCache)
		{
			/* signal to callers that the entry is no longer valid */
			entryByName->isValid = false;
		}

		/* signal to callers that the entry is no longer valid */
		entryById->isValid = false;
	}
}


/*
 * CopyMongoCollection returns a copy of given MongoCollection.
 */
MongoCollection *
CopyMongoCollection(const MongoCollection *collection)
{
	MongoCollection *copiedCollection = palloc0(sizeof(MongoCollection));

	*copiedCollection = *collection;
	copiedCollection->shardKey = !copiedCollection->shardKey ? NULL :
								 CopyPgbsonIntoMemoryContext(copiedCollection->shardKey,
															 CurrentMemoryContext);
	copiedCollection->viewDefinition = !copiedCollection->viewDefinition ? NULL :
									   CopyPgbsonIntoMemoryContext(
		copiedCollection->viewDefinition,
		CurrentMemoryContext);

	copiedCollection->schemaValidator.validator =
		!copiedCollection->schemaValidator.validator ? NULL :
		CopyPgbsonIntoMemoryContext(
			copiedCollection->schemaValidator.validator,
			CurrentMemoryContext);

	return copiedCollection;
}


/*
 * GetMongoCollectionByColId() gets the MongoCollection metadata by collectionId.
 */
MongoCollection *
GetMongoCollectionByColId(uint64 collectionId, LOCKMODE lockMode)
{
	/* make sure hashes exist */
	InitializeCollectionsHash();

	Oid documentsTableOid = GetRelationIdForCollectionId(collectionId, lockMode);

	if (!OidIsValid(documentsTableOid))
	{
		/* table was dropped */
		return NULL;
	}

	bool foundInCache = false;

	RelationIdToCollectionCacheEntry *entryByRelId =
		hash_search(RelationIdToCollectionHash, &documentsTableOid, HASH_FIND,
					&foundInCache);

	if (foundInCache)
	{
		/* now that we have a lock, check for invalidations */
		AcceptInvalidationMessages();

		/*
		 * After acquiring the lock on the table, we may have received an invalidation
		 * that could indicate a rename. In that case, CollectionCacheIsValid or
		 * entry->isValid is set to false and we treat this as a cache miss by
		 * continuing below.
		 */
		if (!CollectionCacheIsValid)
		{
			InitializeCollectionsHash();
		}
		else if (entryByRelId->isValid)
		{
			return CopyMongoCollection(&(entryByRelId->collection));
		}
	}

	MongoCollection collection;
	memset(&collection, 0, sizeof(collection));

	/*
	 * Temporarily disable unimportant logs related to collection catalog lookup
	 * so that regression test outputs don't become flaky (e.g.: due to commands
	 * being executed by Citus locally).
	 */
	int savedGUCLevel = NewGUCNestLevel();
	SetGUCLocally("client_min_messages", "WARNING");

	/* Read the collection metadata from ApiCatalogSchemaName.collections */
	bool collectionExists =
		GetMongoCollectionFromCatalogById(collectionId, documentsTableOid,
										  &collection);

	/* rollback the GUC change that we made for client_min_messages */
	RollbackGUCChange(savedGUCLevel);

	if (!collectionExists)
	{
		/* no collection record with this id */
		return NULL;
	}

	/* get the relation ID and lock the table */
	if (collection.viewDefinition != NULL)
	{
		/* Views not supported in this path */
		return NULL;
	}

	collection.relationId = GetRelationIdForCollectionTableName(collection.tableName,
																lockMode);

	if (!OidIsValid(collection.relationId))
	{
		/* record exists, but table was dropped (maybe just after reading the record) */
		return NULL;
	}

	/* if we experience OOM below, reset the cache to prevent corruption */
	CollectionCacheIsValid = false;

	/* copy the shard key BSON into the cache memory context */
	if (collection.shardKey != NULL)
	{
		collection.shardKey = CopyPgbsonIntoMemoryContext(collection.shardKey,
														  CollectionsCacheContext);
	}

	collection.mongoDataCreationTimeVarAttrNumber =
		GetMongoDataCreationTimeVarAttrNumber(collection.relationId);

	if (collection.schemaValidator.validator != NULL)
	{
		collection.schemaValidator.validator =
			CopyPgbsonIntoMemoryContext(collection.schemaValidator.validator,
										CollectionsCacheContext);
	}

	/* add to relation ID -> collection hash */
	entryByRelId = hash_search(RelationIdToCollectionHash, &(collection.relationId),
							   HASH_ENTER, &foundInCache);

	entryByRelId->collection = collection;
	entryByRelId->isValid = true;

	/* no OOMs, keep the cache */
	CollectionCacheIsValid = true;

	return CopyMongoCollection(&(entryByRelId->collection));
}


/* Tries to extract the collection ID from a relation OID, returns true if it could and set it to the out param, false otherwise. */
bool
TryGetCollectionIdByRelationOid(Oid relationId, uint64 *collectionId, bool
								requireShardTable)
{
	if (collectionId == NULL)
	{
		return false;
	}

	*collectionId = 0;
	if (get_rel_namespace(relationId) != ApiDataNamespaceOid())
	{
		return false;
	}

	/* Get the relation's name using the relationId. */
	const char *relationName = get_rel_name(relationId);

	/* If the relation lookup failed, return false. */
	if (relationName == NULL)
	{
		return false;
	}

	/* Get the collection ID from the shard table name. */
	if (!CheckRelNameValidity(relationName, collectionId, requireShardTable))
	{
		return false;
	}

	return *collectionId != 0;
}


/*
 * GetMongoCollectionByRelationOid checks if a collection
 * with the specified relationId exists in the ApiCatalogSchemaName.
 */
MongoCollection *
GetMongoCollectionByRelationOid(Oid relationId, bool requireShardTable)
{
	uint64 collectionId = 0;

	/*
	 * If the collection Id can't be parsed from relation's name,
	 * then this is not a valid shard of a collection.
	 */
	if (!TryGetCollectionIdByRelationOid(relationId, &collectionId, requireShardTable))
	{
		return NULL;
	}

	/*
	 * Look up the collection using it's ID and return it. we can directly
	 * lookup from the RelationIdToCollectionHash here, but this collection may
	 * have been evicted from cache, so an invalidation may be needed. So
	 * we call the GetMongoCollectionByColId which handles this already.
	 */
	return GetMongoCollectionByColId(collectionId, NoLock);
}


/*
 * GetMongoCollectionOrViewByNameDatum returns collection metadata by database and
 * collection name or NULL if the collection does not exist. Also returns views
 * if applicable
 */
MongoCollection *
GetMongoCollectionOrViewByNameDatum(Datum databaseNameDatum, Datum collectionNameDatum,
									LOCKMODE lockMode)
{
	return GetMongoCollectionByNameDatumCore(databaseNameDatum, collectionNameDatum,
											 lockMode);
}


/*
 * GetMongoCollectionByName returns collection metadata by database and
 * collection name or NULL if the collection does not exist.
 */
MongoCollection *
GetMongoCollectionByNameDatum(Datum databaseNameDatum, Datum collectionNameDatum,
							  LOCKMODE lockMode)
{
	MongoCollection *collection =
		GetMongoCollectionByNameDatumCore(databaseNameDatum, collectionNameDatum,
										  lockMode);

	if (collection != NULL && collection->viewDefinition != NULL)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_COMMANDNOTSUPPORTEDONVIEW),
						errmsg(
							"The namespace %s.%s refers to a view object rather than a collection",
							collection->name.databaseName,
							collection->name.collectionName)));
	}

	return collection;
}


/*
 * Given a collection, retrieves the shard OIDs and shard names that are associated with
 * that table on the current node.
 */
bool
GetMongoCollectionShardOidsAndNames(MongoCollection *collection, ArrayType **shardIdArray,
									ArrayType **shardNames)
{
	Datum *resultDatums = NULL;
	Datum *resultNameDatums = NULL;
	int32_t shardCount = 0;
	GetShardIdsAndNamesForCollection(collection->relationId, collection->tableName,
									 &resultDatums, &resultNameDatums, &shardCount);

	if (shardCount == 0)
	{
		return false;
	}

	*shardIdArray = construct_array(resultDatums, shardCount, OIDOID,
									sizeof(Oid), true,
									TYPALIGN_INT);
	*shardNames = construct_array(resultNameDatums, shardCount, TEXTOID, -1,
								  false,
								  TYPALIGN_INT);
	pfree(resultDatums);
	pfree(resultNameDatums);
	return true;
}


/*
 * Given a collection, tries to get a shardOID if one is applicable (there is
 * a single shard corresponding to that collection) and if it's available
 * locally. If such a shard is found, locks it with the given lock mode
 * and returns it to the caller.
 */
Oid
TryGetCollectionShardTable(MongoCollection *collection, LOCKMODE lockMode)
{
	/* Cannot proceed without a local shard available. */
	if (collection->shardTableName[0] == '\0')
	{
		return InvalidOid;
	}

	if (!UseLocalExecutionShardQueries)
	{
		return InvalidOid;
	}

	/* Don't allow direct shard access in a multi-statement transaction
	 * This is because switching states between remote execution and local execution can produce
	 * isolation issues. So only support this if we're a single command.
	 */
	if (IsTransactionBlock() && !ForceLocalExecutionShardQueries)
	{
		return InvalidOid;
	}

	Oid relationShardOid = GetRelationIdForCollectionTableName(collection->shardTableName,
															   lockMode);
	ereport(DEBUG3, (errmsg("Has relation shard: %d", relationShardOid != InvalidOid)));
	return relationShardOid;
}


static void
TrySetCollectionShard(MongoCollection *collection)
{
	if (collection->isShardRemote)
	{
		return;
	}

	int savedGUCLevel = NewGUCNestLevel();
	SetGUCLocally("client_min_messages", "WARNING");

	/* Get shard table name (distributed) or original tableName (single node) */
	const char *shardName = TryGetShardNameForUnshardedCollection(
		collection->relationId, collection->collectionId, collection->tableName);
	RollbackGUCChange(savedGUCLevel);
	if (shardName != NULL)
	{
		if (shardName[0] == '\0')
		{
			collection->shardTableName[0] = '\0';
			collection->isShardRemote = true;
		}
		else
		{
			strcpy(collection->shardTableName, shardName);
			collection->isShardRemote = false;
		}
	}
	else
	{
		collection->isShardRemote = false;
		collection->shardTableName[0] = '\0';
	}
}


/*
 * GetMongoCollectionByName returns collection metadata by database and
 * collection name or NULL if the collection does not exist.
 */
static MongoCollection *
GetMongoCollectionByNameDatumCore(Datum databaseNameDatum, Datum collectionNameDatum,
								  LOCKMODE lockMode)
{
	/* make sure hashes exist */
	InitializeCollectionsHash();

	int databaseNameLength = VARSIZE_ANY_EXHDR(databaseNameDatum);
	if (databaseNameLength >= MAX_DATABASE_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE), errmsg(
							"The provided database name exceeds the permitted length")));
	}

	int collectionNameLength = VARSIZE_ANY_EXHDR(collectionNameDatum);
	if (collectionNameLength >= MAX_COLLECTION_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE), errmsg(
							"The specified collection name exceeds the allowed length")));
	}

	MongoCollectionName qualifiedName;
	memset(&qualifiedName, 0, sizeof(qualifiedName));

	/* copy text bytes directly, buffers are already 0-initialized above */
	memcpy(qualifiedName.databaseName, VARDATA_ANY(databaseNameDatum),
		   databaseNameLength);
	memcpy(qualifiedName.collectionName, VARDATA_ANY(collectionNameDatum),
		   collectionNameLength);

	bool foundInCache = false;

	NameToCollectionCacheEntry *entryByName =
		hash_search(NameToCollectionHash, &qualifiedName, HASH_FIND, &foundInCache);

	if (foundInCache)
	{
		/* refresh relation ID based on the name */
		if (entryByName->collection.viewDefinition == NULL)
		{
			entryByName->collection.relationId =
				GetRelationIdForCollectionTableName(entryByName->collection.tableName,
													lockMode);
			if (!OidIsValid(entryByName->collection.relationId))
			{
				/* table was dropped */
				return NULL;
			}

			/* Collection is found and valid - check for local shard */
			if (entryByName->collection.shardKey == NULL &&
				!entryByName->collection.isShardRemote &&
				entryByName->collection.shardTableName[0] == '\0')
			{
				TrySetCollectionShard(&entryByName->collection);
			}
		}

		/* now that we have a lock, check for invalidations */
		AcceptInvalidationMessages();

		/*
		 * After acquiring the lock on the table, we may have received an invalidation
		 * that could indicate a rename. In that case, CollectionCacheIsValid or
		 * entry->isValid is set to false and we treat this as a cache miss by
		 * continuing below.
		 */
		if (!CollectionCacheIsValid)
		{
			InitializeCollectionsHash();
		}
		else if (entryByName->isValid)
		{
			return CopyMongoCollection(&(entryByName->collection));
		}
	}

	MongoCollection collection;
	memset(&collection, 0, sizeof(collection));

	/*
	 * Temporarily disable unimportant logs related to collection catalog lookup
	 * so that regression test outputs don't become flaky (e.g.: due to commands
	 * being executed by Citus locally).
	 */
	int savedGUCLevel = NewGUCNestLevel();
	SetGUCLocally("client_min_messages", "WARNING");

	/*
	 * Read the collection metadata from ApiCatalogSchemaName.collections or error
	 * out if the collection does not exist. (We do not cache negative entries,
	 * since we expect them to be rare)
	 */
	bool collectionExists =
		GetMongoCollectionFromCatalogByNameDatum(databaseNameDatum,
												 collectionNameDatum,
												 &collection);

	/* rollback the GUC change that we made for client_min_messages */
	RollbackGUCChange(savedGUCLevel);

	if (!collectionExists)
	{
		/* no collection record with this name */
		return NULL;
	}

	/* get the relation ID and lock the table */
	if (collection.viewDefinition == NULL)
	{
		collection.relationId =
			GetRelationIdForCollectionTableName(collection.tableName, lockMode);

		if (!OidIsValid(collection.relationId))
		{
			/* record exists, but table was dropped (maybe just after reading the record) */
			return NULL;
		}

		if (collection.shardKey == NULL)
		{
			TrySetCollectionShard(&collection);
		}
	}


	/* if we experience OOM below, reset the cache to prevent corruption */
	CollectionCacheIsValid = false;

	/* copy the shard key BSON into the cache memory context */
	if (collection.shardKey != NULL)
	{
		collection.shardKey = CopyPgbsonIntoMemoryContext(collection.shardKey,
														  CollectionsCacheContext);
	}

	if (collection.viewDefinition != NULL)
	{
		collection.viewDefinition = CopyPgbsonIntoMemoryContext(collection.viewDefinition,
																CollectionsCacheContext);
	}
	else
	{
		collection.mongoDataCreationTimeVarAttrNumber =
			GetMongoDataCreationTimeVarAttrNumber(collection.relationId);
	}

	/* get schema validation meta*/
	if (collection.schemaValidator.validator != NULL)
	{
		collection.schemaValidator.validator = CopyPgbsonIntoMemoryContext(
			collection.schemaValidator.validator,
			CollectionsCacheContext);
	}

	/* collection exists, so write a name -> collection cache entry */
	entryByName =
		hash_search(NameToCollectionHash, &qualifiedName, HASH_ENTER, &foundInCache);

	/*
	 * We lazily copy the whole collection struct instead of trying to be overly
	 * clever about keeping only a single copy.
	 */
	entryByName->collection = collection;
	entryByName->isValid = true;

	/* also add to relation ID -> collection hash */
	RelationIdToCollectionCacheEntry *entryByRelId =
		hash_search(RelationIdToCollectionHash, &(collection.relationId),
					HASH_ENTER, &foundInCache);

	entryByRelId->collection = collection;
	entryByRelId->isValid = true;

	/* no OOMs, keep the cache */
	CollectionCacheIsValid = true;

	return CopyMongoCollection(&(entryByName->collection));
}


MongoCollection *
GetTempMongoCollectionByNameDatum(Datum databaseNameDatum, Datum collectionNameDatum,
								  char *collectionName,
								  LOCKMODE lockMode)
{
	MongoCollection *collection = palloc0(sizeof(MongoCollection));

	int databaseNameLength = VARSIZE_ANY_EXHDR(databaseNameDatum);
	if (databaseNameLength >= MAX_DATABASE_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_BADVALUE), errmsg(
							"The provided database name exceeds the permitted length")));
	}

	int collectionNameLength = VARSIZE_ANY_EXHDR(collectionNameDatum);
	if (collectionNameLength >= MAX_COLLECTION_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_BADVALUE), errmsg(
							"The specified collection name exceeds the allowed length")));
	}

	/* copy text bytes directly, buffers are already 0-initialized above */
	memcpy(collection->name.databaseName, VARDATA_ANY(databaseNameDatum),
		   databaseNameLength);
	memcpy(collection->name.collectionName, VARDATA_ANY(collectionNameDatum),
		   collectionNameLength);

	collection->shardKey = NULL;
	collection->viewDefinition = NULL;
	collection->collectionId = UINT64_MAX; /*unused */
	collection->relationId = InvalidOid; /* unused */
	sprintf(collection->tableName, "documents_temp");
	collection->shardTableName[0] = '\0';

	return collection;
}


/*
 * IsDataTableCreatedWithinCurrentXact returns true if data table for given
 * collection has been created within the current transaction.
 */
bool
IsDataTableCreatedWithinCurrentXact(const MongoCollection *collection)
{
	HeapTuple pgCatalogTuple =
		SearchSysCache1(RELOID, ObjectIdGetDatum(collection->relationId));
	if (!HeapTupleIsValid(pgCatalogTuple))
	{
		ereport(ERROR, (errmsg(
							"The data table associated with the collection identified by "
							UINT64_FORMAT
							" cannot be found.",
							collection->collectionId)));
	}

	bool dataTableCreatedWithinCurrentXact =
		HeapTupleHeaderGetXmin(pgCatalogTuple->t_data) == GetCurrentTransactionId();

	ReleaseSysCache(pgCatalogTuple);

	return dataTableCreatedWithinCurrentXact;
}


/*
 * GetMongoCollectionFromCatalogById returns whether a collection
 * with the given id exist in ApiCatalogSchemaName.collections and writes
 * the metadata to the collection struct.
 */
static bool
GetMongoCollectionFromCatalogById(uint64 collectionId, Oid relationId,
								  MongoCollection *collection)
{
	bool collectionExists = false;

	StringInfoData query;
	const int argCount = 1;
	Oid argTypes[1];
	Datum argValues[1];

	memset(collection, 0, sizeof(MongoCollection));
	MemoryContext outerContext = CurrentMemoryContext;
	SPI_connect();

	initStringInfo(&query);
	appendStringInfo(&query,
					 "SELECT * FROM %s.collections WHERE collection_id = $1",
					 ApiCatalogSchemaName);

	argTypes[0] = INT8OID;
	argValues[0] = UInt64GetDatum(collectionId);

	SPI_execute_with_args(query.data, argCount, argTypes, argValues, NULL, false, 1);
	if (SPI_processed > 0)
	{
		TupleDesc tupleDescriptor = SPI_tuptable->tupdesc;
		HeapTuple tuple = SPI_tuptable->vals[0];
		bool isNull = false;

		Datum databaseNameDatum = heap_getattr(tuple, 1, tupleDescriptor, &isNull);
		if (isNull)
		{
			ereport(ERROR, (errmsg("database_name should not be NULL in catalog")));
		}

		memcpy(collection->name.databaseName, VARDATA_ANY(databaseNameDatum),
			   VARSIZE_ANY_EXHDR(databaseNameDatum));

		Datum collectionNameDatum = heap_getattr(tuple, 2, tupleDescriptor, &isNull);
		if (isNull)
		{
			ereport(ERROR, (errmsg("collection name must not be NULL")));
		}

		memcpy(collection->name.collectionName, VARDATA_ANY(collectionNameDatum),
			   VARSIZE_ANY_EXHDR(collectionNameDatum));

		/* Attribute 3 refers to the collection_id */
		collection->collectionId = collectionId;

		Datum shardKeyDatum = heap_getattr(tuple, 4, tupleDescriptor, &isNull);
		if (!isNull)
		{
			pgbson *shardKeyBson = DatumGetPgBson(shardKeyDatum);
			collection->shardKey =
				CopyPgbsonIntoMemoryContext(shardKeyBson, outerContext);
		}

		/* Attr 5 is the collection_uuid */
		Datum uuidDatum = heap_getattr(tuple, 5, tupleDescriptor, &isNull);
		if (!isNull)
		{
			collection->collectionUUID = *DatumGetUUIDP(uuidDatum);
		}

		if (tupleDescriptor->natts >= 6)
		{
			Datum viewDatum = heap_getattr(tuple, 6, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *viewDefinition = DatumGetPgBson(viewDatum);
				collection->viewDefinition =
					CopyPgbsonIntoMemoryContext(viewDefinition, outerContext);
			}
		}

		/* Attr 7,8,9 are for Schema Validation */
		if (tupleDescriptor->natts >= 9)
		{
			/* validator stored as pgbson */
			Datum validatorDatum = heap_getattr(tuple, 7, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *validator = DatumGetPgBson(validatorDatum);
				collection->schemaValidator.validator =
					CopyPgbsonIntoMemoryContext(validator, outerContext);
			}

			/* validation level stored as text */
			Datum validationLevelDatum = heap_getattr(tuple, 8, tupleDescriptor, &isNull);
			if (!isNull)
			{
				const char *validationLevelText = TextDatumGetCString(
					validationLevelDatum);
				collection->schemaValidator.validationLevel =
					strcmp(validationLevelText, "off") == 0 ? ValidationLevel_Off :
					strcmp(validationLevelText, "strict") == 0 ? ValidationLevel_Strict :
					strcmp(validationLevelText, "moderate") == 0 ?
					ValidationLevel_Moderate :
					ValidationLevel_Invalid;
			}

			/* validation action stored as text */
			Datum validationActionDatum = heap_getattr(tuple, 9, tupleDescriptor,
													   &isNull);
			if (!isNull)
			{
				const char *validationActionText = TextDatumGetCString(
					validationActionDatum);
				collection->schemaValidator.validationAction =
					strcmp(validationActionText, "warn") == 0 ? ValidationAction_Warn :
					strcmp(validationActionText, "error") == 0 ? ValidationAction_Error :
					ValidationAction_Invalid;
			}
		}

		/* 10 is collection options */
		if (tupleDescriptor->natts >= 10)
		{
			/* options stored as pgbson */
			Datum optionsDatum = heap_getattr(tuple, 10, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *options = DatumGetPgBson(optionsDatum);
				CopyCollectionOptions(options, collection);
			}
		}

		/* table name is: documents_<collection id> */
		snprintf(collection->tableName, NAMEDATALEN, DOCUMENT_DATA_TABLE_NAME_FORMAT,
				 collection->collectionId);

		collection->collectionId = collectionId;
		collection->relationId = relationId;
		if (collection->shardKey == NULL && collection->viewDefinition == NULL)
		{
			TrySetCollectionShard(collection);
		}

		collectionExists = true;
	}

	pfree(query.data);

	SPI_finish();

	return collectionExists;
}


/*
 * GetMongoCollectionFromCatalogByNameDatum returns whether a collection
 * with the given name (as database and collection name datums) exist in
 * ApiCatalogSchemaName.collections and writes the metadata to the collection
 * struct.
 */
static bool
GetMongoCollectionFromCatalogByNameDatum(Datum databaseNameDatum,
										 Datum collectionNameDatum,
										 MongoCollection *collection)
{
	bool collectionExists = false;

	StringInfoData query;
	const int argCount = 2;
	Oid argTypes[2];
	Datum argValues[2];
	MemoryContext outerContext = CurrentMemoryContext;

	memset(collection, 0, sizeof(MongoCollection));

	SPI_connect();

	initStringInfo(&query);
	appendStringInfo(&query,
					 "SELECT * FROM %s.collections WHERE database_name = $1 AND collection_name = $2",
					 ApiCatalogSchemaName);

	argTypes[0] = TEXTOID;
	argValues[0] = databaseNameDatum;

	argTypes[1] = TEXTOID;
	argValues[1] = collectionNameDatum;

	SPI_execute_with_args(query.data, argCount, argTypes, argValues, NULL, false, 1);
	if (SPI_processed > 0)
	{
		TupleDesc tupleDescriptor = SPI_tuptable->tupdesc;
		HeapTuple tuple = SPI_tuptable->vals[0];
		bool isNull = false;

		/* Attr 1 is database_name */
		/* Attr 2 is collection_name */

		Datum collectionIdDatum = heap_getattr(tuple, 3, tupleDescriptor, &isNull);
		if (isNull)
		{
			ereport(ERROR, (errmsg("collection_id should not be NULL in catalog")));
		}

		collection->collectionId = Int64GetDatum(collectionIdDatum);

		/* copy text bytes, buffers are already 0-initialized by memset above */
		memcpy(collection->name.databaseName, VARDATA_ANY(databaseNameDatum),
			   VARSIZE_ANY_EXHDR(databaseNameDatum));
		memcpy(collection->name.collectionName, VARDATA_ANY(collectionNameDatum),
			   VARSIZE_ANY_EXHDR(collectionNameDatum));

		/* table name is: documents_<collection id> */
		snprintf(collection->tableName, NAMEDATALEN, DOCUMENT_DATA_TABLE_NAME_FORMAT,
				 collection->collectionId);

		Datum shardKeyDatum = heap_getattr(tuple, 4, tupleDescriptor, &isNull);
		if (!isNull)
		{
			pgbson *shardKeyBson = DatumGetPgBson(shardKeyDatum);
			collection->shardKey = CopyPgbsonIntoMemoryContext(shardKeyBson,
															   outerContext);
		}

		/* Attr 5 is collection_uuid */
		Datum uuidDatum = heap_getattr(tuple, 5, tupleDescriptor, &isNull);
		if (!isNull)
		{
			collection->collectionUUID = *DatumGetUUIDP(uuidDatum);
		}

		if (tupleDescriptor->natts >= 6)
		{
			Datum viewDatum = heap_getattr(tuple, 6, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *viewDefinition = DatumGetPgBson(viewDatum);
				collection->viewDefinition =
					CopyPgbsonIntoMemoryContext(viewDefinition, outerContext);
			}
		}

		/* Attr 7,8,9 are for Schema Validation */
		if (tupleDescriptor->natts >= 9)
		{
			/* validator stored as pgbson */
			Datum validatorDatum = heap_getattr(tuple, 7, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *validator = DatumGetPgBson(validatorDatum);
				collection->schemaValidator.validator =
					CopyPgbsonIntoMemoryContext(validator, outerContext);
			}

			/* validation level stored as text */
			Datum validationLevelDatum = heap_getattr(tuple, 8, tupleDescriptor, &isNull);
			if (!isNull)
			{
				const char *validationLevelText = TextDatumGetCString(
					validationLevelDatum);
				collection->schemaValidator.validationLevel =
					strcmp(validationLevelText, "off") == 0 ? ValidationLevel_Off :
					strcmp(validationLevelText, "strict") == 0 ? ValidationLevel_Strict :
					strcmp(validationLevelText, "moderate") == 0 ?
					ValidationLevel_Moderate :
					ValidationLevel_Invalid;
			}

			/* validation action stored as text */
			Datum validationActionDatum = heap_getattr(tuple, 9, tupleDescriptor,
													   &isNull);
			if (!isNull)
			{
				const char *validationActionText = TextDatumGetCString(
					validationActionDatum);
				collection->schemaValidator.validationAction =
					strcmp(validationActionText, "warn") == 0 ? ValidationAction_Warn :
					strcmp(validationActionText, "error") == 0 ? ValidationAction_Error :
					ValidationAction_Invalid;
			}
		}

		/* 10 is collection options */
		if (tupleDescriptor->natts >= 10)
		{
			/* options stored as pgbson */
			Datum optionsDatum = heap_getattr(tuple, 10, tupleDescriptor, &isNull);
			if (!isNull)
			{
				pgbson *options = DatumGetPgBson(optionsDatum);
				CopyCollectionOptions(options, collection);
			}
		}

		collectionExists = true;
	}

	pfree(query.data);

	SPI_finish();

	return collectionExists;
}


/*
 * GetRelationIdForCollectionId returns the OID of the Postgres relation backing
 * the Mongo collection with given collection table id.
 *
 * Returns InvalidOid if no such collection exists.
 */
Oid
GetRelationIdForCollectionId(uint64 collectionId, LOCKMODE lockMode)
{
	char tableName[NAMEDATALEN] = { 0 };
	sprintf(tableName, DOCUMENT_DATA_TABLE_NAME_FORMAT, collectionId);
	Oid relationId = GetRelationIdForCollectionTableName(tableName,
														 lockMode);

	return relationId;
}


/*
 * GetRelationIdForCollectionTableName returns the OID of the Postgres relation backing
 * the Mongo collection with given collection table name.
 *
 * Returns InvalidOid if no such collection exists.
 */
static Oid
GetRelationIdForCollectionTableName(char *collectionTableName, LOCKMODE lockMode)
{
	bool missingOK = true;
	RangeVar *rangeVar = makeRangeVar(ApiDataSchemaName, collectionTableName, -1);

	return RangeVarGetRelid(rangeVar, lockMode, missingOK);
}


/*
 * GetMongoDataCreationTimeVarAttrNumber returns the attribute number of creation_time column.
 */
static AttrNumber
GetMongoDataCreationTimeVarAttrNumber(Oid collectionOid)
{
	HeapTuple tuple = SearchSysCacheAttName(collectionOid, "creation_time");

	if (!HeapTupleIsValid(tuple))
	{
		/* creation_time column is not present */
		return (AttrNumber) - 1;
	}

	Form_pg_attribute targetatt = (Form_pg_attribute) GETSTRUCT(tuple);
	int16 attnum = targetatt->attnum;
	ReleaseSysCache(tuple);

	return attnum;
}


/*
 * Check if DB exists. Check is done case insensitively. If exists, return
 * TRUE and populates the output parameter dbNameInTable with the db name
 * from the catalog table, else FALSE
 */
bool
TryGetDBNameByDatum(Datum databaseNameDatum, char *dbNameInTable)
{
	bool dbExists = false;
	StringInfoData query;
	const int argCount = 1;
	Oid argTypes[1];
	Datum argValues[1];

	SPI_connect();

	initStringInfo(&query);
	appendStringInfo(&query,
					 "SELECT database_name FROM %s.collections WHERE LOWER(database_name) = LOWER($1) LIMIT 1",
					 ApiCatalogSchemaName);

	argTypes[0] = TEXTOID;
	argValues[0] = databaseNameDatum;

	SPI_execute_with_args(query.data, argCount, argTypes, argValues, NULL, false, 1);
	if (SPI_processed > 0)
	{
		TupleDesc tupleDescriptor = SPI_tuptable->tupdesc;
		HeapTuple tuple = SPI_tuptable->vals[0];
		bool isNull = false;

		Datum databaseNameDatumInt = heap_getattr(tuple, 1, tupleDescriptor, &isNull);
		if (isNull)
		{
			ereport(ERROR, (errmsg("database_name should not be NULL in catalog")));
		}

		memcpy(dbNameInTable, VARDATA_ANY(databaseNameDatumInt),
			   VARSIZE_ANY_EXHDR(databaseNameDatumInt));
		dbExists = true;
	}

	pfree(query.data);

	SPI_finish();

	return dbExists;
}


/*
 * Checks if the given collection belongs to the group of Non writable system
 * namespace. If yes, an ereport is done.
 */
void
ValidateCollectionNameForUnauthorizedSystemNs(const char *collectionName,
											  Datum databaseNameDatum)
{
	/* Collection name cannot be empty*/
	if (strlen(collectionName) == 0)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("An invalid and empty namespace has been specified")));
	}
	for (int i = 0; i < NonWritableSystemCollectionNamesLength; i++)
	{
		if (strcmp(collectionName, NonWritableSystemCollectionNames[i]) == 0)
		{
			StringView databaseView = {
				.length = VARSIZE_ANY_EXHDR(databaseNameDatum),
				.string = VARDATA_ANY(databaseNameDatum)
			};

			/* Need to disallow user writes on NonWritableSystemCollectionNames */
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
							errmsg("Unable to write data to specified location %.*s.%s",
								   databaseView.length, databaseView.string,
								   NonWritableSystemCollectionNames[i])));
		}
	}
}


/*
 * Checks if the given collection name belongs to a valid system namespace
 */
void
ValidateCollectionNameForValidSystemNamespace(StringView *collectionView,
											  Datum databaseNameDatum)
{
	char *collectionName = CreateStringFromStringView(collectionView);
	if (StringViewStartsWithStringView(collectionView, &SystemPrefix))
	{
		bool found = false;
		for (int i = 0; i < ValidSystemCollectionNamesLength; i++)
		{
			if (strcmp(collectionName, ValidSystemCollectionNames[i]) == 0)
			{
				found = true;
				ValidateCollectionNameForUnauthorizedSystemNs(
					collectionName, databaseNameDatum);
				break;
			}
		}

		if (!found)
		{
			StringView databaseView = {
				.length = VARSIZE_ANY_EXHDR(databaseNameDatum),
				.string = VARDATA_ANY(databaseNameDatum)
			};
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
							errmsg("System namespace provided is invalid: %.*s.%.*s",
								   databaseView.length, databaseView.string,
								   collectionView->length, collectionView->string)));
		}
	}

	if (EnableNewNamespaceValidation)
	{
		ValidateSystemNamespace(text_to_cstring(DatumGetTextP(databaseNameDatum)),
								collectionName);
	}
}


static int
CompareStringPointers(const void *a, const void *b)
{
	const char *key = *(const char **) a;
	const char *elem = *(const char **) b;
	return strcmp(key, elem);
}


/*
 * this function restricts the usage of certain namespaces that are reserved for internal use.
 * For example, we reserve the 'config' database for internal use and disallow users from creating collections with certain names in that database.
 * TODO : add restrictions for 'admin' and 'local' database
 */
static void
ValidateSystemNamespace(const char *databaseName, const char *collectionName)
{
	/* restrict limited collection in 'config' db (sorted for bsearch) */
	const char *limitedConfigCollections[] = {
		"_shards", "chunks", "collections", "databases", "settings", "shards", "version"
	};
	const int limitedConfigCollectionsLength = sizeof(limitedConfigCollections) /
											   sizeof(char *);
	if (strcmp(databaseName, "config") == 0)
	{
		if (bsearch(&collectionName, limitedConfigCollections,
					limitedConfigCollectionsLength, sizeof(const char *),
					CompareStringPointers) != NULL)
		{
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_ILLEGALOPERATION),
							errmsg(
								"The namespace %s.%s is reserved and cannot be used",
								databaseName, collectionName)));
		}
	}
}


/*
 * command_collection_table returns the OID of a table that stores the data
 * for a given Mongo collection and locks the table for reads.
 */
Datum
command_collection_table(PG_FUNCTION_ARGS)
{
	Datum databaseNameDatum = PG_GETARG_DATUM(0);
	Datum collectionNameDatum = PG_GETARG_DATUM(1);

	MongoCollection *collection = GetMongoCollectionByNameDatum(databaseNameDatum,
																collectionNameDatum,
																AccessShareLock);
	if (collection == NULL)
	{
		PG_RETURN_NULL();
	}

	PG_RETURN_OID(collection->relationId);
}


/*
 * command_invalidate_collection_cache sends an invalidation message that clears
 * the collection caches.
 */
Datum
command_invalidate_collection_cache(PG_FUNCTION_ARGS)
{
	if (EnableOnlyCollectionCacheInvalidateOnCollectionChanges)
	{
		InvalidateCollectionsCache();
	}
	else
	{
		CacheInvalidateRelcacheAll();
	}

	PG_RETURN_VOID();
}


/*
 * command_get_collection_or_view returns the output of the ApiCatalogSchemaName.collections
 * table (all attributes) by database and collection name whether the collection is a
 * collection or a view
 */
Datum
command_get_collection_or_view(PG_FUNCTION_ARGS)
{
	bool allowViews = true;
	Datum returnedDatum = GetCollectionOrViewCore(fcinfo, allowViews);
	PG_RETURN_DATUM(returnedDatum);
}


/*
 * validate_collection_cache_entry is a C function used in regression tests to validate
 * that the collection cache is properly populated. It retrieves the same collection
 * metadata from the cache by both name and id and compares the results to ensure they match.
 * There are two ways to get collections today one is by name and the other is by id, if any new access method is
 * added it should be added to the function as well, so that we can ensure that all the access methods are properly populating the cache.
 */
Datum
validate_collection_cache_entry(PG_FUNCTION_ARGS)
{
	Datum databaseNameDatum = PG_GETARG_DATUM(0);
	Datum collectionNameDatum = PG_GETARG_DATUM(1);

	MongoCollection *collectionByName = GetMongoCollectionByNameDatum(databaseNameDatum,
																	  collectionNameDatum,
																	  AccessShareLock);

	/* Invalidate collection cache, to fetch a new entry by id */
	InvalidateCollectionsCache();

	MongoCollection *collectionById = GetMongoCollectionByColId(
		collectionByName->collectionId,
		AccessShareLock);

	int comparisonResult = memcmp(collectionByName, collectionById,
								  sizeof(MongoCollection));

	PG_RETURN_BOOL(comparisonResult == 0);
}


/*
 * command_get_collection returns the output of the ApiCatalogSchemaName.collections
 * table (all attributes) by database and collection name if the result is a collection.
 * The query fails if the result is a view.
 */
Datum
command_get_collection(PG_FUNCTION_ARGS)
{
	bool allowViews = false;
	Datum returnedDatum = GetCollectionOrViewCore(fcinfo, allowViews);
	PG_RETURN_DATUM(returnedDatum);
}


/*
 * GetCollectionOrViewCore applies the core logic of returning the ApiCatalogSchemaName.collections
 * by database and collection name. Applies filtering based on whether allowViews is specified.
 */
static Datum
GetCollectionOrViewCore(PG_FUNCTION_ARGS, bool allowViews)
{
	Datum databaseDatum = PG_GETARG_DATUM(0);
	Datum collectionName = PG_GETARG_DATUM(1);

	Oid *resultTypeId = NULL;
	TupleDesc resultTupDesc = NULL;
	get_call_result_type(fcinfo, resultTypeId, &resultTupDesc);

	Datum *resultValues = palloc0(sizeof(Datum) * resultTupDesc->natts);
	bool *resultIsNulls = palloc0(sizeof(bool) * resultTupDesc->natts);

	SPI_connect();

	int tupleCountLimit = 1;
	const char *query =
		FormatSqlQuery(
			"SELECT * FROM %s.collections WHERE database_name = $1 and collection_name = $2",
			ApiCatalogSchemaName);
	int nargs = 2;
	Oid argTypes[2] = { TEXTOID, TEXTOID };
	Datum argValues[2] = { databaseDatum, collectionName };
	char *argNulls = NULL;
	bool readOnly = true;
	if (SPI_execute_with_args(query, nargs, argTypes, argValues, argNulls,
							  readOnly, tupleCountLimit) != SPI_OK_SELECT)
	{
		ereport(ERROR, (errmsg("could not run SPI query")));
	}

	bool hasCollection = false;
	if (SPI_processed >= 1 && SPI_tuptable)
	{
		hasCollection = true;
		AttrNumber i = 1;
		for (i = 1; i <= SPI_tuptable->tupdesc->natts && i <= resultTupDesc->natts; i++)
		{
			int tupleNumber = 0;
			bool isNull = true;
			Datum resultDatum = SPI_getbinval(SPI_tuptable->vals[tupleNumber],
											  SPI_tuptable->tupdesc, i, &isNull);
			if (isNull)
			{
				resultIsNulls[i - 1] = true;
				resultValues[i - 1] = (Datum) 0;
			}
			else
			{
				resultIsNulls[i - 1] = false;
				resultValues[i - 1] = SPI_datumTransfer(resultDatum,
														TupleDescAttr(
															SPI_tuptable->tupdesc, i -
															1)->attbyval,
														TupleDescAttr(
															SPI_tuptable->tupdesc, i -
															1)->attlen);
			}
		}

		for (; i <= resultTupDesc->natts; i++)
		{
			resultIsNulls[i - 1] = true;
			resultValues[i - 1] = (Datum) 0;
		}
	}

	SPI_finish();

	if (hasCollection)
	{
		if (!allowViews && resultTupDesc->natts > 5 && !resultIsNulls[5])
		{
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_COMMANDNOTSUPPORTEDONVIEW),
							errmsg(
								"The namespace %s.%s refers to a view object rather than a collection",
								TextDatumGetCString(databaseDatum),
								TextDatumGetCString(collectionName))));
		}

		HeapTuple resultTup = heap_form_tuple(resultTupDesc, resultValues, resultIsNulls);
		PG_RETURN_DATUM(HeapTupleGetDatum(resultTup));
	}
	else
	{
		PG_RETURN_NULL();
	}
}


/*
 * CreateCollection is a C wrapper around create_collection. It returns
 * whether a new collection was created (in case of a race, another
 * transaction may have created it).
 */
bool
CreateCollection(Datum dbNameDatum, Datum collectionNameDatum)
{
	char *collectionNameStr = TextDatumGetCString(collectionNameDatum);
	if (collectionNameStr != NULL && strlen(collectionNameStr) == 0)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("An invalid and empty namespace has been specified")));
	}

	const char *cmdStr = FormatSqlQuery("SELECT %s.create_collection($1, $2)",
										ApiSchemaName);

	Oid argTypes[CREATE_COLLECTION_FUNC_NARGS] = { TEXTOID, TEXTOID };
	Datum argValues[CREATE_COLLECTION_FUNC_NARGS] = {
		dbNameDatum,
		collectionNameDatum,
	};

	/* all args are non-null */
	char *argNulls = NULL;

	bool isNull = true;
	bool readOnly = false;
	Datum resultDatum = ExtensionExecuteQueryWithArgsViaSPI(cmdStr,
															CREATE_COLLECTION_FUNC_NARGS,
															argTypes, argValues, argNulls,
															readOnly, SPI_OK_SELECT,
															&isNull);
	if (isNull)
	{
		ereport(ERROR, (errmsg("create_collection unexpectedly "
							   "returned NULL datum")));
	}

	return DatumGetBool(resultDatum);
}


/*
 * RenameCollection is a C wrapper around rename_collection. It returns
 * whether whether the collection was renamed.
 */
void
RenameCollection(Datum dbNameDatum, Datum srcCollectionNameDatum, Datum
				 destCollectionNameDatum, bool dropTarget)
{
	const char *cmdStr = FormatSqlQuery("SELECT %s.rename_collection($1, $2, $3, $4)",
										ApiSchemaName);

	Oid argTypes[4] = { TEXTOID, TEXTOID, TEXTOID, BOOLOID };
	Datum argValues[4] = {
		dbNameDatum,
		srcCollectionNameDatum,
		destCollectionNameDatum,
		BoolGetDatum(dropTarget)
	};

	/* all args are non-null */
	char *argNulls = NULL;

	bool isNull = true;
	bool readOnly = false;
	ExtensionExecuteQueryWithArgsViaSPI(cmdStr,
										4,
										argTypes, argValues, argNulls,
										readOnly, SPI_OK_SELECT,
										&isNull);
	if (isNull)
	{
		ereport(ERROR, (errmsg("rename_collection unexpectedly "
							   "returned NULL datum")));
	}
}


/*
 * command_get_next_collection_id returns next unique collection id based on
 * the value of ApiGucPrefix.next_collection_id GUC if it is set.
 *
 * Otherwise, uses the next value of collections_collection_id_seq sequence.
 *
 * Note that ApiGucPrefix.next_collection_id GUC is only expected to be set in
 * regression tests to ensure consistent collection ids when running tests
 * in parallel.
 */
Datum
command_get_next_collection_id(PG_FUNCTION_ARGS)
{
	if (NextCollectionId != NEXT_COLLECTION_ID_UNSET)
	{
		int collectionId = NextCollectionId++;
		PG_RETURN_DATUM(UInt64GetDatum(collectionId));
	}

	PG_RETURN_DATUM(SequenceGetNextValAsUser(ApiCatalogCollectionIdSequenceId(),
											 DocumentDBApiExtensionOwner()));
}


/*
 * Validation function that ensures that the database/collections created in
 * documentdb_api are valid.
 */
Datum
command_ensure_valid_db_coll(PG_FUNCTION_ARGS)
{
	ValidateDatabaseCollection(PG_GETARG_DATUM(0), PG_GETARG_DATUM(1));
	PG_RETURN_BOOL(true);
}


/*
 * Validation function that ensures that the database/collections created in
 * documentdb_api are valid.
 */
void
ValidateDatabaseCollection(Datum databaseDatum, Datum collectionDatum)
{
	text *databaseName = DatumGetTextP(databaseDatum);
	text *collectionName = DatumGetTextP(collectionDatum);

	StringView databaseView = {
		.length = VARSIZE_ANY_EXHDR(databaseName), .string = VARDATA_ANY(databaseName)
	};
	StringView collectionView = {
		.length = VARSIZE_ANY_EXHDR(collectionName), .string = VARDATA_ANY(collectionName)
	};

	if (databaseView.length >= MAX_DATABASE_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("Database %.*s must be less than 64 characters",
							   databaseView.length, databaseView.string)));
	}

	for (int i = 0; i < CharactersNotAllowedInDatabaseNamesLength; i++)
	{
		if (StringViewContains(&databaseView, CharactersNotAllowedInDatabaseNames[i]))
		{
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
							errmsg("Database %.*s has an invalid character %c",
								   databaseView.length, databaseView.string,
								   CharactersNotAllowedInDatabaseNames[i])));
		}
	}

	if (collectionView.string == NULL || collectionView.length == 0)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("The specified namespace provided '%.*s.' is invalid.",
							   databaseView.length, databaseView.string)));
	}

	if (StringViewStartsWith(&collectionView, '.'))
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg(
							"Collection names must not begin with the '.' character: %.*s",
							collectionView.length, collectionView.string)));
	}

	for (int i = 0; i < CharactersNotAllowedInCollectionNamesLength; i++)
	{
		if (StringViewContains(&collectionView, CharactersNotAllowedInCollectionNames[i]))
		{
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
							errmsg("Collection name provided is invalid: %.*s",
								   collectionView.length, collectionView.string)));
		}
	}

	if (databaseView.length + collectionView.length + 1 > MaxDatabaseCollectionLength)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("Full namespace must not exceed %u bytes.",
							   MaxDatabaseCollectionLength)));
	}

	ValidateCollectionNameForValidSystemNamespace(&collectionView,
												  PointerGetDatum(databaseName));
}


/*
 * Validation function that ensures that the database name is unique
 * case insensitively
 */
Datum
validate_dbname(PG_FUNCTION_ARGS)
{
	if (PG_ARGISNULL(0))
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("Database name cannot be NULL")));
	}

	text *databaseName = PG_GETARG_TEXT_P(0);

	StringView databaseView = {
		.length = VARSIZE_ANY_EXHDR(databaseName), .string = VARDATA_ANY(databaseName)
	};

	if (databaseView.length >= MAX_DATABASE_NAME_LENGTH)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDNAMESPACE),
						errmsg("Database %.*s must be less than 64 characters",
							   databaseView.length, databaseView.string)));
	}

	char dbNameInTable[MAX_DATABASE_NAME_LENGTH] = { 0 };
	if (TryGetDBNameByDatum(PG_GETARG_DATUM(0), (char *) dbNameInTable))
	{
		if (!StringViewEqualsCString(&databaseView, (char *) dbNameInTable))
		{
			ereport(ERROR,
					(errcode(ERRCODE_DOCUMENTDB_DBALREADYEXISTS),
					 errmsg(
						 "Database already exists but with a different case; existing entry: [%s], attempted creation: [%.*s]",
						 dbNameInTable,
						 databaseView.length, databaseView.string)));
		}
	}

	PG_RETURN_VOID();
}


/*
 * Updates the schema validation information of an existing collection
 */
void
UpsertSchemaValidation(Datum databaseDatum,
					   Datum collectionNameDatum,
					   const bson_value_t *validator, char *validationLevel,
					   char *validationAction)
{
	StringInfo query = makeStringInfo();
	appendStringInfo(query, "UPDATE %s.collections set ", ApiCatalogSchemaName);

	Oid argsTypes[5] = { TEXTOID, TEXTOID };
	Datum argValues[5] = {
		databaseDatum, collectionNameDatum
	};

	int nargs = 2;
	char argNulls[5] = { ' ', ' ', 'n', 'n', 'n' };
	bool isNullIgnore = false;

	bool isFirst = true;

	/* todo: add spec validation for validator, like unsupported keywrods, etc. */
	if (validator != NULL && validator->value_type != BSON_TYPE_EOD)
	{
		appendStringInfo(query, "validator = $%d ", ++nargs);
		argNulls[nargs - 1] = ' ';
		argsTypes[nargs - 1] = BsonTypeId();
		argValues[nargs - 1] = PointerGetDatum(PgbsonInitFromDocumentBsonValue(
												   validator));
		isFirst = false;
	}
	if (!isFirst && validationLevel != NULL)
	{
		appendStringInfo(query, ", ");
	}
	if (validationLevel != NULL)
	{
		appendStringInfo(query, "validation_level = $%d ", ++nargs);
		argNulls[nargs - 1] = ' ';
		argsTypes[nargs - 1] = TEXTOID;
		argValues[nargs - 1] = CStringGetTextDatum(validationLevel);
		isFirst = false;
	}
	if (!isFirst && validationAction != NULL)
	{
		appendStringInfo(query, ", ");
	}
	if (validationAction != NULL)
	{
		appendStringInfo(query, "validation_action = $%d ", ++nargs);
		argNulls[nargs - 1] = ' ';
		argsTypes[nargs - 1] = TEXTOID;
		argValues[nargs - 1] = CStringGetTextDatum(validationAction);
	}

	appendStringInfo(query, "WHERE database_name = $1 AND collection_name = $2");

	RunQueryWithCommutativeWrites(query->data, nargs, argsTypes, argValues, argNulls,
								  SPI_OK_UPDATE, &isNullIgnore);
}


static void
CheckSyntaxForValidator(const bson_iter_t *validator)
{
	if (bson_iter_type(validator) != BSON_TYPE_DOCUMENT)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_TYPEMISMATCH),
						errmsg("$jsonSchema must be of object type")));
	}

	bson_iter_t schemaIter;
	if (bson_iter_recurse(validator, &schemaIter))
	{
		SchemaTreeState localTreeState = { };
		BuildSchemaTree(&localTreeState, &schemaIter);
	}
}


/*
 * This function processes the $jsonSchema in the bson iter.
 * It checks if the $jsonSchema is present and calls CheckSyntaxForValidator
 * to validate the syntax of the validator.
 */
static void
ProcessJsonschemaInIter(bson_iter_t *iter)
{
	CHECK_FOR_INTERRUPTS();
	check_stack_depth();

	bson_iter_t childIter, arrayIter;

	if (bson_iter_recurse(iter, &childIter))
	{
		while (bson_iter_next(&childIter))
		{
			if (strcmp(bson_iter_key(&childIter), "$jsonSchema") == 0)
			{
				CheckSyntaxForValidator(&childIter);
				return;
			}
			else if (BSON_ITER_HOLDS_ARRAY(&childIter) && bson_iter_recurse(&childIter,
																			&arrayIter))
			{
				while (bson_iter_next(&arrayIter))
				{
					if (BSON_ITER_HOLDS_DOCUMENT(&arrayIter))
					{
						bson_iter_t subIter;
						if (bson_iter_recurse(&arrayIter, &subIter) &&
							bson_iter_find(&subIter, "$jsonSchema"))
						{
							CheckSyntaxForValidator(&subIter);
							return;
						}
					}
				}
			}
		}
	}
}


/*
 * This function parses and checks the bson value for "validator" option
 * given in "create"/"collMod" command
 */
const bson_value_t *
ParseAndGetValidatorSpec(bson_iter_t *iter, const char *validatorName, bool *hasValue)
{
	if (!EnableSchemaValidation)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_COMMANDNOTSUPPORTED),
						errmsg("validator not supported yet")));
	}

	/* "validator" values of "null" and "undefined" are valid */
	if (BSON_ITER_HOLDS_UNDEFINED(iter) ||
		BSON_ITER_HOLDS_NULL(iter))
	{
		return NULL;
	}

	EnsureTopLevelFieldType(validatorName, iter, BSON_TYPE_DOCUMENT);

	/* Copy the bson value to make sure the validator is valid beyond the lifetime of iter */
	bson_value_t *validator = palloc(sizeof(bson_value_t));
	bson_value_copy(bson_iter_value(iter), validator);

	/* Large and overly complex validation rules can impact database performance, especially during write operations. */
	if (validator->value.v_doc.data_len > (uint32_t) MaxSchemaValidatorSize)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_BADVALUE),
						errmsg(
							"validator of size > %dKB is not supported. Contact Azure Support if you need to increase this limit.",
							MaxSchemaValidatorSize / 1024),
						errmsg(
							"validator of size > %dKB is not supported. Contact Azure Support if you need to increase this limit.",
							MaxSchemaValidatorSize / 1024)));
	}

	ProcessJsonschemaInIter(iter);

	*hasValue = true;
	return validator;
}


/*
 * This function parses and checks the bson value for "validationAction" option
 * given in "create"/"collMod" command
 */
char *
ParseAndGetValidationActionOption(bson_iter_t *iter, const char *validationActionName,
								  bool *hasValue)
{
	if (!EnableSchemaValidation)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_COMMANDNOTSUPPORTED),
						errmsg("validator not supported yet")));
	}

	/* "validationAction" values of "null" and "undefined" are valid. */
	if (BSON_ITER_HOLDS_UNDEFINED(iter) ||
		BSON_ITER_HOLDS_NULL(iter))
	{
		return NULL;
	}

	EnsureTopLevelFieldType(validationActionName, iter, BSON_TYPE_UTF8);
	const char *validationAction = bson_iter_utf8(iter, NULL);
	if (strcmp(validationAction, "warn") == 0 ||
		strcmp(validationAction, "error") == 0
		)
	{
		*hasValue = true;
		return pstrdup(validationAction);
	}
	else
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_BADVALUE),
						errmsg(
							"The enumeration value '%s' provided for the field '%s' is invalid.",
							validationAction, validationActionName),
						errdetail_log(
							"The enumeration value '%s' provided for the field '%s' is invalid.",
							validationAction, validationActionName)));
	}
}


/*
 * This function parses and checks the bson value for "validationLevel" option
 * given in "create"/"collMod" command
 */
char *
ParseAndGetValidationLevelOption(bson_iter_t *iter, const char *validationLevelName,
								 bool *hasValue)
{
	if (!EnableSchemaValidation)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_COMMANDNOTSUPPORTED),
						errmsg("validator not supported yet")));
	}

	/* "validationLevel" values of "null" and "undefined" are valid. */
	if (BSON_ITER_HOLDS_UNDEFINED(iter) ||
		BSON_ITER_HOLDS_NULL(iter))
	{
		return NULL;
	}

	EnsureTopLevelFieldType(validationLevelName, iter, BSON_TYPE_UTF8);
	const char *validationLevel = bson_iter_utf8(iter, NULL);
	if (strcmp(validationLevel, "off") == 0 ||
		strcmp(validationLevel, "strict") == 0 ||
		strcmp(validationLevel, "moderate") == 0
		)
	{
		*hasValue = true;
		return pstrdup(validationLevel);
	}
	else
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_BADVALUE),
						errmsg(
							"The enumeration value '%s' provided for the field '%s' is invalid.",
							validationLevel, validationLevelName),
						errdetail_log(
							"The enumeration value '%s' provided for the field '%s' is invalid.",
							validationLevel, validationLevelName)));
	}
}


/*
 * This utility function updates the `MongCollection` struct using the `collectionId` and `shardOid`.
 */
void
UpdateMongoCollectionUsingIds(MongoCollection *mongoCollection, uint64 collectionId, Oid
							  shardOid)
{
	const char *shardTableName = NULL;
	if (shardOid != InvalidOid && UseLocalExecutionShardQueries)
	{
		shardTableName = get_rel_name(shardOid);
	}

	mongoCollection->collectionId = collectionId;

	if (shardTableName != NULL)
	{
		strcpy(mongoCollection->shardTableName, shardTableName);
	}

	snprintf(mongoCollection->tableName, NAMEDATALEN, DOCUMENT_DATA_TABLE_NAME_FORMAT,
			 collectionId);
}


/*
 * CheckRelNameValidity verifies if a relation name follows the DocumentDB collection
 * naming convention and extracts the collection ID if valid.
 */
bool
CheckRelNameValidity(const char *relName, uint64_t *collectionId, bool validateShardTable)
{
	if (relName == NULL ||
		strncmp(relName, "documents_", 10) != 0)
	{
		return false;
	}

	/* We use strtoull since it returns the first character that didn't match
	 * We expect this to return the '_' character when it's a collection shard
	 * like ApiDataSchemaName.documents_1_111 and the parsed value will be 1.
	 * Alternatively, this will be \0 and the parsed value will be 1 if the
	 * table is documents_1 (parent table).
	 */
	char *numEndPointer = NULL;
	uint64 parsedCollectionId = strtoull(&relName[10], &numEndPointer, 10);

	if (!validateShardTable)
	{
		*collectionId = parsedCollectionId;
		return true;
	}

	if (IsShardTableForDocumentDbTable(relName, numEndPointer))
	{
		*collectionId = parsedCollectionId;
		return true;
	}

	return false;
}


bool
ParseChangeStreamPreAndPostImageOption(const bson_value_t *optionValue,
									   const char *commandPrefix)
{
	bson_iter_t csIter;
	BsonValueInitIterator(optionValue, &csIter);
	bool foundEnabled = false;
	bool enabled = false;
	while (bson_iter_next(&csIter))
	{
		const char *csKey = bson_iter_key(&csIter);
		if (strcmp(csKey, "enabled") == 0)
		{
			EnsureTopLevelFieldIsBooleanLike(
				psprintf("%s.enabled", commandPrefix), &csIter);
			enabled = BsonValueAsBool(bson_iter_value(&csIter));
			foundEnabled = true;
		}
		else
		{
			ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_UNKNOWNBSONFIELD),
							errmsg("The BSON field '%s.%s'"
								   " is not recognized.",
								   commandPrefix, csKey)));
		}
	}
	if (!foundEnabled)
	{
		ereport(ERROR, (errcode(ERRCODE_DOCUMENTDB_INVALIDOPTIONS),
						errmsg(
							"The 'enabled' field is required"
							" in changeStreamPreAndPostImages")));
	}

	return enabled;
}


/*
 * UpdateChangeStreamPreAndPostImages checks the data table's current replica
 * identity from pg_class syscache and ALTERs it to FULL (when enabled) or
 * DEFAULT (when disabled) only if the current state differs. REPLICA IDENTITY
 * FULL causes PostgreSQL to include the full old row in WAL for UPDATE/DELETE,
 * which is needed for generating pre-images in change streams.
 */
void
UpdateChangeStreamPreAndPostImages(MongoCollection *collection,
								   bool enabled)
{
	if (collection == NULL)
	{
		return;
	}

	/* Resolve the data table's OID from schema + table name */
	Oid nspOid = get_namespace_oid(ApiDataSchemaName, false);
	Oid relOid = get_relname_relid(collection->tableName, nspOid);
	if (!OidIsValid(relOid))
	{
		return;
	}

	/* Look up the current replica identity from pg_class syscache */
	HeapTuple classTuple = SearchSysCache1(RELOID, ObjectIdGetDatum(relOid));
	if (!HeapTupleIsValid(classTuple))
	{
		return;
	}

	Form_pg_class classForm = (Form_pg_class) GETSTRUCT(classTuple);
	char currentReplIdent = classForm->relreplident;
	ReleaseSysCache(classTuple);

	bool isCurrentlyFull = (currentReplIdent == REPLICA_IDENTITY_FULL);

	/* Only ALTER if the desired state differs from the current state */
	if (enabled == isCurrentlyFull)
	{
		return;
	}

	StringInfo command = makeStringInfo();
	appendStringInfo(command,
					 "ALTER TABLE %s.%s REPLICA IDENTITY %s",
					 ApiDataSchemaName,
					 collection->tableName,
					 enabled ? "FULL" : "DEFAULT");

	bool readOnly = false;
	bool resultIsNull = false;
	ExtensionExecuteQueryViaSPI(command->data, readOnly,
								SPI_OK_UTILITY, &resultIsNull);

	pfree(command->data);

	pgbson_writer changeStreamOptions;
	PgbsonWriterInit(&changeStreamOptions);
	if (enabled)
	{
		PgbsonWriterAppendBool(&changeStreamOptions,
							   "changeStreamPreAndPostImages.enabled",
							   strlen("changeStreamPreAndPostImages.enabled"),
							   enabled);
	}
	else
	{
		/* Create the bson_dollar_unset spec to remove the option. */
		PgbsonWriterAppendUtf8(&changeStreamOptions,
							   "", 0, "changeStreamPreAndPostImages");
	}
	pgbson *changeStreamOptionsBson = PgbsonWriterGetPgbson(&changeStreamOptions);

	if (enabled)
	{
		UpdateCollectionOptions(collection, changeStreamOptionsBson);
	}
	else
	{
		RemoveCollectionOptions(collection, changeStreamOptionsBson);
	}
}


/*
 * UpdateEnableUpdateDescription stores or removes the enableUpdateDescription
 * option in the collection's options column. Unlike changeStreamPreAndPostImages,
 * this is a plain boolean value (not a nested document).
 */
void
UpdateEnableUpdateDescription(MongoCollection *collection, bool enabled)
{
	if (collection == NULL)
	{
		return;
	}

	pgbson_writer optionWriter;
	PgbsonWriterInit(&optionWriter);
	if (enabled)
	{
		PgbsonWriterAppendBool(&optionWriter,
							   "enableUpdateDescription",
							   strlen("enableUpdateDescription"),
							   enabled);
	}
	else
	{
		/* Create the bson_dollar_unset spec to remove the option. */
		PgbsonWriterAppendUtf8(&optionWriter,
							   "", 0, "enableUpdateDescription");
	}
	pgbson *optionBson = PgbsonWriterGetPgbson(&optionWriter);

	if (enabled)
	{
		UpdateCollectionOptions(collection, optionBson);
	}
	else
	{
		RemoveCollectionOptions(collection, optionBson);
	}
}


/*
 * UpdateCollectionStatsEnabledOption stores or removes the statsEnabled
 * option in the collection's options column. This is used to track whether
 * planner statistics are enabled for the collection without requiring a
 * FOR UPDATE lock on collection_indexes.
 */
void
UpdateCollectionStatsEnabledOption(uint64 collectionId, bool enabled)
{
	MongoCollection *collection = GetMongoCollectionByColId(collectionId, NoLock);
	if (collection == NULL)
	{
		return;
	}

	pgbson_writer optionWriter;
	PgbsonWriterInit(&optionWriter);

	if (enabled)
	{
		PgbsonWriterAppendBool(&optionWriter, "statsEnabled",
							   strlen("statsEnabled"), true);
		pgbson *updateSpec = PgbsonWriterGetPgbson(&optionWriter);
		UpdateCollectionOptions(collection, updateSpec);
	}
	else
	{
		PgbsonWriterAppendUtf8(&optionWriter, "", 0, "statsEnabled");
		pgbson *removeSpec = PgbsonWriterGetPgbson(&optionWriter);
		RemoveCollectionOptions(collection, removeSpec);
	}
}


static void
UpdateCollectionOptions(MongoCollection *collection, const pgbson *updateSpec)
{
	bool isNull = false;
	bool readOnly = false;
	int numArgs = 2;
	char isNulls[2] = { ' ', ' ' };
	Datum datum[2];
	datum[0] = PointerGetDatum(updateSpec);
	datum[1] = PointerGetDatum(PgbsonInitEmpty());

	Oid argOid[2] = { BsonTypeId(), BsonTypeId() };

	StringInfo updateOptionQuery = makeStringInfo();

	/*
	 * We use CTE to perform bson_dollar_set only once and then check if
	 * the resulting bson is empty (length = 5) to decide whether to set options
	 * to NULL or the updated bson with the specified field unset.
	 */
	appendStringInfo(updateOptionQuery,
					 "WITH updated AS ( "
					 "  SELECT %s.bson_dollar_set( "
					 "    COALESCE(options, $2::%s.bson), "
					 "    $1::%s.bson) AS val "
					 "  FROM %s.collections "
					 "  WHERE collection_id = " UINT64_FORMAT
					 " ) UPDATE %s.collections SET options = "
					 "  CASE WHEN updated.val OPERATOR(%s.=) $2::%s.bson "
					 "    THEN NULL ELSE updated.val END "
					 " FROM updated "
					 " WHERE collection_id = "
					 UINT64_FORMAT,
					 ApiCatalogSchemaName, CoreSchemaName,
					 CoreSchemaName,
					 ApiCatalogSchemaName, collection->collectionId,
					 ApiCatalogSchemaName, CoreSchemaName,
					 CoreSchemaName, collection->collectionId);

	ExtensionExecuteQueryWithArgsViaSPI(updateOptionQuery->data, numArgs,
										argOid, datum, isNulls, readOnly,
										SPI_OK_UPDATE, &isNull);
}


/*
 * Utility function to remove a single or multiple options from collection options based on the path
 * provided in removeArraySpec.
 * `removeArraySpec` can accept valid specification as array of paths or a single path.
 */
static void
RemoveCollectionOptions(MongoCollection *collection, const pgbson *removeArraySpec)
{
	bool isNull = false;
	bool readOnly = false;
	int numArgs = 2;
	char isNulls[2] = { ' ', ' ' };
	Datum datum[2];
	datum[0] = PointerGetDatum(removeArraySpec);
	datum[1] = PointerGetDatum(PgbsonInitEmpty());

	Oid argOid[2] = { BsonTypeId(), BsonTypeId() };

	StringInfo updateOptionQuery = makeStringInfo();

	/*
	 * We use CTE to perform bson_dollar_unset only once and then check if
	 * the resulting bson is empty (length = 5) to decide whether to set options
	 * to NULL or the updated bson with the specified field unset.
	 */
	appendStringInfo(updateOptionQuery,
					 "WITH updated AS ( "
					 "  SELECT %s.bson_dollar_unset( "
					 "    COALESCE(options, $2::%s.bson), "
					 "    $1::%s.bson) AS val "
					 "  FROM %s.collections "
					 "  WHERE collection_id = " UINT64_FORMAT
					 " ) UPDATE %s.collections SET options = "
					 "  CASE WHEN updated.val OPERATOR(%s.=) $2::%s.bson "
					 "    THEN NULL ELSE updated.val END "
					 " FROM updated "
					 " WHERE collection_id = "
					 UINT64_FORMAT,
					 ApiCatalogSchemaName, CoreSchemaName,
					 CoreSchemaName,
					 ApiCatalogSchemaName, collection->collectionId,
					 ApiCatalogSchemaName, CoreSchemaName,
					 CoreSchemaName, collection->collectionId);

	ExtensionExecuteQueryWithArgsViaSPI(updateOptionQuery->data, numArgs,
										argOid, datum, isNulls, readOnly,
										SPI_OK_UPDATE, &isNull);
}


/*
 * Copies all the required options from collection to in memory MongoCollection struct.
 *
 * See also sql/collection_cache_validity_test.sql to add options for cache validation
 */
static void
CopyCollectionOptions(const pgbson *options, MongoCollection *collection)
{
	if (options == NULL)
	{
		return;
	}

	bson_iter_t optionsIter;
	PgbsonInitIterator(options, &optionsIter);
	while (bson_iter_next(&optionsIter))
	{
		const char *key = bson_iter_key(&optionsIter);
		if (strcmp(key, "changeStreamPreAndPostImages") == 0)
		{
			/*
			 * top level field is only present when the option is enabled,
			 * so we can infer that pre and post images are enabled.
			 */
			collection->options.changeStreamPreAndPostImagesEnabled = true;
		}
		else if (strcmp(key, "enableUpdateDescription") == 0)
		{
			/*
			 * enableUpdateDescription is a plain boolean in the options column.
			 * Its presence with value true means update descriptions are enabled
			 * for this collection.
			 */
			collection->options.updateDescriptionEnabled =
				BsonValueAsBool(bson_iter_value(&optionsIter));
		}
		else if (strcmp(key, "statsEnabled") == 0)
		{
			collection->options.statsEnabled =
				BsonValueAsBool(bson_iter_value(&optionsIter));
		}
	}
}
