Repository: DataDog/dd-agent

Issue:
pg8000 client doesn't work for our postgres/pgbouncer db setup
We'd like to use the postgres custom metrics feature with our postgres db setup, but can't because of exceptions. To resolve this, swapping the pg8000 client with the psycopg2 one works. See below. I basically made a custom checker and overrode the pg.connect method.

``` python
# pull in datadog's postgresql AgentCheck
import imp
imp.load_source("datadog_checksd_postgres", "/opt/datadog-agent/agent/checks.d/postgres.py")

import sys
import pprint
import psycopg2
import datadog_checksd_postgres
from datadog_checksd_postgres import PostgreSql

###
# replace pg8000 with psycopg2 because pg8000 does not work with our arch.
# the postgres agent fails with the following exception:
#
#2016-08-12 15:50:30 UTC | ERROR | dd.collector | checks.postgres(__init__.py:767) | Check 'postgres' instance #0 failed
# Traceback (most recent call last):
#   File "/opt/datadog-agent/agent/checks/__init__.py", line 750, in run
#     self.check(copy.deepcopy(instance))
#   File "/opt/datadog-agent/agent/checks.d/postgres.py", line 674, in check
#     version = self._get_version(key, db)
#   File "/opt/datadog-agent/agent/checks.d/postgres.py", line 306, in _get_version
#     cursor.execute('SHOW SERVER_VERSION;')
#   File "/opt/datadog-agent/embedded/lib/python2.7/site-packages/pg8000/core.py", line 567, in execute
#     self._c.execute(self, "begin transaction", None)
#   File "/opt/datadog-agent/embedded/lib/python2.7/site-packages/pg8000/core.py", line 1684, in execute
#     self.handle_messages(cursor)
#   File "/opt/datadog-agent/embedded/lib/python2.7/site-packages/pg8000/core.py", line 1761, in handle_messages
#     raise self.error
# ProgrammingError: ('ERROR', '26000', 'prepared statement "pg8000_statement_0" does not exist')
datadog_checksd_postgres.pg = psycopg2
# need to rewire these -- InterfaceError routes to reconnect, ProgrammingError logs a message
datadog_checksd_postgres.InterfaceError = InterfaceError
datadog_checksd_postgres.ProgrammingError = ProgrammingError

###
# remove ssl from the kwargs, otherwise PostgreSql's get_connection fails with the following exception
#
#2016-08-12 16:10:02 UTC | ERROR | dd.collector | checks.customMetricQuery(__init__.py:767) | Check 'customMetricQuery' instance #0 failed
# Traceback (most recent call last):
#   File "/opt/datadog-agent/agent/checks/__init__.py", line 750, in run
#     self.check(copy.deepcopy(instance))
#   File "/opt/datadog-agent/agent/checks.d/postgres.py", line 673, in check
#     db = self.get_connection(key, host, port, user, password, dbname, ssl)
#   File "/opt/datadog-agent/agent/checks.d/postgres.py", line 586, in get_connection
#     password=password, database=dbname, ssl=ssl)
#   File "/opt/datadog-agent/embedded/lib/python2.7/site-packages/psycopg2/__init__.py", line 164, in connect
#     conn = _connect(dsn, connection_factory=connection_factory, async=async)
# OperationalError: invalid connection option "ssl"
psycopg2_connect = psycopg2.connect
def connect(*args, **kwargs):
    del kwargs['ssl']
    return psycopg2_connect(*args, **kwargs)
datadog_checksd_postgres.pg.connect = connect

class Psycopg2Check(PostgreSql):
    pass
```



Discussion:
Also ability to add, custom query in puppet module of dd through puppet interface. 

@sethcleveland could you share more details about what's causing the issue ? What's the custom query that you are running that's failing ?
Also what version of PG are you running ?
Are you hitting pgbouncer or postgres directly ?

@bpatelturnitin can you open an issue in the puppet module repo please ?

This query causes challenges -- cursor.execute('SHOW SERVER_VERSION;'). Its in the stack trace above. It's part of the postgres.py datadog check -- https://github.com/DataDog/dd-agent/blob/master/checks.d/postgres.py#L303. That query runs before any of our custom querys and actually prevents the postgres agent from working. We're hitting the pgbouncer. Querying against postgres works fine with pg8000. FWIW, There's chatter on the pg8000 about querying against a pgbouncer -- https://github.com/mfenniak/pg8000/issues/89, but I don't think there's gonna be much traction/improvment.

