Initial commit.

Author: arikfr

.gitignore

@@ -0,0 +1,3 @@
+.idea
+*.pyc
+rd_service/settings.py

rd_service/__init__.py

@@ -0,0 +1 @@
+

rd_service/cli.py

@@ -0,0 +1,44 @@
+import argparse
+import logging
+import urlparse
+import redis
+import time
+import settings
+import data
+
+
+def start_workers(data_manager):
+    try:
+        data_manager.start_workers(settings.WORKERS_COUNT, settings.CONNECTION_STRING,
+                                   settings.MAX_CONNECTIONS)
+        logging.info("Workers started.")
+
+        while True:
+            try:
+                data_manager.refresh_queries()
+            except Exception:
+                logging.error("Something went wrong with refreshing queries...");
+            time.sleep(60)
+    except KeyboardInterrupt:
+        logging.warning("Exiting; waiting for threads")
+        data_manager.stop_workers()
+
+
+if __name__ == '__main__':
+    channel = logging.StreamHandler()
+    logging.getLogger().addHandler(channel)
+    logging.getLogger().setLevel("DEBUG")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command")
+    args = parser.parse_args()
+
+    url = urlparse.urlparse(settings.REDIS_URL)
+    redis_connection = redis.StrictRedis(host=url.hostname, port=url.port, db=0, password=url.password)
+    data_manager = data.Manager(redis_connection, settings.INTERNAL_DB_CONNECTION_STRING, settings.MAX_CONNECTIONS)
+
+    if args.command == "worker":
+        start_workers(data_manager)
+    else:
+        print "Unknown command"
+

rd_service/data/__init__.py

@@ -0,0 +1,4 @@
+from manager import Manager
+from worker import Job
+import models
+import utils

rd_service/data/manager.py

