Merge lp:~lucio.torre/txstatsd/add-distinct-db-plugin into lp:~txstatsd-dev/txstatsd/distinct-plugin

[Lucio Torre]

talk to the database to save the buckets of info

[Sidnei da Silva (sidnei)]

My biggest concern here is that if this fails to connect to postgres, it will go in an infinite loop, with no logging at all and no sleep() between retries, so it can be a very tight loop consumin cpu with no visibility.

It should limit the amount of retries and at least have some logging if postgres cannot be reached.

[Sidnei da Silva (sidnei)]

review: Needs Fixing

[lp:~lucio.torre/txstatsd/add-distinct-db-plugin]

5. By Lucio Torre on 2011-12-09

remove while, use INSERT .. RETURNING

6. By Lucio Torre on 2011-12-09

added test files

[Sidnei da Silva (sidnei)]

Looks good now. My last wish is a Makefile with a 'test' target that performs the steps outlined in the README. So that I can add this as a distinct (no pun intended) project in tarmac and run the tests before merge. You can choose to ignore me though.

review: Approve

=== added file 'README'
--- README	1970-01-01 00:00:00 +0000
+++ README	2011-12-09 20:08:25 +0000
@@ -0,0 +1,10 @@
+To test:
+./bin/start-database.sh
+psql -h `pwd`/tmp/db1 -d distinct < bin/schema.sql
+
+Then:
+trial distinctdb
+
+When done:
+./bin/stop-database.sh
+

=== added directory 'bin'
=== added file 'bin/schema.sql'
--- bin/schema.sql	1970-01-01 00:00:00 +0000
+++ bin/schema.sql	2011-12-09 20:08:25 +0000
@@ -0,0 +1,3 @@
+CREATE TABLE paths (id SERIAL PRIMARY KEY NOT NULL, path TEXT NOT NULL UNIQUE);
+CREATE TABLE points (path_id INTEGER, bucket INTEGER, value TEXT, count INTEGER);
+CREATE INDEX points_idx ON points (path_id, bucket);

=== added file 'bin/start-database.sh'
--- bin/start-database.sh	1970-01-01 00:00:00 +0000
+++ bin/start-database.sh	2011-12-09 20:08:25 +0000
@@ -0,0 +1,63 @@
+#! /bin/bash
+
+ROOTDIR=${ROOTDIR:-`bzr root`}
+if [ ! -d "$ROOTDIR"  ]; then
+    echo "ROOTDIR '$ROOTDIR' doesn't exist" >&2
+    exit 1
+fi
+
+DATABASES="
+distinct
+"
+
+function setup_database() {
+    local TESTDIR=$1
+
+    echo "## Starting postgres in $TESTDIR ##"
+    mkdir -p "$TESTDIR/data"
+    chmod 700 "$TESTDIR/data"
+
+    export PGHOST="$TESTDIR"
+    export PGDATA="$TESTDIR/data"
+    if [ -d /usr/lib/postgresql/8.4 ]; then
+        export PGBINDIR=/usr/lib/postgresql/8.4/bin
+    elif [ -d /usr/lib/postgresql/8.3 ]; then
+        export PGBINDIR=/usr/lib/postgresql/8.3/bin
+    else
+        echo "Cannot find valid parent for PGBINDIR"
+    fi
+    $PGBINDIR/initdb -E UNICODE -D $PGDATA
+    # set up the database options file
+    if [ ! -e $PGDATA/postgresql.conf ]; then
+        echo "PostgreSQL data directory apparently didn't init"
+    else
+    (
+        cat <<EOF
+search_path='\$user,public,ts2'
+add_missing_from=false
+log_statement='all'
+log_line_prefix='[%m] %q%u@%d %c '
+fsync = off
+EOF
+    ) > $PGDATA/postgresql.conf
+    fi
+    $PGBINDIR/initdb -A trust &>/dev/null
+    $PGBINDIR/pg_ctl start -w -D $TESTDIR/data -l $TESTDIR/postgres.log -o "-F -k $TESTDIR -h ''"
+    for db in $DATABASES; do
+        $PGBINDIR/createdb --encoding UNICODE "$db" &>/dev/null
+        $PGBINDIR/createlang plpgsql "$db"
+    done
+    $PGBINDIR/createuser --superuser --createdb "postgres" &>/dev/null
+    # create the additional users we need via a psql script
+    $PGBINDIR/psql -U postgres template1 <<EOF
+CREATE ROLE client INHERIT;
+
+CREATE USER client IN ROLE client;
+EOF
+    echo "To set your environment so psql will connect to this DB instance type:"
+    echo "    export PGHOST=$TESTDIR"
+    echo "## Done. ##"
+    echo -n host=$TESTDIR dbname=distinct > $ROOTDIR/tmp/pg.dsn
+}
+
+setup_database $ROOTDIR/tmp/db1

