New database setting: connect_query

Specifies a query to be executed after connection is established.

Original patch by Teodor Sigaev, applied with minor modifications

diff --git a/doc/config.txt b/doc/config.txt
index 33a3d0f1d7f98acb9242060b4361c5c7d077b980..25c9a6a6a82bf9fd77a96273da3de4ce1513b281 100644 (file)
--- a/doc/config.txt
+++ b/doc/config.txt
@@ -349,13 +349,19 @@ database.
 Otherwise pgbouncer tries to log into destination database with client
 username, meaning that there will be one pool per user.
 
-=== Per-database pool Size ===
+=== Pool configuration ===
 
 ==== pool_size ====
 
 Set maximum size of pools for this database.  If not set,
 the default_pool_size is used.
 
+==== connect_query ====
+
+Query to be executed after connecttion is established, but before
+taking the connection into use by clients.  If the query raises errors,
+they are logged but ignored otherwise.
+
 === Extra parameters ===
 
 They allow setting default parameters on server connection.

diff --git a/doc/todo.txt b/doc/todo.txt
index 42a2f062fcbf91d714de99ff037b2205e5f066f3..3fc33b9850bf1253623a078171e382d1041fe1e7 100644 (file)
--- a/doc/todo.txt
+++ b/doc/todo.txt
@@ -20,6 +20,7 @@
  * SHOW TOTALS that shows stats summary (as goes to log) plus mem usage.
  * removing db from config and reload; works - kill connections
  * SHOW ACTIVE_SOCKETS - like show sockets; but filter only active ones
+ * Per-database connect_query (Teodor Sigaev).
 
 == Minor features ==
 

diff --git a/etc/pgbouncer.ini b/etc/pgbouncer.ini
index 69bb173c9e718492267418925cb62580087f3590..fef3bdcdec16a2d5a2813b6ca868910ae5be9853 100644 (file)
--- a/etc/pgbouncer.ini
+++ b/etc/pgbouncer.ini
@@ -8,7 +8,7 @@ foodb =
 bardb = host=127.0.0.1 dbname=bazdb
 
 ; acceess to dest database will go with single user
-forcedb = host=127.0.0.1 port=300 user=baz password=foo client_encoding=UNICODE datestyle=ISO
+forcedb = host=127.0.0.1 port=300 user=baz password=foo client_encoding=UNICODE datestyle=ISO connect_query='SELECT 1'
 
 ;; Configuation section
 [pgbouncer]

diff --git a/include/bouncer.h b/include/bouncer.h
index 42e7344269508a2628f66c7e582f587cbfc9a863..83cd1a5a1a2728ab4b9fb70e79c087916e74a0db 100644 (file)
--- a/include/bouncer.h
+++ b/include/bouncer.h
@@ -233,8 +233,12 @@ struct PgDatabase {
        int pool_size;          /* max server connections in one pool */
 
        const char *dbname;     /* server-side name, pointer to inside startup_msg */
+
+       /* startup commands to send to server after connect. malloc-ed */
+       const char *connect_query;
 };
 
+
 /*
  * A client or server connection.
  *
@@ -252,6 +256,7 @@ struct PgSocket {
        bool ready:1;           /* server: accepts new query */
        bool close_needed:1;    /* server: this socket must be closed ASAP */
        bool setting_vars:1;    /* server: setting client vars */
+       bool exec_on_connect:1; /* server: executing connect_query */
 
        bool wait_for_welcome:1;/* client: no server yet in pool, cannot send welcome msg */
 

diff --git a/src/janitor.c b/src/janitor.c
index 824c9a4cd32b01ac0920ff582aac04782fb9403d..ef43ef70a942f5202001e4d5223f7397cc3b4c34 100644 (file)
--- a/src/janitor.c
+++ b/src/janitor.c
@@ -544,6 +544,8 @@ static void kill_database(PgDatabase *db)
        }
        if (db->forced_user)
                obj_free(user_cache, db->forced_user);
+       if (db->connect_query)
+               free((void *)db->connect_query);
        statlist_remove(&db->head, &database_list);
        obj_free(db_cache, db);
 }