@@ -0,0 +1,171 @@
+import collections
+from contextlib import contextmanager
+import json
+import logging
+import psycopg2
+import psycopg2.pool
+import qr
+import redis
+import time
+import query_runner
+import worker
+from utils import gen_query_hash
+
+
+class QueryResult(collections.namedtuple('QueryData', 'id query data runtime retrieved_at query_hash')):
+    def to_dict(self, parse_data=False):
+        d = self._asdict()
+
+        if parse_data and d['data']:
+            d['data'] = json.loads(d['data'])
+
+        return d
+
+
+class Manager(object):
+    def __init__(self, redis_connection, db_connection_string, db_max_connections):
+        self.redis_connection = redis_connection
+        self.workers = []
+        self.db_connection_pool = psycopg2.pool.ThreadedConnectionPool(1, db_max_connections,
+                                                                       db_connection_string)
+        self.queue = qr.PriorityQueue("jobs")
+        self.max_retries = 5
+
+    # TODO: Use our Django Models
+    def get_query_result_by_id(self, query_result_id):
+        with self.db_transaction() as cursor:
+            sql = "SELECT id, query, data, runtime, retrieved_at, query_hash FROM query_results " \
+                  "WHERE id=%s LIMIT 1"
+            cursor.execute(sql, (query_result_id,))
+            query_result = cursor.fetchone()
+
+        if query_result:
+            query_result = QueryResult(*query_result)
+
+        return query_result
+
+    def get_query_result(self, query, ttl=0):
+        query_hash = gen_query_hash(query)
+
+        with self.db_transaction() as cursor:
+            sql = "SELECT id, query, data, runtime, retrieved_at, query_hash FROM query_results " \
+                  "WHERE query_hash=%s " \
+                  "AND retrieved_at < now() at time zone 'utc' - interval '%s second'" \
+                  "ORDER BY retrieved_at DESC LIMIT 1"
+            cursor.execute(sql, (query_hash, psycopg2.extensions.AsIs(ttl)))
+            query_result = cursor.fetchone()
+
+        if query_result:
+            query_result = QueryResult(*query_result)
+
+        return query_result
+
+    def add_job(self, query, priority):
+        query_hash = gen_query_hash(query)
+        logging.info("[Manager][%s] Inserting job with priority=%s", query_hash, priority)
+        try_count = 0
+        job = None
+
+        while try_count < self.max_retries:
+            try_count += 1
+
+            pipe = self.redis_connection.pipeline()
+            try:
+                pipe.watch('query_hash_job:%s' % query_hash)
+                job_id = pipe.get('query_hash_job:%s' % query_hash)
+                if job_id:
+                    logging.info("[Manager][%s] Found existing job: %s", query_hash, job_id)
+                    job = worker.Job.load(self, job_id)
+                else:
+                    job = worker.Job(self, query, priority)
+                    pipe.multi()
+                    job.save(pipe)
+                    logging.info("[Manager][%s] Created new job: %s", query_hash, job.id)
+                    self.queue.push(job.id, job.priority)
+                break
+
+            except redis.WatchError:
+                continue
+
+        if not job:
+            logging.error("[Manager][%s] Failed adding job for query.", query_hash)
+
+        return job
+
+    def refresh_queries(self):
+        sql = """SELECT queries.query, queries.ttl, retrieved_at
+        FROM (SELECT query, min(ttl) as ttl FROM queries  WHERE ttl > 0 GROUP by query) queries
+        JOIN (SELECT query, max(retrieved_at) as retrieved_at
+        FROM query_results
+        GROUP BY query) query_results on query_results.query=queries.query
+        WHERE queries.ttl > 0
+          AND query_results.retrieved_at + ttl * interval '1 second' < now() at time zone 'utc';"""
+
+        queries = self.run_query(sql)
+        for query, ttl, retrieved_at in queries:
+            self.add_job(query, worker.Job.LOW_PRIORITY)
+
+    def store_query_result(self, query, data, run_time, retrieved_at):
+        query_result_id = None
+        query_hash = gen_query_hash(query)
+        sql = "INSERT INTO query_results (query_hash, query, data, runtime, retrieved_at) " \
+              "VALUES (%s, %s, %s, %s, %s) RETURNING id"
+        with self.db_transaction() as cursor:
+            cursor.execute(sql, (query_hash, query, data, run_time, retrieved_at))
+            if cursor.rowcount == 1:
+                query_result_id = cursor.fetchone()[0]
+                logging.info("[Manager][%s] Inserted query data; id=%s", query_hash, query_result_id)
+
+                sql = "UPDATE queries SET latest_query_data_id=%s WHERE query_hash=%s"
+                cursor.execute(sql, (query_result_id, query_hash))
+
+                logging.info("[Manager][%s] Updated %s queries.", query_hash, cursor.rowcount)
+            else:
+                logging.error("[Manager][%s] Failed inserting query data.", query_hash)
+        return query_result_id
+
+    def run_query(self, *args):
+        sql = args[0]
+        logging.debug("running query: %s %s", sql, args[1:])
+
+        with self.db_transaction() as cursor:
+            cursor.execute(sql, args[1:])
+            if cursor.description:
+                data = list(cursor)
+            else:
+                data = cursor.rowcount
+
+        return data
+
+    def start_workers(self, workers_count, connection_string, max_connections):
+        if self.workers:
+            return self.workers
+
+        # TODO: who closes the connection pool?
+        pg_connection_pool = psycopg2.pool.ThreadedConnectionPool(1, max_connections, connection_string)
+        runner = query_runner.redshift(pg_connection_pool)
+
+        self.workers = [worker.Worker(self, runner) for i in range(workers_count)]
+        for w in self.workers:
+            w.start()
+
+        return self.workers
+
+    def stop_workers(self):
+        for w in self.workers:
+            w.continue_working = False
+            w.join()
+
+    @contextmanager
+    def db_transaction(self):
+        connection = self.db_connection_pool.getconn()
+        cursor = connection.cursor()
+        try:
+            yield cursor
+        except:
+            connection.rollback()
+            raise
+        else:
+            connection.commit()
+        finally:
+            self.db_connection_pool.putconn(connection)

rd_service/data/models.py