=== added file 'bin/stop-database.sh'
--- bin/stop-database.sh	1970-01-01 00:00:00 +0000
+++ bin/stop-database.sh	2011-12-09 20:08:25 +0000
@@ -0,0 +1,22 @@
+#! /bin/bash
+
+ROOTDIR=${ROOTDIR:-`bzr root`}
+if [ ! -d "$ROOTDIR"  ]; then
+    echo "ROOTDIR '$ROOTDIR' doesn't exist" >&2
+    exit 1
+fi
+
+if [ -d /usr/lib/postgresql/8.4 ]; then
+    PGBINDIR=/usr/lib/postgresql/8.4/bin
+elif [ -d /usr/lib/postgresql/8.3 ]; then
+    PGBINDIR=/usr/lib/postgresql/8.3/bin
+else
+    echo "Cannot find valid parent for PGBINDIR"
+fi
+
+# setting PGDATA tells pg_ctl which DB to talk to
+export PGDATA=$ROOTDIR/tmp/db1/data/
+$PGBINDIR/pg_ctl status > /dev/null
+if [ $? = 0 ]; then
+    $PGBINDIR/pg_ctl stop -t 60 -w -m fast
+fi

=== renamed directory 'distinctplugin' => 'distinctdb'
=== modified file 'distinctdb/distinctmetric.py'
--- distinctplugin/distinctmetric.py	2011-12-06 19:39:21 +0000
+++ distinctdb/distinctmetric.py	2011-12-09 20:08:25 +0000
@@ -1,17 +1,23 @@
+import time
+
+import psycopg2
+
+from zope.interface import implements
+from twisted.internet.threads import deferToThread
+from txstatsd.itxstatsd import IMetric
+
+ONEDAY = 60 * 60 * 24
+
+
 class DistinctMetricReporter(object):
     """
-    Keeps an estimate of the distinct numbers of items seen on various
-    sliding windows of time.
+    Keeps an mesurement of the distinct numbers of items seen and the times
+    it has seen each one.
     """
     implements(IMetric)
 
-    MESSAGE = (
-        "$prefix%(key)s.count_1min %(count_1min)s %(timestamp)s\n"
-        "$prefix%(key)s.count_1hour %(count_1hour)s %(timestamp)s\n"
-        "$prefix%(key)s.count_1day %(count_1day)s %(timestamp)s\n"
-        "$prefix%(key)s.count %(count)s %(timestamp)s\n")
-
-    def __init__(self, name, wall_time_func=time.time, prefix=""):
+    def __init__(self, name, wall_time_func=time.time, prefix="",
+            bucket_size=ONEDAY, dsn=None):
         """Construct a metric we expect to be periodically updated.
 
         @param name: Indicates what is being instrumented.
@@ -21,36 +27,70 @@
         """
         self.name = name
         self.wall_time_func = wall_time_func
-        self.counter = SlidingDistinctCounter(32, 32)
         if prefix:
             prefix += '.'
