Use peewee in data.Manager.

redash/__init__.py

@@ -5,8 +5,7 @@ from flask.ext.restful import Api
 from flask_peewee.db import Database
 
 import redis
-from redash import settings
-from redash.data import utils
+from redash import settings, utils
 
 
 app = Flask(__name__,
@@ -18,10 +17,8 @@ app = Flask(__name__,
 api = Api(app)
 
 # configure our database
-app.config['DATABASE'] = {
-    'name': 'postgres',
-    'engine': 'peewee.PostgresqlDatabase',
-}
+settings.DATABASE_CONFIG.update({'threadlocals': True})
+app.config['DATABASE'] = settings.DATABASE_CONFIG
 db = Database(app)
 
 from redash.authentication import setup_authentication
@@ -36,6 +33,8 @@ def json_representation(data, code, headers=None):
 
 redis_url = urlparse.urlparse(settings.REDIS_URL)
 redis_connection = redis.StrictRedis(host=redis_url.hostname, port=redis_url.port, db=0, password=redis_url.password)
-data_manager = data.Manager(redis_connection, settings.INTERNAL_DB_CONNECTION_STRING, settings.MAX_CONNECTIONS)
+
+from redash import data
+data_manager = data.Manager(redis_connection, db)
 
 from redash import controllers

redash/controllers.py

@@ -16,8 +16,7 @@ from flask import g, render_template, send_from_directory, make_response, reques
 from flask.ext.restful import Resource, abort
 
 import sqlparse
-from redash import settings
-from redash.data import utils
+from redash import settings, utils
 from redash import data
 
 from redash import app, auth, api, redis_connection, data_manager

redash/data/__init__.py

@@ -1,3 +1,2 @@
 from manager import Manager
 from worker import Job
-import utils

redash/data/manager.py

@@ -1,18 +1,17 @@
 """
 Data manager. Used to manage and coordinate execution of queries.
 """
-import collections
 from contextlib import contextmanager
+import collections
 import json
+import time
 import logging
 import psycopg2
 import qr
 import redis
-import time
-import worker
-import settings
-from utils import gen_query_hash
-
+from redash import settings
+from redash.data import worker
+from redash.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):
@@ -25,10 +24,10 @@ class QueryResult(collections.namedtuple('QueryData', 'id query data runtime ret
 
 
 class Manager(object):
-    def __init__(self, redis_connection, db_connection_string, db_max_connections):
+    def __init__(self, redis_connection, db):
         self.redis_connection = redis_connection
+        self.db = db
         self.workers = []
-        self.db_connection_string = db_connection_string
         self.queue = qr.PriorityQueue("jobs", **self.redis_connection.connection_pool.connection_kwargs)
         self.max_retries = 5
         self.status = {
@@ -176,17 +175,18 @@ class Manager(object):
 
     @contextmanager
     def db_transaction(self):
-        connection = psycopg2.connect(self.db_connection_string)
-        cursor = connection.cursor()
+        self.db.connect_db()
+
+        cursor = self.db.database.get_cursor()
         try:
             yield cursor
         except:
-            connection.rollback()
+            self.db.database.rollback()
             raise
         else:
-            connection.commit()
+            self.db.database.commit()
         finally:
-            connection.close()
+            self.db.close_db(None)
 
     def _save_status(self):
         self.redis_connection.hmset('manager:status', self.status)

redash/data/query_runner.py

@@ -8,10 +8,12 @@ query language (for example: HiveQL).
 """
 import logging
 import json
-import psycopg2
 import sys
 import select
-from .utils import JSONEncoder
+
+import psycopg2
+
+from redash.utils import JSONEncoder
 
 def redshift(connection_string):
     def column_friendly_name(column_name):
@@ -64,4 +66,4 @@ def redshift(connection_string):
  
         return json_data, error
  
-    return query_runner
+    return query_runner

redash/data/worker.py

@@ -11,7 +11,7 @@ import time
 import signal
 import setproctitle
 import redis
-from utils import gen_query_hash
+from redash.utils import gen_query_hash
 
 
 class Job(object):

redash/models.py

@@ -3,9 +3,8 @@ import hashlib
 import time
 import datetime
 from flask.ext.peewee.utils import slugify
-from redash.data import utils
 import peewee
-from redash import db
+from redash import db, utils
 
 
 #class User(db.Model):

redash/settings_example.py

@@ -9,8 +9,11 @@ CONNECTION_ADAPTER = "mysql"
 # -- example mysql CONNECTION_STRING = "Server=;User=;Pwd=;Database="
 # -- example pg    CONNECTION_STRING = "user= password= host= port=5439 dbname="
 CONNECTION_STRING = "user= password= host= port=5439 dbname="
-# Connection string for the operational databases (where we store the queries, results, etc)
-INTERNAL_DB_CONNECTION_STRING = "dbname=postgres"
+# Connection settings for re:dash's own database (where we store the queries, results, etc)
+DATABASE_CONFIG = {
+    'name': 'postgres',
+    'engine': 'peewee.PostgresqlDatabase',
+}
 # Google Apps domain to allow access from; any user with email in this Google Apps will be allowed
 # access
 GOOGLE_APPS_DOMAIN = ""