From be6b61b5b6e57e4f96fd5e710b507fb911f1d2cf Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 4 Aug 2017 10:44:12 +0200
Subject: [PATCH 2/4] Extend lookup routines for FDW and foreign server with
 NULL handling

The cache lookup routines for foreign-data wrappers and foreign servers
are extended with an extra argument able to control if an error or a
NULL object is returned to the caller in the event of an undefined
object.
---
 contrib/dblink/dblink.c             |  2 +-
 contrib/file_fdw/file_fdw.c         |  4 ++--
 contrib/postgres_fdw/connection.c   |  4 ++--
 contrib/postgres_fdw/postgres_fdw.c |  8 ++++----
 doc/src/sgml/fdwhandler.sgml        | 10 ++++++++--
 src/backend/catalog/objectaddress.c | 12 ++++++------
 src/backend/commands/foreigncmds.c  | 14 ++++++++------
 src/backend/commands/tablecmds.c    |  8 ++++----
 src/backend/foreign/foreign.c       | 24 ++++++++++++++++++++----
 src/include/foreign/foreign.h       |  4 ++--
 10 files changed, 57 insertions(+), 33 deletions(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 81136b131c..b5bf5c47d5 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -2775,7 +2775,7 @@ get_connect_string(const char *servername)
 		Oid			userid = GetUserId();
 
 		user_mapping = GetUserMapping(userid, serverid);
-		fdw = GetForeignDataWrapper(fdwid);
+		fdw = GetForeignDataWrapper(fdwid, false);
 
 		/* Check permissions, user must have usage on the server. */
 		aclresult = pg_foreign_server_aclcheck(serverid, userid, ACL_USAGE);
diff --git a/contrib/file_fdw/file_fdw.c b/contrib/file_fdw/file_fdw.c
index 2396bd442f..1d18a17417 100644
--- a/contrib/file_fdw/file_fdw.c
+++ b/contrib/file_fdw/file_fdw.c
@@ -357,8 +357,8 @@ fileGetOptions(Oid foreigntableid,
 	 * Simplify?)
 	 */
 	table = GetForeignTable(foreigntableid);
-	server = GetForeignServer(table->serverid);
-	wrapper = GetForeignDataWrapper(server->fdwid);
+	server = GetForeignServer(table->serverid, false);
+	wrapper = GetForeignDataWrapper(server->fdwid, false);
 
 	options = NIL;
 	options = list_concat(options, wrapper->options);
diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index be4ec07cf9..65f7806484 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -182,7 +182,7 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	 */
 	if (entry->conn == NULL)
 	{
-		ForeignServer *server = GetForeignServer(user->serverid);
+		ForeignServer *server = GetForeignServer(user->serverid, false);
 
 		/* Reset all transient state fields, to be sure all are clean */
 		entry->xact_depth = 0;
@@ -1003,7 +1003,7 @@ pgfdw_reject_incomplete_xact_state_change(ConnCacheEntry *entry)
 	if (!HeapTupleIsValid(tup))
 		elog(ERROR, "cache lookup failed for user mapping %u", entry->key);
 	umform = (Form_pg_user_mapping) GETSTRUCT(tup);
-	server = GetForeignServer(umform->umserver);
+	server = GetForeignServer(umform->umserver, false);
 	ReleaseSysCache(tup);
 
 	ereport(ERROR,
diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index e036704b6d..6800842575 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -506,7 +506,7 @@ postgresGetForeignRelSize(PlannerInfo *root,
 
 	/* Look up foreign-table catalog info. */
 	fpinfo->table = GetForeignTable(foreigntableid);
-	fpinfo->server = GetForeignServer(fpinfo->table->serverid);
+	fpinfo->server = GetForeignServer(fpinfo->table->serverid, false);
 
 	/*
 	 * Extract user-settable option values.  Note that per-table setting of
@@ -2038,7 +2038,7 @@ postgresIsForeignRelUpdatable(Relation rel)
 	updatable = true;
 
 	table = GetForeignTable(RelationGetRelid(rel));
-	server = GetForeignServer(table->serverid);
+	server = GetForeignServer(table->serverid, false);
 
 	foreach(lc, server->options)
 	{
@@ -3588,7 +3588,7 @@ postgresAcquireSampleRowsFunc(Relation relation, int elevel,
 	 * owner, even if the ANALYZE was started by some other user.
 	 */
 	table = GetForeignTable(RelationGetRelid(relation));
-	server = GetForeignServer(table->serverid);
+	server = GetForeignServer(table->serverid, false);
 	user = GetUserMapping(relation->rd_rel->relowner, table->serverid);
 	conn = GetConnection(user, false);
 
@@ -3811,7 +3811,7 @@ postgresImportForeignSchema(ImportForeignSchemaStmt *stmt, Oid serverOid)
 	 * Get connection to the foreign server.  Connection manager will
 	 * establish new connection if necessary.
 	 */
-	server = GetForeignServer(serverOid);
+	server = GetForeignServer(serverOid, false);
 	mapping = GetUserMapping(GetUserId(), server->serverid);
 	conn = GetConnection(mapping, false);
 
diff --git a/doc/src/sgml/fdwhandler.sgml b/doc/src/sgml/fdwhandler.sgml
index dbeaab555d..4e493d1e23 100644
--- a/doc/src/sgml/fdwhandler.sgml
+++ b/doc/src/sgml/fdwhandler.sgml
@@ -1288,25 +1288,31 @@ ShutdownForeignScan(ForeignScanState *node);
     <para>
 <programlisting>
 ForeignDataWrapper *
-GetForeignDataWrapper(Oid fdwid);
+GetForeignDataWrapper(Oid fdwid, bool missing_ok);
 </programlisting>
 
      This function returns a <structname>ForeignDataWrapper</structname>
      object for the foreign-data wrapper with the given OID.  A
      <structname>ForeignDataWrapper</structname> object contains properties
      of the FDW (see <filename>foreign/foreign.h</filename> for details).
+     If <literal>missing_ok</> is true, a <literal>NULL</> result is
+     returned to the caller instead of an error for an undefined
+     foreign-data wrapper.
     </para>
 
     <para>
 <programlisting>
 ForeignServer *
-GetForeignServer(Oid serverid);
+GetForeignServer(Oid serverid, bool missing_ok);
 </programlisting>
 
      This function returns a <structname>ForeignServer</structname> object
      for the foreign server with the given OID.  A
      <structname>ForeignServer</structname> object contains properties
      of the server (see <filename>foreign/foreign.h</filename> for details).
+     If <literal>missing_ok</> is true, a <literal>NULL</> result is
+     returned to the caller instead of an error for an undefined foreign
+     server.
     </para>
 
     <para>
diff --git a/src/backend/catalog/objectaddress.c b/src/backend/catalog/objectaddress.c
index ee1ccd0d12..98fa02aa02 100644
--- a/src/backend/catalog/objectaddress.c
+++ b/src/backend/catalog/objectaddress.c
@@ -3186,7 +3186,7 @@ getObjectDescription(const ObjectAddress *object)
 			{
 				ForeignDataWrapper *fdw;
 
-				fdw = GetForeignDataWrapper(object->objectId);
+				fdw = GetForeignDataWrapper(object->objectId, false);
 				appendStringInfo(&buffer, _("foreign-data wrapper %s"), fdw->fdwname);
 				break;
 			}
@@ -3195,7 +3195,7 @@ getObjectDescription(const ObjectAddress *object)
 			{
 				ForeignServer *srv;
 
-				srv = GetForeignServer(object->objectId);
+				srv = GetForeignServer(object->objectId, false);
 				appendStringInfo(&buffer, _("server %s"), srv->servername);
 				break;
 			}
@@ -3215,7 +3215,7 @@ getObjectDescription(const ObjectAddress *object)
 						 object->objectId);
 				umform = (Form_pg_user_mapping) GETSTRUCT(tup);
 				useid = umform->umuser;
-				srv = GetForeignServer(umform->umserver);
+				srv = GetForeignServer(umform->umserver, false);
 
 				ReleaseSysCache(tup);
 
@@ -4690,7 +4690,7 @@ getObjectIdentityParts(const ObjectAddress *object,
 			{
 				ForeignDataWrapper *fdw;
 
-				fdw = GetForeignDataWrapper(object->objectId);
+				fdw = GetForeignDataWrapper(object->objectId, false);
 				appendStringInfoString(&buffer, quote_identifier(fdw->fdwname));
 				if (objname)
 					*objname = list_make1(pstrdup(fdw->fdwname));
@@ -4701,7 +4701,7 @@ getObjectIdentityParts(const ObjectAddress *object,
 			{
 				ForeignServer *srv;
 
-				srv = GetForeignServer(object->objectId);
+				srv = GetForeignServer(object->objectId, false);
 				appendStringInfoString(&buffer,
 									   quote_identifier(srv->servername));
 				if (objname)
@@ -4724,7 +4724,7 @@ getObjectIdentityParts(const ObjectAddress *object,
 						 object->objectId);
 				umform = (Form_pg_user_mapping) GETSTRUCT(tup);
 				useid = umform->umuser;
-				srv = GetForeignServer(umform->umserver);
+				srv = GetForeignServer(umform->umserver, false);
 
 				ReleaseSysCache(tup);
 
diff --git a/src/backend/commands/foreigncmds.c b/src/backend/commands/foreigncmds.c
index 9ad991507f..33f40ee98f 100644
--- a/src/backend/commands/foreigncmds.c
+++ b/src/backend/commands/foreigncmds.c
@@ -368,7 +368,8 @@ AlterForeignServerOwner_internal(Relation rel, HeapTuple tup, Oid newOwnerId)
 			aclresult = pg_foreign_data_wrapper_aclcheck(form->srvfdw, newOwnerId, ACL_USAGE);
 			if (aclresult != ACLCHECK_OK)
 			{
-				ForeignDataWrapper *fdw = GetForeignDataWrapper(form->srvfdw);
+				ForeignDataWrapper *fdw = GetForeignDataWrapper(form->srvfdw,
+																false);
 
 				aclcheck_error(aclresult, ACL_KIND_FDW, fdw->fdwname);
 			}
@@ -1033,7 +1034,8 @@ AlterForeignServer(AlterForeignServerStmt *stmt)
 
 	if (stmt->options)
 	{
-		ForeignDataWrapper *fdw = GetForeignDataWrapper(srvForm->srvfdw);
+		ForeignDataWrapper *fdw = GetForeignDataWrapper(srvForm->srvfdw,
+														false);
 		Datum		datum;
 		bool		isnull;
 
@@ -1187,7 +1189,7 @@ CreateUserMapping(CreateUserMappingStmt *stmt)
 							stmt->servername)));
 	}
 
-	fdw = GetForeignDataWrapper(srv->fdwid);
+	fdw = GetForeignDataWrapper(srv->fdwid, false);
 
 	/*
 	 * Insert tuple into pg_user_mapping.
@@ -1299,7 +1301,7 @@ AlterUserMapping(AlterUserMappingStmt *stmt)
 		 * Process the options.
 		 */
 
-		fdw = GetForeignDataWrapper(srv->fdwid);
+		fdw = GetForeignDataWrapper(srv->fdwid, false);
 
 		datum = SysCacheGetAttr(USERMAPPINGUSERSERVER,
 								tp,
@@ -1479,7 +1481,7 @@ CreateForeignTable(CreateForeignTableStmt *stmt, Oid relid)
 	if (aclresult != ACLCHECK_OK)
 		aclcheck_error(aclresult, ACL_KIND_FOREIGN_SERVER, server->servername);
 
-	fdw = GetForeignDataWrapper(server->fdwid);
+	fdw = GetForeignDataWrapper(server->fdwid, false);
 
 	/*
 	 * Insert tuple into pg_foreign_table.
@@ -1542,7 +1544,7 @@ ImportForeignSchema(ImportForeignSchemaStmt *stmt)
 	(void) LookupCreationNamespace(stmt->local_schema);
 
 	/* Get the FDW and check it supports IMPORT */
-	fdw = GetForeignDataWrapper(server->fdwid);
+	fdw = GetForeignDataWrapper(server->fdwid, false);
 	if (!OidIsValid(fdw->fdwhandler))
 		ereport(ERROR,
 				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 7859ef13ac..1ecf3ebc0c 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -9391,8 +9391,8 @@ ATExecAlterColumnGenericOptions(Relation rel,
 				 errmsg("foreign table \"%s\" does not exist",
 						RelationGetRelationName(rel))));
 	fttableform = (Form_pg_foreign_table) GETSTRUCT(tuple);
-	server = GetForeignServer(fttableform->ftserver);
-	fdw = GetForeignDataWrapper(server->fdwid);
+	server = GetForeignServer(fttableform->ftserver, false);
+	fdw = GetForeignDataWrapper(server->fdwid, false);
 
 	heap_close(ftrel, AccessShareLock);
 	ReleaseSysCache(tuple);
@@ -12249,8 +12249,8 @@ ATExecGenericOptions(Relation rel, List *options)
 				 errmsg("foreign table \"%s\" does not exist",
 						RelationGetRelationName(rel))));
 	tableform = (Form_pg_foreign_table) GETSTRUCT(tuple);
-	server = GetForeignServer(tableform->ftserver);
-	fdw = GetForeignDataWrapper(server->fdwid);
+	server = GetForeignServer(tableform->ftserver, false);
+	fdw = GetForeignDataWrapper(server->fdwid, false);
 
 	memset(repl_val, 0, sizeof(repl_val));
 	memset(repl_null, false, sizeof(repl_null));
diff --git a/src/backend/foreign/foreign.c b/src/backend/foreign/foreign.c
index a113bf540d..bc73300aa6 100644
--- a/src/backend/foreign/foreign.c
+++ b/src/backend/foreign/foreign.c
@@ -32,7 +32,7 @@
  * GetForeignDataWrapper -	look up the foreign-data wrapper by OID.
  */
 ForeignDataWrapper *
-GetForeignDataWrapper(Oid fdwid)
+GetForeignDataWrapper(Oid fdwid, bool missing_ok)
 {
 	Form_pg_foreign_data_wrapper fdwform;
 	ForeignDataWrapper *fdw;
@@ -43,7 +43,11 @@ GetForeignDataWrapper(Oid fdwid)
 	tp = SearchSysCache1(FOREIGNDATAWRAPPEROID, ObjectIdGetDatum(fdwid));
 
 	if (!HeapTupleIsValid(tp))
+	{
+		if (missing_ok)
+			return NULL;
 		elog(ERROR, "cache lookup failed for foreign-data wrapper %u", fdwid);
+	}
 
 	fdwform = (Form_pg_foreign_data_wrapper) GETSTRUCT(tp);
 
@@ -82,7 +86,11 @@ GetForeignDataWrapperByName(const char *fdwname, bool missing_ok)
 	if (!OidIsValid(fdwId))
 		return NULL;
 
-	return GetForeignDataWrapper(fdwId);
+	/*
+	 * missing_ok set to true makes no sense here as a lookup has already
+	 * happened.
+	 */
+	return GetForeignDataWrapper(fdwId, false);
 }
 
 
@@ -90,7 +98,7 @@ GetForeignDataWrapperByName(const char *fdwname, bool missing_ok)
  * GetForeignServer - look up the foreign server definition.
  */
 ForeignServer *
-GetForeignServer(Oid serverid)
+GetForeignServer(Oid serverid, bool missing_ok)
 {
 	Form_pg_foreign_server serverform;
 	ForeignServer *server;
@@ -101,7 +109,11 @@ GetForeignServer(Oid serverid)
 	tp = SearchSysCache1(FOREIGNSERVEROID, ObjectIdGetDatum(serverid));
 
 	if (!HeapTupleIsValid(tp))
+	{
+		if (missing_ok)
+			return NULL;
 		elog(ERROR, "cache lookup failed for foreign server %u", serverid);
+	}
 
 	serverform = (Form_pg_foreign_server) GETSTRUCT(tp);
 
@@ -152,7 +164,11 @@ GetForeignServerByName(const char *srvname, bool missing_ok)
 	if (!OidIsValid(serverid))
 		return NULL;
 
-	return GetForeignServer(serverid);
+	/*
+	 * missing_ok set to true makes no sense here as a lookup has already
+	 * happened.
+	 */
+	return GetForeignServer(serverid, false);
 }
 
 
diff --git a/src/include/foreign/foreign.h b/src/include/foreign/foreign.h
index 2f4c569d1d..13d5fab983 100644
--- a/src/include/foreign/foreign.h
+++ b/src/include/foreign/foreign.h
@@ -69,10 +69,10 @@ typedef struct ForeignTable
 } ForeignTable;
 
 
-extern ForeignServer *GetForeignServer(Oid serverid);
+extern ForeignServer *GetForeignServer(Oid serverid, bool missing_ok);
 extern ForeignServer *GetForeignServerByName(const char *name, bool missing_ok);
 extern UserMapping *GetUserMapping(Oid userid, Oid serverid);
-extern ForeignDataWrapper *GetForeignDataWrapper(Oid fdwid);
+extern ForeignDataWrapper *GetForeignDataWrapper(Oid fdwid, bool missing_ok);
 extern ForeignDataWrapper *GetForeignDataWrapperByName(const char *name,
 							bool missing_ok);
 extern ForeignTable *GetForeignTable(Oid relid);
-- 
2.13.4

