Merge pull request #1273 from DataDog/conor/pg-service-check

[postgres] Always sent OK status, even with cached connection.
DataDog · Jan 2, 2015 · 9e89aaf · 9e89aaf
2 parents 499450a + 57fa171
commit 9e89aaf
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 17 deletions.
diff --git a/checks.d/postgres.py b/checks.d/postgres.py
@@ -19,6 +19,7 @@ class PostgreSql(AgentCheck):
     RATE = AgentCheck.rate
     GAUGE = AgentCheck.gauge
     MONOTONIC = AgentCheck.monotonic_count
+    SERVICE_CHECK_NAME = 'postgres.can_connect'
 
     # turning columns into tags
     DB_METRICS = {
@@ -181,7 +182,7 @@ class PostgreSql(AgentCheck):
 SELECT %s
   FROM pg_stat_database, max_con
 """
-    }        
+    }
 
     def __init__(self, name, init_config, agentConfig):
         AgentCheck.__init__(self, name, init_config, agentConfig)
@@ -322,20 +323,21 @@ def _collect_stats(self, key, db, instance_tags, relations):
             self.log.error("Connection error: %s" % str(e))
             raise ShouldRestartException
 
+    def _get_service_check_tags(self, host, port, dbname):
+        service_check_tags = [
+            "host:%s" % host,
+            "port:%s" % port,
+            "db:%s" % dbname
+        ]
+        return service_check_tags
+
     def get_connection(self, key, host, port, user, password, dbname, 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:
-                service_check_tags = [
-                    "host:%s" % host,
-                    "port:%s" % port
-                ]
-                if dbname:
-                    service_check_tags.append("db:%s" % dbname)
-
                 if host == 'localhost' and password == '':
                     # Use ident method
                     connection = pg.connect("user=%s dbname=%s" % (user, dbname))
@@ -345,14 +347,11 @@ def get_connection(self, key, host, port, user, password, dbname, use_cached=Tru
                 else:
                     connection = pg.connect(host=host, user=user, password=password,
                         database=dbname)
-                status = AgentCheck.OK
-                self.service_check('postgres.can_connect', status, tags=service_check_tags)
-                self.log.info('pg status: %s' % status)
-
-            except Exception:
-                status = AgentCheck.CRITICAL
-                self.service_check('postgres.can_connect', status, tags=service_check_tags)
-                self.log.info('pg status: %s' % status)
+            except Exception as e:
+                message = u'Error establishing postgres connection: %s' % (str(e))
+                service_check_tags = self._get_service_check_tags(host, port, dbname)
+                self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.CRITICAL,
+                    tags=service_check_tags, message=message)
                 raise
         else:
             if not host:
@@ -406,6 +405,10 @@ def check(self, instance):
             self._collect_stats(key, db, tags, relations)
 
         if db is not None:
+            service_check_tags = self._get_service_check_tags(host, port, dbname)
+            message = u'Established connection to postgres://%s:%s/%s' % (host, port, dbname)
+            self.service_check(self.SERVICE_CHECK_NAME, AgentCheck.OK,
+                tags=service_check_tags, message=message)
             try:
                 # commit to close the current query transaction
                 db.commit()

diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
@@ -59,7 +59,7 @@ def testChecks(self):
         service_checks_count = len(service_checks)
         self.assertTrue(type(service_checks) == type([]))
         self.assertTrue(service_checks_count > 0)
-        self.assertEquals(len([sc for sc in service_checks if sc['check'] == "postgres.can_connect"]), 1, service_checks)
+        self.assertEquals(len([sc for sc in service_checks if sc['check'] == "postgres.can_connect"]), 2, 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)