-        self.message = Template(DistinctMetricReporter.MESSAGE).substitute(
-            prefix=prefix)
-
-    def count(self):
-        return self.counter.distinct()
-
-    def count_1min(self, now):
-        return self.counter.distinct(now - 60)
-
-    def count_1hour(self, now):
-        return self.counter.distinct(now - 60 * 60)
-
-    def count_1day(self, now):
-        return self.counter.distinct(now - 60 * 60 * 24)
+        self.prefix = prefix
+        self.bucket_size = bucket_size
+        self.dsn = dsn
+        self.metric_id = None
+        self.build_bucket()
+
+    def build_bucket(self, timestamp=None):
+        self.max = 0
+        self.bucket = {}
+        self.bucket_no = self.get_bucket_no(timestamp)
+
+    def get_bucket_no(self, timestamp=None):
+        if timestamp is None:
+            timestamp = self.wall_time_func()
+        return int(timestamp / (self.bucket_size))
 
     def process(self, fields):
         self.update(fields[0])
 
     def update(self, item):
-        self.counter.add(self.wall_time_func(), item)
+        value = self.bucket.get(item, 0) + 1
+
+        self.bucket[item] = value
+        if value > self.max:
+            self.max = value
+
+    def _save_bucket(self, bucket, bucket_no):
+        path = self.prefix + self.name
+        if self.metric_id is None:
+            c = psycopg2.connect(self.dsn)
+            cr = c.cursor()
+            cr.execute("SELECT * FROM paths WHERE path = %s", (path,))
+            row = cr.fetchone()
+            if row is None:
+                cr.execute("INSERT INTO paths (path) VALUES (%s) "
+                            "RETURNING (id)", (path,))
+                row = cr.fetchone()
+                cr.execute("commit")
+
+            self.metric_id = row[0]
+
+        for i, (k, v) in enumerate(bucket.iteritems()):
+            cr.execute("INSERT INTO points (path_id, bucket, value, count) "
+                "VALUES (%s, %s, %s, %s)", (self.metric_id, bucket_no,
+                                            k, v))
+            if i % 1000 == 0:
+                cr.execute("commit")
+        cr.execute("commit")
+
+    def save_bucket(self, bucket, bucket_no):
+        if self.dsn is not None:
+            deferToThread(self._save_bucket, bucket, bucket_no)
 
     def flush(self, interval, timestamp):
-        now = self.wall_time_func()
-        return self.message % {
-            "key": self.name,
-            "count": self.count(),
-            "count_1min": self.count_1min(now),
-            "count_1hour": self.count_1hour(now),
-            "count_1day": self.count_1day(now),
-            "timestamp": timestamp}
+        current_bucket = self.get_bucket_no(timestamp)
+        if current_bucket != self.bucket_no:
+            self.save_bucket(self.bucket, self.bucket_no)
+            self.build_bucket(timestamp)
+
+        metrics = []
+        items = {".count": len(self.bucket),
+                 ".max": self.max}
+        for item, value in items.iteritems():
+            metrics.append((self.prefix + self.name + item, value, timestamp))
+        return metrics

