Add new function BackgroundWorkerInitializeConnectionByOid.

robertmhaas · robertmhaas · commit 5d2f957f3f9d · 2015-02-02T16:23:59.000-05:00
Sometimes it's useful for a background worker to be able to initialize
its database connection by OID rather than by name, so provide a way
to do that.
diff --git a/doc/src/sgml/bgworker.sgml b/doc/src/sgml/bgworker.sgml
@@ -146,14 +146,17 @@ typedef struct BackgroundWorker
   </para>
 
   <para>Once running, the process can connect to a database by calling
-   <function>BackgroundWorkerInitializeConnection(<parameter>char *dbname</parameter>, <parameter>char *username</parameter>)</function>.
+   <function>BackgroundWorkerInitializeConnection(<parameter>char *dbname</parameter>, <parameter>char *username</parameter>)</function> or
+   <function>BackgroundWorkerInitializeConnectionByOid(<parameter>Oid dboid</parameter>, <parameter>Oid useroid</parameter>)</function>.
    This allows the process to run transactions and queries using the
-   <literal>SPI</literal> interface.  If <varname>dbname</> is NULL,
-   the session is not connected to any particular database, but shared catalogs
-   can be accessed.  If <varname>username</> is NULL, the process will run as
-   the superuser created during <command>initdb</>.
-   BackgroundWorkerInitializeConnection can only be called once per background
-   process, it is not possible to switch databases.
+   <literal>SPI</literal> interface.  If <varname>dbname</> is NULL or
+   <varname>dboid</> is <literal>InvalidOid</>, the session is not connected
+   to any particular database, but shared catalogs can be accessed.
+   If <varname>username</> is NULL or <varname>useroid</> is
+   <literal>InvalidOid</>, the process will run as the superuser created
+   during <command>initdb</>.
+   A background worker can only call one of these two functions, and only
+   once.  It is not possible to switch databases.
   </para>
 
   <para>
diff --git a/src/backend/bootstrap/bootstrap.c b/src/backend/bootstrap/bootstrap.c
@@ -467,7 +467,7 @@ BootstrapModeMain(void)
 	 */
 	InitProcess();
 
-	InitPostgres(NULL, InvalidOid, NULL, NULL);
+	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, NULL);
 
 	/* Initialize stuff for bootstrap-file processing */
 	for (i = 0; i < MAXATTR; i++)
diff --git a/src/backend/postmaster/autovacuum.c b/src/backend/postmaster/autovacuum.c
@@ -450,7 +450,7 @@ AutoVacLauncherMain(int argc, char *argv[])
 	InitProcess();
 #endif
 
-	InitPostgres(NULL, InvalidOid, NULL, NULL);
+	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, NULL);
 
 	SetProcessingMode(NormalProcessing);
 
@@ -1620,7 +1620,7 @@ AutoVacWorkerMain(int argc, char *argv[])
 		 * Note: if we have selected a just-deleted database (due to using
 		 * stale stats info), we'll fail and exit here.
 		 */
-		InitPostgres(NULL, dbid, NULL, dbname);
+		InitPostgres(NULL, dbid, NULL, InvalidOid, dbname);
 		SetProcessingMode(NormalProcessing);
 		set_ps_display(dbname, false);
 		ereport(DEBUG1,
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
@@ -5313,7 +5313,30 @@ BackgroundWorkerInitializeConnection(char *dbname, char *username)
 				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
 				 errmsg("database connection requirement not indicated during registration")));
 
