Merge branch 'pc-add-password' into 'master'

Add password as an argument

See merge request !1

6 files changed, 57 insertions, 9 deletions

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
new file mode 100644
index 0000000..8ec7a11
--- /dev/null
+++ b/.gitlab-ci.yml
@@ -0,0 +1,15 @@
+image: python:3.3
+
+before_script:
+  - python3 -v
+  - which python
+  - pip3 install psycopg2
+
+unittest:
+  script:
+    - python3 -m unittest discover tests
+
+install:
+  script:
+    - python3 setup.py install
+    - /usr/local/bin/prometheus-pgbouncer-exporter --version
diff --git a/prometheus_pgbouncer_exporter/__init__.py b/prometheus_pgbouncer_exporter/__init__.py
index 45b81bf..923225e 100644
--- a/prometheus_pgbouncer_exporter/__init__.py
+++ b/prometheus_pgbouncer_exporter/__init__.py
@@ -1,3 +1,3 @@
 __ver_major__ = 1
-__ver_minor__ = 7
+__ver_minor__ = 8
 __version__ = "%d.%d" % (__ver_major__, __ver_minor__)
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..ea69f46 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.items() 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),
diff --git a/setup.py b/setup.py
index 3035917..72e2f24 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,9 @@
+import sys
 from setuptools import setup, find_packages
 
+if sys.version_info.major < 3:
+    raise RuntimeError('Installing requires Python 3 or newer')
+
 setup(
     name="prometheus-pgbouncer-exporter",
     packages=find_packages(),
@@ -8,4 +12,9 @@ setup(
             'prometheus-pgbouncer-exporter = prometheus_pgbouncer_exporter.cli:main',
         ],
     },
+    install_requires=[
+        'psycopg2',
+        'ConfigArgParse',
+        'prometheus_client',
+    ],
 )
diff --git a/tests/test_connection.py b/tests/test_connection.py
new file mode 100644
index 0000000..d4fb134
--- /dev/null
+++ b/tests/test_connection.py
@@ -0,0 +1,18 @@
+import unittest
+
+from unittest.mock import patch, Mock
+
+from prometheus_pgbouncer_exporter import utils
+
+class ConnectionTest(unittest.TestCase):
+
+    @patch('prometheus_pgbouncer_exporter.utils.psycopg2.connect')
+    def test_get_connection_with_host_and_dbname(self, connect):
+        conn = utils.get_connection(host='/tmp/', dbname='template1')
+        connect.assert_called_once_with(host='/tmp/', dbname='template1')
+        connect.return_value.set_session.assert_called_once_with(autocommit=True)
+
+    @patch('prometheus_pgbouncer_exporter.utils.psycopg2.connect')
+    def test_get_connection_with_old_args_plus_password_works(self, connect):
+        conn = utils.get_connection(user='pablo', port=5432, host='localhost', password='mypassword')
+        connect.assert_called_once_with(dbname='pgbouncer', user='pablo', port=5432, host='localhost', password='mypassword')