From: William Grant <william.grant@canonical.com>
Date: Wed, 12 Sep 2012 09:59:31 +0000 (+1000)
Subject: Add DISABLE/ENABLE commands to prevent new connections.
X-Git-Tag: pgbouncer_1_6_rc1~50^2~3
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=1858855a780f82f3c3d8dc9e4e3bdb1e8e02d515;p=pgbouncer

Add DISABLE/ENABLE commands to prevent new connections.

DISABLE causes new connections to a particular database to be rejected
while existing ones continue. ENABLE allows new connections again.
---

diff --git a/NEWS b/NEWS
index 859ba8c..ae0f986 100644
--- a/NEWS
+++ b/NEWS
@@ -1,6 +1,10 @@
 2012-07-xx  -  PgBouncer 1.5.3  -  ""
 
-  = Minor Feature =
+  = Features =
+
+    * ENABLE/DISABLE commands to forbid or allow access to one database.
+
+  = Minor Features =
 
     * max_packet_size - config parameter to tune maximum packet size
       that is allowed through.  Default is kept same: (2G-1), but now
diff --git a/include/bouncer.h b/include/bouncer.h
index 687126a..5dc47eb 100644
--- a/include/bouncer.h
+++ b/include/bouncer.h
@@ -253,6 +253,7 @@ struct PgDatabase {
 	bool db_paused;		/* PAUSE <db>; was issued */
 	bool db_dead;		/* used on RELOAD/SIGHUP to later detect removed dbs */
 	bool db_auto;		/* is the database auto-created by autodb_connstr */
+	bool db_disabled;	/* is the database accepting new connections? */
 	bool admin;		/* internal console db */
 
 	struct PktBuf *startup_params; /* partial StartupMessage (without user) be sent to server */
diff --git a/src/admin.c b/src/admin.c
index 7bf9e80..9b2c7a6 100644
--- a/src/admin.c
+++ b/src/admin.c
@@ -980,6 +980,56 @@ static bool admin_cmd_pause(PgSocket *admin, const char *arg)
 	return true;
 }
 
+/* Command: DISABLE */
+static bool admin_cmd_disable(PgSocket *admin, const char *arg)
+{
+	PgDatabase *db;
+
+	if (!admin->admin_user)
+		return admin_error(admin, "admin access needed");
+
+	if (!arg[0])
+		return admin_error(admin, "a database is required");
+
+	log_info("DISABLE '%s' command issued", arg);
+	db = find_database(arg);
+	if (db == NULL) {
+		db = register_auto_database(arg);
+		if (db == NULL) {
+			return admin_error(admin, "no such database: %s", arg);
+		} else {
+			slog_info(admin, "registered new auto-database for DISABLE: %s", arg);
+		}
+	}
+	if (db == admin->pool->db)
+		return admin_error(admin, "cannot disable admin db: %s", arg);
+
+	db->db_disabled = 1;
+	return admin_ready(admin, "DISABLE");
+}
+
+/* Command: ENABLE */
+static bool admin_cmd_enable(PgSocket *admin, const char *arg)
+{
+	PgDatabase *db;
+
+	if (!admin->admin_user)
+		return admin_error(admin, "admin access needed");
+
+	if (!arg[0])
+		return admin_error(admin, "a database is required");
+
+	log_info("ENABLE '%s' command issued", arg);
+	db = find_database(arg);
+	if (db == NULL)
+		return admin_error(admin, "no such database: %s", arg);
+	if (db == admin->pool->db)
+		return admin_error(admin, "cannot disable admin db: %s", arg);
+
+	db->db_disabled = 0;
+	return admin_ready(admin, "ENABLE");
+}
+
 /* Command: KILL */
 static bool admin_cmd_kill(PgSocket *admin, const char *arg)
 {
@@ -1133,6 +1183,8 @@ static bool admin_cmd_show(PgSocket *admin, const char *arg)
 }
 
 static struct cmd_lookup cmd_list [] = {
+	{"disable", admin_cmd_disable},
+	{"enable", admin_cmd_enable},
 	{"kill", admin_cmd_kill},
 	{"pause", admin_cmd_pause},
 	{"reload", admin_cmd_reload},
diff --git a/src/client.c b/src/client.c
index 056fb96..c28bd7f 100644
--- a/src/client.c
+++ b/src/client.c
@@ -105,6 +105,10 @@ bool set_pool(PgSocket *client, const char *dbname, const char *username)
 		disconnect_client(client, true, "no memory for pool");
 		return false;
 	}
+	if (client->pool->db->db_disabled) {
+		disconnect_client(client, true, "pgbouncer database is disabled");
+		return false;
+	}
 
 	return check_fast_fail(client);
 }