-	InitPostgres(dbname, InvalidOid, username, NULL);
+	InitPostgres(dbname, InvalidOid, username, InvalidOid, NULL);
+
+	/* it had better not gotten out of "init" mode yet */
+	if (!IsInitProcessingMode())
+		ereport(ERROR,
+				(errmsg("invalid processing mode in background worker")));
+	SetProcessingMode(NormalProcessing);
+}
+
+/*
+ * Connect background worker to a database using OIDs.
+ */
+void
+BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid)
+{
+	BackgroundWorker *worker = MyBgworkerEntry;
+
+	/* XXX is this the right errcode? */
+	if (!(worker->bgw_flags & BGWORKER_BACKEND_DATABASE_CONNECTION))
+		ereport(FATAL,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("database connection requirement not indicated during registration")));
+
+	InitPostgres(NULL, dboid, NULL, useroid, NULL);
 
 	/* it had better not gotten out of "init" mode yet */
 	if (!IsInitProcessingMode())
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
@@ -3714,7 +3714,7 @@ PostgresMain(int argc, char *argv[],
 	 * it inside InitPostgres() instead.  In particular, anything that
 	 * involves database access should be there, not here.
 	 */
-	InitPostgres(dbname, InvalidOid, username, NULL);
+	InitPostgres(dbname, InvalidOid, username, InvalidOid, NULL);
 
 	/*
 	 * If the PostmasterContext is still around, recycle the space; we don't
diff --git a/src/backend/utils/init/miscinit.c b/src/backend/utils/init/miscinit.c
@@ -453,11 +453,10 @@ has_rolreplication(Oid roleid)
  * Initialize user identity during normal backend startup
  */
 void
-InitializeSessionUserId(const char *rolename)
+InitializeSessionUserId(const char *rolename, Oid roleid)
 {
 	HeapTuple	roleTup;
 	Form_pg_authid rform;
-	Oid			roleid;
 
 	/*
 	 * Don't do scans if we're bootstrapping, none of the system catalogs
@@ -468,7 +467,10 @@ InitializeSessionUserId(const char *rolename)
 	/* call only once */
 	AssertState(!OidIsValid(AuthenticatedUserId));
 
-	roleTup = SearchSysCache1(AUTHNAME, PointerGetDatum(rolename));
+	if (rolename != NULL)
+		roleTup = SearchSysCache1(AUTHNAME, PointerGetDatum(rolename));
+	else
+		roleTup = SearchSysCache1(AUTHOID, ObjectIdGetDatum(roleid));
 	if (!HeapTupleIsValid(roleTup))
 		ereport(FATAL,
 				(errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION),
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
@@ -523,6 +523,9 @@ BaseInit(void)
  * name can be returned to the caller in out_dbname.  If out_dbname isn't
  * NULL, it must point to a buffer of size NAMEDATALEN.
  *
+ * Similarly, the username can be passed by name, using the username parameter,
+ * or by OID using the useroid parameter.
+ *
  * In bootstrap mode no parameters are used.  The autovacuum launcher process
  * doesn't use any parameters either, because it only goes far enough to be
  * able to read pg_database; it doesn't connect to any particular database.
@@ -537,7 +540,7 @@ BaseInit(void)
  */
 void
 InitPostgres(const char *in_dbname, Oid dboid, const char *username,
-			 char *out_dbname)
+			 Oid useroid, char *out_dbname)
 {
 	bool		bootstrap = IsBootstrapProcessingMode();
 	bool		am_superuser;
@@ -692,18 +695,18 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 					(errcode(ERRCODE_UNDEFINED_OBJECT),
 					 errmsg("no roles are defined in this database system"),
 					 errhint("You should immediately run CREATE USER \"%s\" SUPERUSER;.",
-							 username)));
+							 username != NULL ? username : "postgres")));
 	}
 	else if (IsBackgroundWorker)
 	{
-		if (username == NULL)
+		if (username == NULL && !OidIsValid(useroid))
 		{
 			InitializeSessionUserIdStandalone();
 			am_superuser = true;
 		}
 		else
 		{
-			InitializeSessionUserId(username);
+			InitializeSessionUserId(username, useroid);
 			am_superuser = superuser();
 		}
 	}
@@ -712,7 +715,7 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 		/* normal multiuser case */
 		Assert(MyProcPort != NULL);
 		PerformAuthentication(MyProcPort);
-		InitializeSessionUserId(username);
+		InitializeSessionUserId(username, useroid);
 		am_superuser = superuser();
 	}
 
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
@@ -307,7 +307,7 @@ extern bool InLocalUserIdChange(void);
 extern bool InSecurityRestrictedOperation(void);
 extern void GetUserIdAndContext(Oid *userid, bool *sec_def_context);
 extern void SetUserIdAndContext(Oid userid, bool sec_def_context);