=== added directory 'distinctdb/tests'
=== added file 'distinctdb/tests/__init__.py'
=== added file 'distinctdb/tests/test_distinct.py'
--- distinctdb/tests/test_distinct.py	1970-01-01 00:00:00 +0000
+++ distinctdb/tests/test_distinct.py	2011-12-09 20:08:25 +0000
@@ -0,0 +1,147 @@
+# Copyright (C) 2011 Canonical
+# All Rights Reserved
+
+import ConfigParser
+from cStringIO import StringIO
+import os
+import time
+import subprocess
+
+import psycopg2
+
+from twisted.trial.unittest import TestCase
+from twisted.plugin import getPlugins
+from twisted.plugins import distinctdbplugin
+from txstatsd.itxstatsd import IMetricFactory
+from txstatsd import service
+
+from distinctdb import distinctmetric as distinct
+
+
+class TestDistinctMetricReporter(TestCase):
+
+    def test_get_bucket_no(self):
+        _wall_time = [0]
+
+        def _time():
+            return _wall_time[0]
+
+        dmr = distinct.DistinctMetricReporter("test", wall_time_func=_time)
+        self.assertEquals(dmr.get_bucket_no(), 0)
+        _wall_time = [60 * 60 * 24 + 1]
+        dmr.update("one")
+        self.assertEquals(dmr.get_bucket_no(), 1)
+
+    def test_max(self):
+        _wall_time = [0]
+
+        def _time():
+            return _wall_time[0]
+
+        dmr = distinct.DistinctMetricReporter("test", wall_time_func=_time)
+        self.assertEquals(dmr.get_bucket_no(), 0)
+        self.assertEquals(dmr.max, 0)
+        dmr.update("one")
+        dmr.update("one")
+        dmr.update("two")
+        self.assertEquals(dmr.max, 2)
+        dmr.flush(1, 60 * 60 * 24 + 1)
+        dmr.update("one")
+        self.assertEquals(dmr.max, 1)
+
+    def test_reports(self):
+        _wall_time = [0]
+
+        def _time():
+            return _wall_time[0]
+
+        result = {}
+
+        dmr = distinct.DistinctMetricReporter("test", wall_time_func=_time)
+
+        def save(b, b_no):
+            result["bucket"] = b
+            result["bucket_no"] = b_no
+        dmr.save_bucket = save
+        dmr.update("one")
+        dmr.update("one")
+        dmr.update("two")
+        day = 60 * 60 * 24 + 1
+        dmr.flush(1, day)
+        dmr.update("three")
+        self.assertEquals(result,
+            {"bucket": {"one": 2, "two": 1}, "bucket_no": 0})
+        self.assertEquals(dmr.flush(1, day),
+            [("test.max", 1, day), ("test.count", 1, day)])
+
+    def test_configure(self):
+        class TestOptions(service.OptionsGlue):
+            optParameters = [["test", "t", "default", "help"]]
+            config_section = "statsd"
+
+        o = TestOptions()
+        config_file = ConfigParser.RawConfigParser()
+        config_file.readfp(StringIO("[statsd]\n\n[plugin_distinctdb]\n"
+            "dsn = dbdsn\nbucket_size = 100"))
+        o.configure(config_file)
+        dmf = distinctdbplugin.DistinctMetricFactory()
+        dmf.configure(o)
+        dmr = dmf.build_metric("foo", "bar", time.time)
+        self.assertEquals(dmr.bucket_size, 100)
+        self.assertEquals(dmr.dsn, "dbdsn")
+
+
+class TestPlugin(TestCase):
+
+    def test_factory(self):
+        self.assertTrue(distinctdbplugin.distinct_metric_factory in \
+                        list(getPlugins(IMetricFactory)))
+
+
+class TestDatabase(TestCase):
+
+    def setUp(self):
+        rootdir = subprocess.check_output(["bzr", "root"]).strip()
+        dsn_file = os.path.join(rootdir, "tmp", "pg.dsn")
+        self.dsn = open(dsn_file).read()
+        self.conn = psycopg2.connect(self.dsn)
+
+    def tearDown(self):
+        cr = self.conn.cursor()
+        cr.execute("rollback")
+        cr.execute("DELETE FROM paths")
+        cr.execute("DELETE FROM points")
+        cr.execute("commit")
+
+    def test_connect(self):
+        cr = self.conn.cursor()
+        cr.execute("SELECT 0")
+        result = cr.fetchall()
+        self.assertTrue(result, [(0,)])
+
+    def test_create_metric_id(self):
+        dmr = distinct.DistinctMetricReporter("test", dsn=self.dsn)
+        dmr._save_bucket({}, 0)
+        cr = self.conn.cursor()
+        cr.execute("SELECT * FROM paths WHERE path = 'test'")
+        cr.execute("SELECT * FROM paths")
+        self.assertEquals(len(cr.fetchall()), 1)
+
+    def test_find_saved_data(self):
+        dmr = distinct.DistinctMetricReporter("test", dsn=self.dsn)
+        dmr.update("one")
+        dmr.update("one")
+        dmr.update("two")
+        dmr._save_bucket(dmr.bucket, 0)
+        cr = self.conn.cursor()
+        cr.execute("SELECT * FROM points ORDER BY value")
+        rows = cr.fetchall()
+        self.assertEquals(rows, [(dmr.metric_id, 0, "one", 2),
+                                 (dmr.metric_id, 0, "two", 1)])
+
+    def test_load_metric_id(self):
+        dmr = distinct.DistinctMetricReporter("test", dsn=self.dsn)
+        dmr._save_bucket({}, 0)
+        dmr2 = distinct.DistinctMetricReporter("test", dsn=self.dsn)
+        dmr2._save_bucket({}, 0)
+        self.assertEquals(dmr.metric_id, dmr2.metric_id)

