Add password as a connection option

2 files changed, 14 insertions, 8 deletions

diff --git a/prometheus_pgbouncer_exporter/cli.py b/prometheus_pgbouncer_exporter/cli.py
index 4f55f1e..91bbd0a 100644
--- a/prometheus_pgbouncer_exporter/cli.py
+++ b/prometheus_pgbouncer_exporter/cli.py
@@ -75,6 +75,12 @@ def main():
         env_var='PGBOUNCER_USER',
     )
     p.add(
+        '--pgbouncer-password',
+        default=None,
+        help="Password to connect to pgbouncer with",
+        env_var='PGBOUNCER_PASSWORD',
+    )
+    p.add(
         '--pgbouncer-host',
         default=None,
         help="Host on which to connect to pgbouncer",
@@ -106,9 +112,10 @@ def main():
     logging.info(p.format_values())
 
     connection = get_connection(
-        options.pgbouncer_user,
-        options.pgbouncer_port,
-        options.pgbouncer_host,
+        user = options.pgbouncer_user,
+        port = options.pgbouncer_port,
+        host = options.pgbouncer_host,
+        password = options.pgbouncer_password,
     )
 
     REGISTRY.register(StatsCollector(
diff --git a/prometheus_pgbouncer_exporter/utils.py b/prometheus_pgbouncer_exporter/utils.py
index 7f9aae0..3a99695 100644
--- a/prometheus_pgbouncer_exporter/utils.py
+++ b/prometheus_pgbouncer_exporter/utils.py
@@ -14,12 +14,11 @@
 import psycopg2
 
 
-def get_connection(user, port, host):
+def get_connection(user=None, port=None, host=None, dbname='pgbouncer', password=None):
+    kwargs = { 'user': user, 'port': port, 'host': host, 'dbname': dbname, 'password': password }
+    kwargs = dict([(k, v) for k, v in kwargs.iteritems() if v])
     connection = psycopg2.connect(
-        database='pgbouncer',
-        user=user,
-        port=port,
-        host=host,
+        **kwargs
     )
 
     # pgbouncer does not support transactions (as it does not make sense to),