It'd be nice to configure which python postgres client to use, or just use psycopg2 and remove pg8000. Is there a reason for using pg8000? 

Also wanted to mention, we tried the pgbouncer here -- https://github.com/DataDog/dd-agent/blob/master/checks.d/pgbouncer.py#L10. It actually uses the psycopg2 client. It errors for completely different reasons -- something to the effect of "Not Allowed." I'm assuming that's a permission issues. I didn't dig much further because that check lacks the custom metrics ability and the postgres built-in metrics. It'd be nice to "marry" these two modules together. I'm no expert on the field, but at least from our usage, this needs a little tender loving care.

Made a tweak to the workaround above to ensure reconnection.

@sethcleveland thanks for the feedback.

Is there a reason you cannot connect to postgres directly ? 
We'd like to keep using pg8000 as it's a pure python module, which ensure it will work on all systems that can run python while psycopg2 might be unavailable in some environments.

Did some investigating. All signs point to not connecting directly to postgres for custom metrics. Basically, we're setup to go through pgbouncer for HA, security, and performance reasons. There are other architectural/design reasons that are beyond my understanding at the moment as well.

@sethcleveland Thanks for the details!

Using pgbouncer for HA is not a great idea for monitoring as you might be querying different PG instances resulting in incoherent metrics. I recommend pointing the agent directly to the local PG instance (it just uses one connection). 

If that's not an option, you should run your modified version as a custom check (just drop it in `/etc/dd-agent/checks.d`).

I'm going to close this issue as switching to psycopg2 for this check will prevent some users from running the check.

We don't want to prevent other users from running the check. As it stands, my company can't use the built-in datadog postgres stuff because the checker's postgres client doesn't work with our database configuration. We are a customer and have a support request for this challenge.

Supporting our needs, enables datadog to potentially support others. There's also a cascading effect when we can't use the default stuff, ex. the datadog provided puppet stuff doesn't work anymore. We have to replicate and adjust that. The datadog puppet stuff's missing a deployment model for custom checks.

I would like datadog to provide metrics (including custom metric support) independent of the database client. In other words, we pick the metrics we want along with a client for a database configuration.  By default, the status quo can be maintained with the option to adjust according to different customer needs. For example, a customer can select a subset of the postgres metrics, add a few custom metrics, and pick the right postgres client for the job, with the default pg8000. 