=== modified file 'distinctdb/version.py'
--- distinctplugin/version.py	2011-12-06 19:39:21 +0000
+++ distinctdb/version.py	2011-12-09 20:08:25 +0000
@@ -1,1 +1,1 @@
-distinctplugin = "0.0.1"
\ No newline at end of file
+distinctplugin = "0.0.1"

=== modified file 'setup.py'
--- setup.py	2011-12-06 19:39:21 +0000
+++ setup.py	2011-12-09 20:08:25 +0000
@@ -41,7 +41,7 @@
         list(getPlugins(IPlugin))
 
 setup(
-    cmdclass = {'install': TxPluginInstaller},
+    cmdclass={'install': TxPluginInstaller},
     name="distinctplugin",
     version=version.distinctplugin,
     description="A txstatsd plugin for distinct counts",

=== removed file 'twisted/plugins/distinct_plugin.py'
--- twisted/plugins/distinct_plugin.py	2011-12-06 19:39:21 +0000
+++ twisted/plugins/distinct_plugin.py	1970-01-01 00:00:00 +0000
@@ -1,20 +0,0 @@
-from zope.interface import implements
-
-from twisted.plugin import IPlugin
-from txstatsd.itxstatsd import IMetricFactory
-from distinctcount.distinctmetric import DistinctMetricReporter
-
-class DistinctMetricFactory(object):
-    implements(IMetricFactory, IPlugin)
-
-    name = "distinct"
-    metric_type = "d"
-
-    def build_metric(self, prefix, name, wall_time_func=None):
-        return DistinctMetricReporter(name, prefix=prefix,
-                                      wall_time_func=wall_time_func)
-
-    def configure(self, options):
-        pass
-
-distinct_metric_factory = DistinctMetricFactory()

=== added file 'twisted/plugins/distinctdbplugin.py'
--- twisted/plugins/distinctdbplugin.py	1970-01-01 00:00:00 +0000
+++ twisted/plugins/distinctdbplugin.py	2011-12-09 20:08:25 +0000
@@ -0,0 +1,33 @@
+from zope.interface import implements
+
+from twisted.plugin import IPlugin
+from txstatsd.itxstatsd import IMetricFactory
+from distinctdb.distinctmetric import DistinctMetricReporter, ONEDAY
+
+
+class DistinctMetricFactory(object):
+    implements(IMetricFactory, IPlugin)
+
+    name = "distinct"
+    metric_type = "d"
+
+    bucket_size = None
+    dsn = None
+    metric_ids = None
+
+    def build_metric(self, prefix, name, wall_time_func=None):
+        return DistinctMetricReporter(name, prefix=prefix,
+                                      wall_time_func=wall_time_func,
+                                      bucket_size=self.bucket_size,
+                                      dsn=self.dsn)
+
+    def configure(self, options):
+        self.section = dict(options.get("plugin_distinctdb", {}))
+        try:
+            self.bucket_size = int(self.section.get("bucket_size", ONEDAY))
+        except ValueError:
+            self.bucket_size = ONEDAY
+
+        self.dsn = self.section.get("dsn", None)
+
+distinct_metric_factory = DistinctMetricFactory()