@@ -0,0 +1,137 @@
+import json
+from django.db import models
+from django.template.defaultfilters import slugify
+import utils
+
+
+class QueryResult(models.Model):
+    id = models.AutoField(primary_key=True)
+    query_hash = models.CharField(max_length=32)
+    query = models.TextField()
+    data = models.TextField()
+    runtime = models.FloatField()
+    retrieved_at = models.DateTimeField()
+
+    class Meta:
+        app_label = 'redash'
+        db_table = 'query_results'
+
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'query_hash': self.query_hash,
+            'query': self.query,
+            'data': json.loads(self.data),
+            'runtime': self.runtime,
+            'retrieved_at': self.retrieved_at
+        }
+
+    def __unicode__(self):
+        return u"%d | %s | %s" % (self.id, self.query_hash, self.retrieved_at)
+
+
+class Query(models.Model):
+    id = models.AutoField(primary_key=True)
+    latest_query_data = models.ForeignKey(QueryResult)
+    name = models.CharField(max_length=255)
+    description = models.CharField(max_length=4096)
+    query = models.TextField()
+    query_hash = models.CharField(max_length=32)
+    ttl = models.IntegerField()
+    user = models.CharField(max_length=360)
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    class Meta:
+        app_label = 'redash'
+        db_table = 'queries'
+
+    def to_dict(self, with_result=True):
+        d = {
+            'id': self.id,
+            'latest_query_data_id': self.latest_query_data_id,
+            'name': self.name,
+            'description': self.description,
+            'query': self.query,
+            'query_hash': self.query_hash,
+            'ttl': self.ttl,
+            'user': self.user,
+            'created_at': self.created_at,
+        }
+
+        if with_result and self.latest_query_data_id:
+            d['latest_query_data'] = self.latest_query_data.to_dict()
+
+        return d
+
+    def save(self, *args, **kwargs):
+        self.query_hash = utils.gen_query_hash(self.query)
+        super(Query, self).save(*args, **kwargs)
+
+    def __unicode__(self):
+        return unicode(self.id)
+
+
+class Dashboard(models.Model):
+    id = models.AutoField(primary_key=True)
+    slug = models.CharField(max_length=140)
+    name = models.CharField(max_length=100)
+    user = models.CharField(max_length=360)
+    layout = models.TextField()
+    is_archived = models.BooleanField(default=False)
+
+    class Meta:
+        app_label = 'redash'
+        db_table = 'dashboards'
+
+    def to_dict(self, with_widgets=False):
+        layout = json.loads(self.layout)
+
+        if with_widgets:
+            widgets = {w.id: w.to_dict() for w in self.widgets.all()}
+            widgets_layout = map(lambda row: map(lambda widget_id: widgets.get(widget_id, None), row), layout)
+        else:
+            widgets_layout = None
+
+        return {
+            'id': self.id,
+            'slug': self.slug,
+            'name': self.name,
+            'user': self.user,
+            'layout': layout,
+            'widgets': widgets_layout
+        }
+
+    def save(self, *args, **kwargs):
+        # TODO: make sure slug is unique
+        if not self.slug:
+            self.slug = slugify(self.name)
+        super(Dashboard, self).save(*args, **kwargs)
+
+    def __unicode__(self):
+        return u"%s=%s" % (self.id, self.name)
+
+
+class Widget(models.Model):
+    id = models.AutoField(primary_key=True)
+    query = models.ForeignKey(Query)
+    type = models.CharField(max_length=100)
+    width = models.IntegerField()
+    options = models.TextField()
+    dashboard = models.ForeignKey(Dashboard, related_name='widgets')
+
+    class Meta:
+        app_label = 'redash'
+        db_table = 'widgets'
+
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'query': self.query.to_dict(),
+            'type': self.type,
+            'width': self.width,
+            'options': json.loads(self.options),
+            'dashboard_id': self.dashboard_id
+        }
+
+    def __unicode__(self):
+        return u"%s=>%s" % (self.id, self.dashboard_id)