More updates for the review

checks.d/pgbouncer.py

@@ -16,6 +16,7 @@ class PgBouncer(AgentCheck):
     RATE = AgentCheck.rate
     GAUGE = AgentCheck.gauge
     DB_NAME = 'pgbouncer'
+    SERVICE_CHECK_NAME = 'pgbouncer.can_connect'
 
     STATS_METRICS = {
         'descriptors': [
@@ -56,7 +57,15 @@ def __init__(self, name, init_config, agentConfig, instances=None):
         AgentCheck.__init__(self, name, init_config, agentConfig, instances)
         self.dbs = {}
 
-    def _collect_stats(self, key, db, instance_tags):
+    def _get_service_checks_tags(self, host, port):
+        service_checks_tags = [
+            "host:%s" % host,
+            "port:%s" % port,
+            "db:%s" % self.DB_NAME
+        ]
+        return service_checks_tags
+
+    def _collect_stats(self, db, instance_tags):
         """Query pgbouncer for various metrics
         """
 
@@ -107,11 +116,7 @@ def _get_connection(self, key, host, port, user, password, use_cached=True):
 
         elif host != "" and user != "":
             try:
-                service_check_tags = [
-                    "host:%s" % host,
-                    "port:%s" % port
-                ]
-                service_check_tags.append("db:%s" % self.DB_NAME)
+
 
                 if host == 'localhost' and password == '':
                     # Use ident method
@@ -122,22 +127,22 @@ def _get_connection(self, key, host, port, user, password, use_cached=True):
                 else:
                     connection = pg.connect(host=host, user=user, password=password,
                         database=self.DB_NAME)
-                status = AgentCheck.OK
-                self.service_check('pgbouncer.can_connect', status, tags=service_check_tags)
-                self.log.debug('pgbouncer status: %s' % status)
+
+                connection.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+                self.log.debug('pgbouncer status: %s' % AgentCheck.OK)
 
             except Exception:
-                status = AgentCheck.CRITICAL
-                self.service_check('pgbouncer.can_connect', status, tags=service_check_tags)
-                self.log.debug('pgbouncer status: %s' % status)
-                raise
+                message = u'Cannot establish connection to pgbouncer://%s:%s/%s' % (host, port, self.DB_NAME)
+                self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.CRITICAL,
+                    tags=self._get_service_checks_tags(host, port), message=message)
+                self.log.debug('pgbouncer status: %s' % AgentCheck.CRITICAL)
+                pass
         else:
             if not host:
                 raise CheckException("Please specify a PgBouncer host to connect to.")
             elif not user:
                 raise CheckException("Please specify a user to connect to PgBouncer as.")
 
-        connection.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 
         self.dbs[key] = connection
         return connection
@@ -149,7 +154,7 @@ def check(self, instance):
         password = instance.get('password', '')
         tags = instance.get('tags', [])
 
-        key = '%s:%s:%s' % (host, port, self.DB_NAME)
+        key = '%s:%s' % (host, port)
 
         if tags is None:
             tags = []
@@ -158,8 +163,13 @@ def check(self, instance):
 
         try:
             db = self._get_connection(key, host, port, user, password)
-            self._collect_stats(key, db, tags)
+            self._collect_stats(db, tags)
         except ShouldRestartException:
             self.log.info("Resetting the connection")
             db = self._get_connection(key, host, port, user, password, use_cached=False)
-            self._collect_stats(key, db, tags)
+            self._collect_stats(db, tags)
+
+        message = u'Established connection to pgbouncer://%s:%s/%s' % (host, port, self.DB_NAME)
+        self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.OK,
+            tags=self._get_service_checks_tags(host, port), message=message)
+        self.log.debug('pgbouncer status: %s' % AgentCheck.OK)

ci/pgbouncer.rb