diff --git a/src/loader.c b/src/loader.c
index 6092bc009083e98aea1dc1abdd6f0e0fe45353d9..b0b58e6ba3d8cc1f4432c64d30ac1b2d4e18a4c7 100644 (file)
--- a/src/loader.c
+++ b/src/loader.c
@@ -127,6 +127,31 @@ static char * cstr_get_pair(char *p,
        return cstr_skip_ws(p);
 }
 
+static void set_connect_query(PgDatabase *db, const char *new)
+{
+       const char *old = db->connect_query;
+       char *val = NULL;
+
+       if (old && new) {
+               if (strcmp(old, new) == 0)
+                       return;
+               val = strdup(new);
+               if (val) {
+                       free((void *)old);
+                       db->connect_query = val;
+               }
+       } else if (new) {
+               val = strdup(new);
+               db->connect_query = val;
+       } else {
+               free((void *)db->connect_query);
+               db->connect_query = NULL;
+       }
+
+       if (new && !val)
+               log_error("no memory, cannot assign connect_query for %s", db->name);
+}
+
 /* fill PgDatabase from connstr */
 void parse_database(char *name, char *connstr)
 {
@@ -143,6 +168,7 @@ void parse_database(char *name, char *connstr)
        char *client_encoding = NULL;
        char *datestyle = NULL;
        char *timezone = NULL;
+       char *connect_query = NULL;
        char *unix_dir = "";
 
        in_addr_t v_addr = INADDR_NONE;
@@ -175,6 +201,8 @@ void parse_database(char *name, char *connstr)
                        timezone = val;
                else if (strcmp("pool_size", key) == 0)
                        pool_size = atoi(val);
+               else if (strcmp("connect_query", key) == 0)
+                       connect_query = val;
                else {
                        log_error("skipping database %s because"
                                  " of unknown parameter in connstring: %s", name, key);
@@ -258,6 +286,10 @@ void parse_database(char *name, char *connstr)
                        changed = true;
                else if (strcmp(db->unix_socket_dir, unix_dir) != 0)
                        changed = true;
+               else if ((db->connect_query && !connect_query)
+                        || (!db->connect_query && connect_query)
+                        || (connect_query && strcmp(connect_query, db->connect_query) != 0))
+                       changed = true;
 
                if (changed)
                        tag_database_dirty(db);
@@ -273,6 +305,9 @@ void parse_database(char *name, char *connstr)
        if (host)
                log_debug("%s: host=%s/%s", name, host, inet_ntoa(db->addr.ip_addr));
 
+       /* assign connect_query */
+       set_connect_query(db, connect_query);
+
        pktbuf_static(&buf, db->startup_params, sizeof(db->startup_params));
 
        pktbuf_put_string(&buf, "database");

diff --git a/src/server.c b/src/server.c
index 8ceb0634209a8a113df929682a64f3de66acf100..314f2cd0c935306f16c7fa888f623c12479201c4 100644 (file)
--- a/src/server.c
+++ b/src/server.c
@@ -86,6 +86,20 @@ static bool handle_server_startup(PgSocket *server, PktHdr *pkt)
                return false;
        }
 
+       /* ignore most that happens during connect_query */
+       if (server->exec_on_connect) {
+               switch (pkt->type) {
+               case 'Z':
+               case 'S':       /* handle them below */
+                       break;
+
+               case 'E':       /* log & ignore errors */
+                       log_server_error("S: error while executing exec_on_query", pkt);
+               default:        /* ignore rest */
+                       sbuf_prepare_skip(sbuf, pkt->len);
+                       return true;
+               }
+       }
 
        switch (pkt->type) {
        default:
@@ -115,6 +129,18 @@ static bool handle_server_startup(PgSocket *server, PktHdr *pkt)
                break;
 
        case 'Z':               /* ReadyForQuery */
+               if (server->exec_on_connect) {
+                       server->exec_on_connect = 0;
+                       /* deliberately ignore transaction status */
+               } else if (server->pool->db->connect_query) {
+                       server->exec_on_connect = 1;
+                       slog_debug(server, "server conect ok, send exec_on_connect");
+                       SEND_generic(res, server, 'Q', "s", server->pool->db->connect_query);
+                       if (!res)
+                               disconnect_server(server, false, "exec_on_connect query failed");
+                       break;
+               }
+
                /* login ok */
                slog_debug(server, "server login ok, start accepting queries");
                server->ready = 1;