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

[Lucio Torre]

missing test files

Unmerged revisions

6. By Lucio Torre on 2011-12-09

added test files

5. By Lucio Torre on 2011-12-09

remove while, use INSERT .. RETURNING

4. By Lucio Torre on 2011-12-08

linted

3. By Lucio Torre on 2011-12-08

added some files

2. By Lucio Torre on 2011-12-08

almost done

1. By Lucio Torre on 2011-12-06

empty plugin skeleton

=== added file 'README'
--- README	1970-01-01 00:00:00 +0000
+++ README	2012-02-07 12:32:37 +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
+
 
=== renamed file 'README' => 'README.moved'
=== added directory 'bin'
=== added file 'bin/schema.sql'
--- bin/schema.sql	1970-01-01 00:00:00 +0000
+++ bin/schema.sql	2012-02-07 12:32:37 +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	2012-02-07 12:32:37 +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	2012-02-07 12:32:37 +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
 
=== added directory 'distinctdb'
=== added file 'distinctdb/__init__.py'
=== added file 'distinctdb/distinctmetric.py'
--- distinctdb/distinctmetric.py	1970-01-01 00:00:00 +0000
+++ distinctdb/distinctmetric.py	2012-02-07 12:32:37 +0000
@@ -0,0 +1,96 @@
+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 mesurement of the distinct numbers of items seen and the times
+    it has seen each one.
+    """
+    implements(IMetric)
+
+    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.
+        @param wall_time_func: Function for obtaining wall time.
+        @param prefix: If present, a string to prepend to the message
+            composed when C{report} is called.
+        """
+        self.name = name
+        self.wall_time_func = wall_time_func
+        if prefix:
+            prefix += '.'
+        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):
+        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):
+        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	2012-02-07 12:32:37 +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)
 
=== added file 'distinctdb/version.py'
--- distinctdb/version.py	1970-01-01 00:00:00 +0000
+++ distinctdb/version.py	2012-02-07 12:32:37 +0000
@@ -0,0 +1,1 @@
+distinctplugin = "0.0.1"
 
=== added file 'setup.py'
--- setup.py	1970-01-01 00:00:00 +0000
+++ setup.py	2012-02-07 12:32:37 +0000
@@ -0,0 +1,64 @@
+from distutils.core import setup
+from distutils.command.install import install
+import os
+
+from twisted.plugin import IPlugin, getPlugins
+
+from distinctplugin import version
+
+# If setuptools is present, use it to find_packages(), and also
+# declare our dependency on epsilon.
+extra_setup_args = {}
+try:
+    import setuptools
+    from setuptools import find_packages
+except ImportError:
+    def find_packages():
+        """
+        Compatibility wrapper.
+
+        Taken from storm setup.py.
+        """
+        packages = []
+        for directory, subdirectories, files in os.walk("distinctplugin"):
+            if '__init__.py' in files:
+                packages.append(directory.replace(os.sep, '.'))
+        return packages
+
+long_description = """
+A plugin for txstatsd to count disctinct values using redis and postgres.
+"""
+
+
+class TxPluginInstaller(install):
+    def run(self):
+        install.run(self)
+        # Make sure we refresh the plugin list when installing, so we know
+        # we have enough write permissions.
+        # see http://twistedmatrix.com/documents/current/core/howto/plugin.html
+        # "when installing or removing software which provides Twisted plugins,
+        # the site administrator should be sure the cache is regenerated"
+        list(getPlugins(IPlugin))
+
+setup(
+    cmdclass={'install': TxPluginInstaller},
+    name="distinctplugin",
+    version=version.distinctplugin,
+    description="A txstatsd plugin for distinct counts",
+    author="txStatsD Developers",
+    url="https://launchpad.net/txstatsd",
+    license="MIT",
+    packages=find_packages() + ["twisted.plugins"],
+    long_description=long_description,
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: System Administrators",
+        "Intended Audience :: Information Technology",
+        "Programming Language :: Python",
+        "Topic :: Database",
+        "Topic :: Internet :: WWW/HTTP",
+        "License :: OSI Approved :: MIT License",
+       ],
+    **extra_setup_args
+    )
 
=== renamed file 'setup.py' => 'setup.py.moved'
=== added directory 'twisted'
=== renamed directory 'twisted' => 'twisted.moved'
=== added directory 'twisted/plugins'
=== added file 'twisted/plugins/distinctdbplugin.py'
--- twisted/plugins/distinctdbplugin.py	1970-01-01 00:00:00 +0000
+++ twisted/plugins/distinctdbplugin.py	2012-02-07 12:32:37 +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()