Merge pull request DataDog#1391 from DataDog/dorian/pgbouncer

Add pgbouncer check and tests

.travis.yml

@@ -65,6 +65,7 @@ env:
     - TRAVIS_FLAVOR=fluentd
     - TRAVIS_FLAVOR=rabbitmq
     - TRAVIS_FLAVOR=etcd
+    - TRAVIS_FLAVOR=pgbouncer
 
 # Override travis defaults with empty jobs
 before_install: echo "OVERRIDING TRAVIS STEPS"

Rakefile

@@ -18,6 +18,7 @@ require './ci/memcache'
 require './ci/mongo'
 require './ci/mysql'
 require './ci/nginx'
+require './ci/pgbouncer'
 require './ci/postgres'
 require './ci/rabbitmq'
 require './ci/redis'

checks.d/pgbouncer.py

@@ -0,0 +1,175 @@
+"""Pgbouncer check
+
+Collects metrics from the pgbouncer database.
+"""
+from checks import AgentCheck, CheckException
+
+import psycopg2 as pg
+
+
+class ShouldRestartException(Exception): pass
+
+
+class PgBouncer(AgentCheck):
+    """Collects metrics from pgbouncer
+    """
+    RATE = AgentCheck.rate
+    GAUGE = AgentCheck.gauge
+    DB_NAME = 'pgbouncer'
+    SERVICE_CHECK_NAME = 'pgbouncer.can_connect'
+
+    STATS_METRICS = {
+        'descriptors': [
+            ('database', 'db'),
+        ],
+        'metrics': {
+            'total_requests':       ('pgbouncer.stats.requests_per_second', RATE),
+            'total_received':       ('pgbouncer.stats.bytes_received_per_second', RATE),
+            'total_sent':           ('pgbouncer.stats.bytes_sent_per_second', RATE),
+            'total_query_time':     ('pgbouncer.stats.total_query_time', GAUGE),
+            'avg_req':              ('pgbouncer.stats.avg_req', GAUGE),
+            'avg_recv':             ('pgbouncer.stats.avg_recv', GAUGE),
+            'avg_sent':             ('pgbouncer.stats.avg_sent', GAUGE),
+            'avg_query':            ('pgbouncer.stats.avg_query', GAUGE),
+        },
+        'query': """SHOW STATS""",
+    }
+
+    POOLS_METRICS = {
+        'descriptors': [
+            ('database', 'db'),
+            ('user', 'user'),
+        ],
+        'metrics': {
+            'cl_active':            ('pgbouncer.pools.cl_active', GAUGE),
+            'cl_waiting':           ('pgbouncer.pools.cl_waiting', GAUGE),
+            'sv_active':            ('pgbouncer.pools.sv_active', GAUGE),
+            'sv_idle':              ('pgbouncer.pools.sv_idle', GAUGE),
+            'sv_used':              ('pgbouncer.pools.sv_used', GAUGE),
+            'sv_tested':            ('pgbouncer.pools.sv_tested', GAUGE),
+            'sv_login':             ('pgbouncer.pools.sv_login', GAUGE),
+            'maxwait':              ('pgbouncer.pools.maxwait', GAUGE),
+        },
+        'query': """SHOW POOLS""",
+    }
+
+    def __init__(self, name, init_config, agentConfig, instances=None):
+        AgentCheck.__init__(self, name, init_config, agentConfig, instances)
+        self.dbs = {}
+
+    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
+        """
+
+        metric_scope = [self.STATS_METRICS, self.POOLS_METRICS]
+
+        try:
+            cursor = db.cursor()
+            for scope in metric_scope:
+
+                cols = scope['metrics'].keys()
+
+                try:
+                    query = scope['query']
+                    self.log.debug("Running query: %s" % query)
+                    cursor.execute(query)
+
+                    results = cursor.fetchall()
+                except pg.Error, e:
+                    self.log.warning("Not all metrics may be available: %s" % str(e))
+                    continue
+
+                for row in results:
+                    if row[0] == self.DB_NAME:
+                        continue
+
+                    desc = scope['descriptors']
+                    assert len(row) == len(cols) + len(desc)
+
+                    tags = instance_tags[:]
+                    tags += ["%s:%s" % (d[0][1], d[1]) for d in zip(desc, row[:len(desc)])]
+
+                    values = zip([scope['metrics'][c] for c in cols], row[len(desc):])
+
+                    [v[0][1](self, v[0][0], v[1], tags=tags) for v in values]
+
+            if not results:
+                self.warning('No results were found for query: "%s"' % query)
+
+            cursor.close()
+        except pg.Error, e:
+            self.log.error("Connection error: %s" % str(e))
+            raise ShouldRestartException
+
+    def _get_connection(self, key, host, port, user, password, use_cached=True):
+        "Get and memoize connections to instances"
+        if key in self.dbs and use_cached:
+            return self.dbs[key]
+
+        elif host != "" and user != "":
+            try:
+
+
+                if host == 'localhost' and password == '':
+                    # Use ident method
+                    connection = pg.connect("user=%s dbname=%s" % (user, self.DB_NAME))
+                elif port != '':
+                    connection = pg.connect(host=host, port=port, user=user,
+                        password=password, database=self.DB_NAME)
+                else:
+                    connection = pg.connect(host=host, user=user, password=password,
+                        database=self.DB_NAME)
+
+                connection.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+                self.log.debug('pgbouncer status: %s' % AgentCheck.OK)
+
+            except Exception:
+                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.")
+
+
+        self.dbs[key] = connection
+        return connection
+
+    def check(self, instance):
+        host = instance.get('host', '')
+        port = instance.get('port', '')
+        user = instance.get('username', '')
+        password = instance.get('password', '')
+        tags = instance.get('tags', [])
+
+        key = '%s:%s' % (host, port)
+
+        if tags is None:
+            tags = []
+        else:
+            tags = list(set(tags))
+
+        try:
+            db = self._get_connection(key, host, port, user, password)
+            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(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

@@ -0,0 +1,76 @@
+require './ci/common'
+require './ci/postgres'
+
+def pgb_rootdir
+  "#{ENV['INTEGRATIONS_DIR']}/pgbouncer"
+end
+
+
+namespace :ci do
+  namespace :pgbouncer do |flavor|
+    task :before_install => ['ci:common:before_install']
+
+    task :install do
+      Rake::Task['ci:postgres:install'].invoke
+      unless Dir.exist? File.expand_path(pgb_rootdir)
+        sh %(wget -O $VOLATILE_DIR/pgbouncer-1.5.4.tar.gz https://s3.amazonaws.com/travis-archive/pgbouncer-1.5.4.tar.gz)
+        sh %(mkdir -p $VOLATILE_DIR/pgbouncer)
+        sh %(tar xzf $VOLATILE_DIR/pgbouncer-1.5.4.tar.gz\
+             -C $VOLATILE_DIR/pgbouncer --strip-components=1)
+        sh %(mkdir -p #{pgb_rootdir})
+        sh %(cd $VOLATILE_DIR/pgbouncer\
+             && ./configure --prefix=#{pgb_rootdir}\
+             && make\
+             && cp pgbouncer #{pgb_rootdir})
+      end
+    end
+
+    task :before_script do
+      Rake::Task['ci:postgres:before_script'].invoke
+      sh %(cp $TRAVIS_BUILD_DIR/ci/resources/pgbouncer/pgbouncer.ini\
+           #{pgb_rootdir}/pgbouncer.ini)
+      sh %(cp $TRAVIS_BUILD_DIR/ci/resources/pgbouncer/users.txt\
+           #{pgb_rootdir}/users.txt)
+      sh %(#{pgb_rootdir}/pgbouncer -d #{pgb_rootdir}/pgbouncer.ini)
+      sleep_for 3
+      sh %(PGPASSWORD=datadog #{pg_rootdir}/bin/psql\
+           -h localhost -p 15433 -U datadog -w\
+           -c "SELECT * FROM persons"\
+           datadog_test)
+      sleep_for 3
+    end
+
+    task :script do
+      this_provides = [
+        'pgbouncer'
+      ]
+      Rake::Task['ci:common:run_tests'].invoke(this_provides)
+    end
+
+    task :cleanup do
+      sh %(rm -rf $VOLATILE_DIR/pgbouncer*)
+      sh %(killall pgbouncer)
+      Rake::Task['ci:postgres:cleanup'].invoke
+    end
+
+    task :execute do
+      exception = nil
+      begin
+        %w(before_install install before_script script).each do |t|
+          Rake::Task["#{flavor.scope.path}:#{t}"].invoke
+        end
+      rescue => e
+        exception = e
+        puts "Failed task: #{e.class} #{e.message}".red
+      end
+      if ENV['SKIP_CLEANUP']
+        puts 'Skipping cleanup, disposable environments are great'.yellow
+      else
+        puts 'Cleaning up'
+        Rake::Task["#{flavor.scope.path}:cleanup"].invoke
+      end
+      fail exception if exception
+    end
+
+  end
+end

ci/resources/pgbouncer/pgbouncer.ini

@@ -0,0 +1,11 @@
+[databases]
+datadog_test = host=127.0.0.1 port=15432 dbname=datadog_test
+
+[pgbouncer]
+listen_port = 15433
+listen_addr = *
+auth_type = md5
+auth_file = embedded/pgbouncer/users.txt
+admin_users = datadog
+logfile = /tmp/pgbouncer.log
+pidfile = /tmp/pgbouncer.pid

ci/resources/pgbouncer/users.txt

@@ -0,0 +1 @@
+"datadog" "datadog"

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

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

requirements.txt

@@ -14,6 +14,7 @@ pysnmp-mibs
 pymysql
 pyvmomi==5.5.0
 pg8000
+psycopg2
 ntplib
 httplib2
 kafka-python==0.9.0-9bed11db98387c0d9e456528130b330631dc50af

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

@@ -12,4 +12,4 @@ requests==2.3.0
 simplejson==3.3.3
 snakebite==1.3.9
 backports.ssl_match_hostname==3.4.0.2
-tornado==3.2.2
+tornado==3.2.2

tests/test_pgbouncer.py

@@ -0,0 +1,71 @@
+from tests.common import load_check, AgentCheckTest
+
+import time
+import psycopg2 as pg
+from nose.plugins.attrib import attr
+from checks import AgentCheck
+
+@attr(requires='pgbouncer')
+class TestPgbouncer(AgentCheckTest):
+    CHECK_NAME = 'pgbouncer'
+
+    def test_checks(self):
+        config = {
+            'init_config': {},
+            'instances': [
+                {
+                'host': 'localhost',
+                'port': 15433,
+                'username': 'datadog',
+                'password': 'datadog'
+                }
+            ]
+        }
+
+        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.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
+        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.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_count = len(self.service_checks)
+        self.assertTrue(type(self.service_checks) == type([]))
+        self.assertTrue(service_checks_count > 0)
+        self.assertServiceCheck(
+            'pgbouncer.can_connect',
+            status=AgentCheck.OK,
+            tags=['host:localhost', 'port:15433', 'db:pgbouncer'],
+            count=service_checks_count)
+
+if __name__ == '__main__':
+    unittest.main()