-extern void InitializeSessionUserId(const char *rolename);
+extern void InitializeSessionUserId(const char *rolename, Oid useroid);
 extern void InitializeSessionUserIdStandalone(void);
 extern void SetSessionAuthorization(Oid userid, bool is_superuser);
 extern Oid	GetCurrentRoleId(void);
@@ -411,7 +411,7 @@ extern AuxProcType MyAuxProcType;
 extern void pg_split_opts(char **argv, int *argcp, char *optstr);
 extern void InitializeMaxBackends(void);
 extern void InitPostgres(const char *in_dbname, Oid dboid, const char *username,
-			 char *out_dbname);
+			 Oid useroid, char *out_dbname);
 extern void BaseInit(void);
 
 /* in utils/init/miscinit.c */
diff --git a/src/include/postmaster/bgworker.h b/src/include/postmaster/bgworker.h
@@ -130,6 +130,9 @@ extern PGDLLIMPORT BackgroundWorker *MyBgworkerEntry;
  */
 extern void BackgroundWorkerInitializeConnection(char *dbname, char *username);
 
+/* Just like the above, but specifying database and user by OID. */
+extern void BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid);
+
 /* Block/unblock signals in a background worker process */
 extern void BackgroundWorkerBlockSignals(void);
 extern void BackgroundWorkerUnblockSignals(void);

Original file line number	Diff line number	Diff line change
`@@ -523,6 +523,9 @@ BaseInit(void)`
`523`	`523`	`* name can be returned to the caller in out_dbname. If out_dbname isn't`
`524`	`524`	`* NULL, it must point to a buffer of size NAMEDATALEN.`
`525`	`525`	`*`
	`526`	`+ * Similarly, the username can be passed by name, using the username parameter,`
	`527`	`+ * or by OID using the useroid parameter.`
	`528`	`+ *`
`526`	`529`	`* In bootstrap mode no parameters are used. The autovacuum launcher process`
`527`	`530`	`* doesn't use any parameters either, because it only goes far enough to be`
`528`	`531`	`* able to read pg_database; it doesn't connect to any particular database.`
`@@ -537,7 +540,7 @@ BaseInit(void)`
`537`	`540`	`*/`
`538`	`541`	`void`
`539`	`542`	`InitPostgres(const char in_dbname, Oid dboid, const char username,`
`540`		`- char *out_dbname)`
	`543`	`+ Oid useroid, char *out_dbname)`
`541`	`544`	`{`
`542`	`545`	`bool bootstrap = IsBootstrapProcessingMode();`
`543`	`546`	`bool am_superuser;`
`@@ -692,18 +695,18 @@ InitPostgres(const char in_dbname, Oid dboid, const char username,`
`692`	`695`	`(errcode(ERRCODE_UNDEFINED_OBJECT),`
`693`	`696`	`errmsg("no roles are defined in this database system"),`
`694`	`697`	`errhint("You should immediately run CREATE USER \"%s\" SUPERUSER;.",`
`695`		`- username)));`
	`698`	`+ username != NULL ? username : "postgres")));`
`696`	`699`	`}`
`697`	`700`	`else if (IsBackgroundWorker)`
`698`	`701`	`{`
`699`		`- if (username == NULL)`
	`702`	`+ if (username == NULL && !OidIsValid(useroid))`
`700`	`703`	`{`
`701`	`704`	`InitializeSessionUserIdStandalone();`
`702`	`705`	`am_superuser = true;`
`703`	`706`	`}`
`704`	`707`	`else`
`705`	`708`	`{`
`706`		`- InitializeSessionUserId(username);`
	`709`	`+ InitializeSessionUserId(username, useroid);`
`707`	`710`	`am_superuser = superuser();`
`708`	`711`	`}`
`709`	`712`	`}`
`@@ -712,7 +715,7 @@ InitPostgres(const char in_dbname, Oid dboid, const char username,`
`712`	`715`	`/* normal multiuser case */`
`713`	`716`	`Assert(MyProcPort != NULL);`
`714`	`717`	`PerformAuthentication(MyProcPort);`
`715`		`- InitializeSessionUserId(username);`
	`718`	`+ InitializeSessionUserId(username, useroid);`
`716`	`719`	`am_superuser = superuser();`
`717`	`720`	`}`
`718`	`721`