Have the client double-check that it is connected as a database superuser before plowing through initialization steps.

schmiddy · schmiddy · commit fbcd24b4e09f · 2012-11-30T20:09:22.000-07:00
Addresses Issue #2.
diff --git a/bin/pg_repack.c b/bin/pg_repack.c
@@ -82,6 +82,7 @@ typedef struct repack_index
 	const char	   *create_index;	/* CREATE INDEX */
 } repack_index;
 
+static bool is_superuser(void);
 static void repack_all_databases(const char *order_by);
 static bool repack_one_database(const char *order_by, const char *table, char *errbuf, size_t errsize);
 static void repack_one_table(const repack_table *table, const char *order_by);
@@ -162,6 +163,31 @@ main(int argc, char *argv[])
 	return 0;
 }
 
+
+/*
+ * Test if the current user is a database superuser.
+ * Borrowed from psql/common.c
+ *
+ * Note: this will correctly detect superuserness only with a protocol-3.0
+ * or newer backend; otherwise it will always say "false".
+ */
+bool
+is_superuser(void)
+{
+	const char *val;
+
+	if (!connection)
+		return false;
+
+	val = PQparameterStatus(connection, "is_superuser");
+	
+	if (val && strcmp(val, "on") == 0)
+		return true;
+
+	return false;
+}
+
+
 /*
  * Call repack_one_database for each database.
  */
@@ -173,6 +199,10 @@ repack_all_databases(const char *orderby)
 
 	dbname = "postgres";
 	reconnect(ERROR);
+
+	if (!is_superuser())
+		elog(ERROR, "You must be a superuser to use %s", PROGRAM_NAME);
+
 	result = execute("SELECT datname FROM pg_database WHERE datallowconn ORDER BY 1;", 0, NULL);
 	disconnect();
 
@@ -230,7 +260,7 @@ static bool
 repack_one_database(const char *orderby, const char *table, char *errbuf, size_t errsize)
 {
 	bool			ret = false;
-	PGresult	   *res;
+	PGresult	   *res = NULL;
 	int				i;
 	int				num;
 	StringInfoData	sql;
@@ -239,6 +269,13 @@ repack_one_database(const char *orderby, const char *table, char *errbuf, size_t
 
 	reconnect(ERROR);
 
+	if (!is_superuser()) {
+		if (errbuf)
+			snprintf(errbuf, errsize, "You must be a superuser to use %s",
+					 PROGRAM_NAME);
+		goto cleanup;
+	}
+
 	/* Query the extension version. Exit if no match */
 	res = execute_elevel("select repack.version(), repack.version_sql()",
 		0, NULL, DEBUG2);
@@ -400,7 +437,8 @@ repack_one_database(const char *orderby, const char *table, char *errbuf, size_t
 	ret = true;
 
 cleanup:
-	PQclear(res);
+	if (res)
+		PQclear(res);
 	disconnect();
 	termStringInfo(&sql);
 	return ret;