@@ -31,9 +31,6 @@ def pgb_rootdir
            #{pgb_rootdir}/pgbouncer.ini)
       sh %(cp $TRAVIS_BUILD_DIR/ci/resources/pgbouncer/users.txt\
            #{pgb_rootdir}/users.txt)
-      sh %(cp $TRAVIS_BUILD_DIR/ci/resources/pgbouncer/pgbouncer_stimulus.sh\
-           #{pgb_rootdir}/pgbouncer_stimulus.sh)
-      sh %(chmod +x #{pgb_rootdir}/pgbouncer_stimulus.sh)
       sh %(#{pgb_rootdir}/pgbouncer -d #{pgb_rootdir}/pgbouncer.ini)
       sleep_for 3
       sh %(PGPASSWORD=datadog #{pg_rootdir}/bin/psql\

ci/sysstat.rb

@@ -18,7 +18,7 @@ def sysstat_rootdir
       unless Dir.exist? File.expand_path(sysstat_rootdir)
         sh %(curl -s -L\
              -o $VOLATILE_DIR/sysstat-#{sysstat_version}.tar.xz\
-             http://perso.orange.fr/sebastien.godard/sysstat-11.0.1.tar.xz)
+             https://s3.amazonaws.com/travis-archive/sysstat-11.0.1.tar.xz)
         sh %(mkdir -p $VOLATILE_DIR/sysstat)
         sh %(mkdir -p #{sysstat_rootdir})
         sh %(mkdir -p #{sysstat_rootdir}/var/log/sa)

conf.d/pgbouncer.yaml.example

@@ -1,10 +1,10 @@
 init_config:
 
 instances:
-#   -   host: localhost
-#       port: 15433
-#       username: my_username
-#       password: my_password
-#       tags:
-#            - optional_tag1
-#            - optional_tag2
+#   - host: localhost
+#     port: 15433
+#     username: my_username
+#     password: my_password
+#     tags:
+#       - optional_tag1
+#       - optional_tag2

source-optional-requirements.txt

@@ -1,6 +1,7 @@
 kazoo==1.3.1
 pycurl==7.19.5
 psutil==2.1.1
+psycopg2==2.6
 pymongo==2.6.3
 pysnmp-mibs==0.1.4
 pysnmp==4.2.5

source-requirements.txt

@@ -4,7 +4,6 @@ httplib2==0.9
 kafka-python==0.9.0-9bed11db98387c0d9e456528130b330631dc50af
 ntplib==0.3.2
 pg8000==1.9.6
-psycopg2==2.6
 PyMySQL==0.6.1
 python-memcached==1.53
 pyyaml==3.11

tests/test_pgbouncer.py

@@ -1,18 +1,17 @@
-import unittest
-import subprocess
-from tests.common import load_check
-
-from nose.plugins.attrib import attr
+from tests.common import load_check, AgentCheckTest
 
 import time
-from pprint import pprint
+import psycopg2 as pg
+from nose.plugins.attrib import attr
+from checks import AgentCheck
 
 @attr(requires='pgbouncer')
-class TestPgbouncer(unittest.TestCase):
+class TestPgbouncer(AgentCheckTest):
+    CHECK_NAME = 'pgbouncer'
 
     def test_checks(self):
-
         config = {
+            'init_config': {},
             'instances': [
                 {
                 'host': 'localhost',
@@ -22,46 +21,51 @@ def test_checks(self):
                 }
             ]
         }
-        agentConfig = {
-            'version': '0.1',
-            'api_key': 'toto'
-        }
 
-        self.check = load_check('pgbouncer', config, agentConfig)
-        self.check.run()
-        metrics = self.check.get_metrics()
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.cl_active'])           >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.cl_waiting'])          >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.sv_active'])           >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.sv_idle'])             >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.sv_used'])             >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.sv_tested'])           >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.sv_login'])            >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.pools.maxwait'])             >= 1, pprint(metrics))
+        self.run_check(config)
+
+        self.assertMetric('pgbouncer.pools.cl_active')
+        self.assertMetric('pgbouncer.pools.cl_waiting')
+        self.assertMetric('pgbouncer.pools.sv_active')
+        self.assertMetric('pgbouncer.pools.sv_idle')
+        self.assertMetric('pgbouncer.pools.sv_used')
+        self.assertMetric('pgbouncer.pools.sv_tested')
+        self.assertMetric('pgbouncer.pools.sv_login')
+        self.assertMetric('pgbouncer.pools.maxwait')
 
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.total_query_time'])    >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.avg_req'])             >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.avg_recv'])            >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.avg_sent'])            >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.avg_query'])           >= 1, pprint(metrics))
+        self.assertMetric('pgbouncer.stats.total_query_time')
+        self.assertMetric('pgbouncer.stats.avg_req')
+        self.assertMetric('pgbouncer.stats.avg_recv')
+        self.assertMetric('pgbouncer.stats.avg_sent')
+        self.assertMetric('pgbouncer.stats.avg_query')
         # Rate metrics, need 2 collection rounds
-        subprocess.call(['embedded/pgbouncer/pgbouncer_stimulus.sh'])
+        try:
+            connection = pg.connect(
+                host='localhost',
+                port='15433',
+                user='datadog',
+                password='datadog',
+                database='datadog_test')
+            connection.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+            cur = connection.cursor()
+            cur.execute('SELECT * FROM persons;')
+        except Exception:
+            pass
         time.sleep(5)
-        self.check.run()
-        metrics = self.check.get_metrics()
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.requests_per_second'])         >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.bytes_received_per_second'])   >= 1, pprint(metrics))
-        self.assertTrue(len([m for m in metrics if m[0] == u'pgbouncer.stats.bytes_sent_per_second'])       >= 1, pprint(metrics))
+        self.run_check(config)
+        self.assertMetric('pgbouncer.stats.requests_per_second')
+        self.assertMetric('pgbouncer.stats.bytes_received_per_second')
+        self.assertMetric('pgbouncer.stats.bytes_sent_per_second')
 
         # Service checks
-        service_checks = self.check.get_service_checks()
-        service_checks_count = len(service_checks)
-        self.assertTrue(type(service_checks) == type([]))
+        service_checks_count = len(self.service_checks)
+        self.assertTrue(type(self.service_checks) == type([]))
         self.assertTrue(service_checks_count > 0)
-        self.assertEquals(len([sc for sc in service_checks if sc['check'] == "pgbouncer.can_connect"]), 1, service_checks)
-        # Assert that all service checks have the proper tags: host, port and db
-        self.assertEquals(len([sc for sc in service_checks if "host:localhost" in sc['tags']]), service_checks_count, service_checks)
-        self.assertEquals(len([sc for sc in service_checks if "port:%s" % config['instances'][0]['port'] in sc['tags']]), service_checks_count, service_checks)
+        self.assertServiceCheck(
+            'pgbouncer.can_connect',
+            status=AgentCheck.OK,
+            tags=['host:localhost', 'port:15433', 'db:pgbouncer'],
+            count=service_checks_count)
 
 if __name__ == '__main__':
     unittest.main()