Or, another approach that works for us. Provide custom metric support in the PGBouncer module, and allow the option to disable those built-in metrics (we don't need them at this time).

Hi @sethcleveland, to be certain we're on the same page, I'm going to have your rep schedule a quick call for us to discuss and find a reasonable way forward for you.

Sounds good. I got the invite.

Good discussion. Thank you for taking the time to talk about this 👍 .


Resolving patch:
diff --git a/checks.d/postgres.py b/checks.d/postgres.py
--- a/checks.d/postgres.py
+++ b/checks.d/postgres.py
@@ -10,8 +10,12 @@
 import socket
 
 # 3rd party
-import pg8000 as pg
-from pg8000 import InterfaceError, ProgrammingError
+try:
+    import psycopg2
+except ImportError:
+    psycopg2 = None
+
+import pg8000
 
 # project
 from checks import AgentCheck, CheckException
@@ -20,6 +24,10 @@
 MAX_CUSTOM_RESULTS = 100
 TABLE_COUNT_LIMIT = 200
 
+def psycopg2_connect(*args, **kwargs):
+    del kwargs['ssl']
+    return psycopg2.connect(*args, **kwargs)
+
 
 class ShouldRestartException(Exception):
     pass
@@ -300,6 +308,16 @@ def __init__(self, name, init_config, agentConfig, instances=None):
         self.replication_metrics = {}
         self.custom_metrics = {}
 
+    def _get_pg_attrs(self, instance):
+        if _is_affirmative(instance.get('use_psycopg2', False)):
+            if psycopg2 is None:
+                self.log.error("Unable to import psycopg2, falling back to pg8000")
+            else:
+                return psycopg2_connect, psycopg2.InterfaceError, psycopg2.ProgrammingError
+
+        # Let's use pg8000
+        return pg8000.connect, pg8000.InterfaceError, pg8000.ProgrammingError
+
     def _get_version(self, key, db):
         if key not in self.versions:
             cursor = db.cursor()
@@ -416,7 +434,7 @@ def _build_relations_config(self, yamlconfig):
                 self.log.warn('Failed to parse config element=%s, check syntax' % str(element))
         return config
 
-    def _collect_stats(self, key, db, instance_tags, relations, custom_metrics, function_metrics, count_metrics):
+    def _collect_stats(self, key, db, instance_tags, relations, custom_metrics, function_metrics, count_metrics, interface_error, programming_error):
         """Query pg_stat_* for various metrics
         If relations is not an empty list, gather per-relation metrics
         on top of that.
@@ -491,7 +509,7 @@ def _collect_stats(self, key, db, instance_tags, relations, custom_metrics, func
                         cursor.execute(query.replace(r'%', r'%%'))
 
                     results = cursor.fetchall()
-                except ProgrammingError as e:
+                except programming_error as e:
                     log_func("Not all metrics may be available: %s" % str(e))
                     continue
 
@@ -556,7 +574,7 @@ def _collect_stats(self, key, db, instance_tags, relations, custom_metrics, func
                         v[0][1](self, v[0][0], v[1], tags=tags)
 
             cursor.close()
-        except InterfaceError as e:
+        except interface_error as e:
             self.log.error("Connection error: %s" % str(e))
             raise ShouldRestartException
         except socket.error as e:
@@ -571,7 +589,7 @@ def _get_service_check_tags(self, host, port, dbname):
         ]
         return service_check_tags
 
-    def get_connection(self, key, host, port, user, password, dbname, ssl, use_cached=True):
+    def get_connection(self, key, host, port, user, password, dbname, ssl, connect_fct, use_cached=True):
         "Get and memoize connections to instances"
         if key in self.dbs and use_cached:
             return self.dbs[key]
@@ -580,12 +598,12 @@ def get_connection(self, key, host, port, user, password, dbname, ssl, use_cache
             try:
                 if host == 'localhost' and password == '':
                     # Use ident method
-                    connection = pg.connect("user=%s dbname=%s" % (user, dbname))
+                    connection = connect_fct("user=%s dbname=%s" % (user, dbname))
                 elif port != '':
-                    connection = pg.connect(host=host, port=port, user=user,
+                    connection = connect_fct(host=host, port=port, user=user,
                         password=password, database=dbname, ssl=ssl)
                 else:
-                    connection = pg.connect(host=host, user=user, password=password,
+                    connection = connect_fct(host=host, user=user, password=password,
                         database=dbname, ssl=ssl)
             except Exception as e:
                 message = u'Error establishing postgres connection: %s' % (str(e))
@@ -667,17 +685,19 @@ def check(self, instance):
         # preset tags to the database name
         db = None
 
+        connect_fct, interface_error, programming_error = self._get_pg_attrs(instance)
+
         # Collect metrics
         try:
             # Check version
-            db = self.get_connection(key, host, port, user, password, dbname, ssl)
+            db = self.get_connection(key, host, port, user, password, dbname, ssl, connect_fct)
             version = self._get_version(key, db)
             self.log.debug("Running check against version %s" % version)
-            self._collect_stats(key, db, tags, relations, custom_metrics, function_metrics, count_metrics)
+            self._collect_stats(key, db, tags, relations, custom_metrics, function_metrics, count_metrics, interface_error, programming_error)
         except ShouldRestartException:
             self.log.info("Resetting the connection")
-            db = self.get_connection(key, host, port, user, password, dbname, ssl, use_cached=False)
-            self._collect_stats(key, db, tags, relations, custom_metrics, function_metrics, count_metrics)
+            db = self.get_connection(key, host, port, user, password, dbname, ssl, connect_fct, use_cached=False)
+            self._collect_stats(key, db, tags, relations, custom_metrics, function_metrics, count_metrics, interface_error, programming_error)
 
         if db is not None:
             service_check_tags = self._get_service_check_tags(host, port, dbname)


Test patch:
diff --git a/tests/checks/integration/test_postgres.py b/tests/checks/integration/test_postgres.py
--- a/tests/checks/integration/test_postgres.py
+++ b/tests/checks/integration/test_postgres.py
@@ -42,7 +42,7 @@ def test_checks(self):
             }
         ]
 
-        self.run_check_twice(dict(instances=instances))
+        self.run_check_twice(dict(instances=instances), force_reload=True)
 
         # Useful to get server version
         # FIXME: Not great, should have a function like that available
@@ -217,3 +217,219 @@ def test_checks(self):
         self.assertServiceMetadata(['version'], count=2)
 
         self.coverage_report()
+        from pg8000.core import Connection
+        self.assertTrue(type(self.check.dbs[key]) == Connection)
+        self.check.dbs[key].close()
+
+    def test_psycopg2(self):
+        host = 'localhost'
+        port = 15432
+        dbname = 'datadog_test'
+
+        instances = [
+            {
+                'host': host,
+                'port': port,
+                'username': 'datadog',
+                'password': 'datadog',
+                'use_psycopg2': 'yes',
+                'dbname': dbname,
+                'relations': ['persons'],
+                'custom_metrics': [{
+                    'descriptors': [('datname', 'customdb')],
+                    'metrics': {
+                        'numbackends': ['custom.numbackends', 'Gauge'],
+                    },
+                    'query': "SELECT datname, %s FROM pg_stat_database WHERE datname = 'datadog_test' LIMIT(1)",
+                    'relation': False,
+                }]
+            },
+            {
+                'host': host,
+                'port': port,
+                'username': 'datadog',
+                'password': 'datadog',
+                'dbname': 'dogs',
+                'relations': ['breed', 'kennel']
+            }
+        ]
+
+        self.run_check_twice(dict(instances=instances), force_reload=True)
+
+        # Useful to get server version
+        # FIXME: Not great, should have a function like that available
+        key = (host, port, dbname)
+        db = self.check.dbs[key]
+
+        # Testing DB_METRICS scope
+        COMMON_METRICS = [
+            'postgresql.connections',
+            'postgresql.commits',
+            'postgresql.rollbacks',
+            'postgresql.disk_read',
+            'postgresql.buffer_hit',
+            'postgresql.rows_returned',
+            'postgresql.rows_fetched',
+            'postgresql.rows_inserted',
+            'postgresql.rows_updated',
+            'postgresql.rows_deleted',
+            'postgresql.database_size',
+        ]
+
+        for mname in COMMON_METRICS:
+            for db in ('datadog_test', 'dogs'):
+                self.assertMetric(mname, count=1, tags=['db:%s' % db])
+
+        NEWER_92_METRICS = [
+            'postgresql.deadlocks',
+            'postgresql.temp_bytes',
+            'postgresql.temp_files',
+        ]
+
+        if self.check._is_9_2_or_above(key, db):
+            for mname in NEWER_92_METRICS:
+                for db in ('datadog_test', 'dogs'):
+                    self.assertMetric(mname, count=1, tags=['db:%s' % db])
+
+        # Testing BGW_METRICS scope
+        COMMON_BGW_METRICS = [
+            'postgresql.bgwriter.checkpoints_timed',
+            'postgresql.bgwriter.checkpoints_requested',
+            'postgresql.bgwriter.buffers_checkpoint',
+            'postgresql.bgwriter.buffers_clean',
+            'postgresql.bgwriter.maxwritten_clean',
+            'postgresql.bgwriter.buffers_backend',
+            'postgresql.bgwriter.buffers_alloc',
+        ]
+
+        for mname in COMMON_BGW_METRICS:
+            self.assertMetric(mname, count=1)
+
+        NEWER_91_BGW_METRICS = [
+            'postgresql.bgwriter.buffers_backend_fsync',
+        ]
+
+        if self.check._is_9_1_or_above(key, db):
+            for mname in NEWER_91_BGW_METRICS:
+                self.assertMetric(mname, count=1)
+
+        NEWER_92_BGW_METRICS = [
+            'postgresql.bgwriter.write_time',
+            'postgresql.bgwriter.sync_time',
+        ]
+
+        if self.check._is_9_2_or_above(key, db):
+            for mname in NEWER_92_BGW_METRICS:
+                self.assertMetric(mname, count=1)
+
+        # FIXME: Test postgresql.locks
+
+        # Relation specific metrics
+        RELATION_METRICS = [
+            'postgresql.seq_scans',
+            'postgresql.seq_rows_read',
+            'postgresql.index_scans',
+            'postgresql.index_rows_fetched',
+            'postgresql.rows_inserted',
+            'postgresql.rows_updated',
+            'postgresql.rows_deleted',
+            'postgresql.rows_hot_updated',
+            'postgresql.live_rows',
+            'postgresql.dead_rows',
+        ]
+
+        SIZE_METRICS = [
+            'postgresql.table_size',
+            'postgresql.index_size',
+            'postgresql.total_size',
+        ]
+
+        STATIO_METRICS = [
+            'postgresql.heap_blocks_read',
+            'postgresql.heap_blocks_hit',
+            'postgresql.index_blocks_read',
+            'postgresql.index_blocks_hit',
+            'postgresql.toast_blocks_read',
+            'postgresql.toast_blocks_hit',
+            'postgresql.toast_index_blocks_read',
+            'postgresql.toast_index_blocks_hit',
+        ]
+
+        for inst in instances:
+            for rel in inst.get('relations', []):
+                expected_tags = ['db:%s' % inst['dbname'], 'table:%s' % rel]
+                expected_rel_tags = ['db:%s' % inst['dbname'], 'table:%s' % rel, 'schema:public']
+                for mname in RELATION_METRICS:
+                    count = 1
+                    # We only build a test index and stimulate it on breed
+                    # in the dogs DB, so the other index metrics shouldn't be
+                    # here.
+                    if 'index' in mname and rel != 'breed':
+                        count = 0
+                    self.assertMetric(mname, count=count, tags=expected_rel_tags)
+
+                for mname in SIZE_METRICS:
+                    self.assertMetric(mname, count=1, tags=expected_tags)
+
+                for mname in STATIO_METRICS:
+                    at_least = None
+                    count = 1
+                    if '.index' in mname and rel != 'breed':
+                        count = 0
+                    # FIXME: toast are not reliable, need to do some more setup
+                    # to get some values here I guess
+                    if 'toast' in mname:
+                        at_least = 0  # how to set easily a flaky metric, w/o impacting coverage
+                        count = None
+                    self.assertMetric(mname, count=count, at_least=at_least, tags=expected_rel_tags)
+
+        # Index metrics
+        IDX_METRICS = [
+            'postgresql.index_scans',
+            'postgresql.index_rows_read',
+            'postgresql.index_rows_fetched',
+        ]
+
+        # we have a single index defined!
+        expected_tags = ['db:dogs', 'table:breed', 'index:breed_names', 'schema:public']
+        for mname in IDX_METRICS:
+            self.assertMetric(mname, count=1, tags=expected_tags)
+
+        # instance connection metrics
+        CONNECTION_METRICS = [
+            'postgresql.max_connections',
+            'postgresql.percent_usage_connections',
+        ]
+        for mname in CONNECTION_METRICS:
+            self.assertMetric(mname, count=1)
+
+        # db level connections
+        for inst in instances:
+            expected_tags = ['db:%s' % inst['dbname']]
+            self.assertMetric('postgresql.connections', count=1, tags=expected_tags)
+
+        # By schema metrics
+        self.assertMetric('postgresql.table.count', value=2, count=1, tags=['schema:public'])
+        self.assertMetric('postgresql.db.count', value=2, count=1)
+
+        # Our custom metric
+        self.assertMetric('custom.numbackends', value=1, tags=['customdb:datadog_test'])
+
+        # Test service checks
+        self.assertServiceCheck('postgres.can_connect',
+            count=1, status=AgentCheck.OK,
+            tags=['host:localhost', 'port:15432', 'db:datadog_test']
+        )
+        self.assertServiceCheck('postgres.can_connect',
+            count=1, status=AgentCheck.OK,
+            tags=['host:localhost', 'port:15432', 'db:dogs']
+        )
+
+        # Assert service metadata
+        self.assertServiceMetadata(['version'], count=2)
+
+        self.coverage_report()
+
+        from psycopg2.extensions import connection
+        self.assertTrue(type(self.check.dbs[key]) == connection)
+        self.check.dbs[key].close()