Split __init__ into several modules and remove flask-peewee dependency.

This should make imports more sensible and with less side effects. Also might reduce the memory footprint of the workers.

manage.py

@@ -2,10 +2,12 @@
 """
 CLI to manage redash.
 """
-from redash import settings, app, db, models, __version__
-from redash.import_export import import_manager
 from flask.ext.script import Manager, prompt_pass
 
+from redash import settings, models, __version__
+from redash.wsgi import app
+from redash.import_export import import_manager
+
 manager = Manager(app)
 database_manager = Manager(help="Manages the database (create/drop tables).")
 users_manager = Manager(help="Users management commands.")
@@ -25,6 +27,7 @@ def runworkers():
 
 @manager.shell
 def make_shell_context():
+    from redash.models import db
     return dict(app=app, db=db, models=models)
 
 @manager.command

redash/__init__.py

@@ -1,16 +1,9 @@
-import json
-import urlparse
 import logging
-from datetime import timedelta
-from flask import Flask, make_response
-from flask.ext.restful import Api
-from flask_peewee.db import Database
+import urlparse
 import redis
 from statsd import StatsClient
-from celery import Celery
 
-import events
-from redash import settings, utils
+from redash import settings, events
 
 __version__ = '0.4.0'
 
@@ -26,48 +19,12 @@ def setup_logging():
 
 setup_logging()
 
-app = Flask(__name__,
-            template_folder=settings.STATIC_ASSETS_PATH,
-            static_folder=settings.STATIC_ASSETS_PATH,
-            static_path='/static')
-
-celery = Celery('redash',
-                broker=settings.CELERY_BROKER,
-                include='redash.tasks')
-
-celery.conf.update(CELERY_RESULT_BACKEND=settings.CELERY_BACKEND,
-                   CELERYBEAT_SCHEDULE={
-                       'refresh_queries': {
-                           'task': 'redash.tasks.refresh_queries',
-                           'schedule': timedelta(seconds=30)
-                       },
-                   },
-                   CELERY_TIMEZONE='UTC')
-
-api = Api(app)
-
-# configure our database
-settings.DATABASE_CONFIG.update({'threadlocals': True})
-app.config['DATABASE'] = settings.DATABASE_CONFIG
-db = Database(app)
-
-from redash.authentication import setup_authentication
-auth = setup_authentication(app)
-
-@api.representation('application/json')
-def json_representation(data, code, headers=None):
-    resp = make_response(json.dumps(data, cls=utils.JSONEncoder), code)
-    resp.headers.extend(headers or {})
-    return resp
-
-
 redis_url = urlparse.urlparse(settings.REDIS_URL)
 if redis_url.path:
     redis_db = redis_url.path[1]
 else:
     redis_db = 0
 
+# TODO: move this to function that create a connection?
 redis_connection = redis.StrictRedis(host=redis_url.hostname, port=redis_url.port, db=redis_db, password=redis_url.password)
 statsd_client = StatsClient(host=settings.STATSD_HOST, port=settings.STATSD_PORT, prefix=settings.STATSD_PREFIX)
-
-from redash import controllers

redash/authentication.py

@@ -5,14 +5,12 @@ import time
 import logging
 
 from flask import request, make_response, redirect, url_for
-from flask.ext.googleauth import GoogleAuth, login
 from flask.ext.login import LoginManager, login_user, current_user
+from flask.ext.googleauth import GoogleAuth, login
 from werkzeug.contrib.fixers import ProxyFix
 
-from models import AnonymousUser
 from redash import models, settings
 
-
 login_manager = LoginManager()
 logger = logging.getLogger('authentication')
 
@@ -99,7 +97,7 @@ def setup_authentication(app):
             openid_auth._OPENID_ENDPOINT = "https://www.google.com/a/%s/o8/ud?be=o8" % settings.GOOGLE_APPS_DOMAIN
 
     login_manager.init_app(app)
-    login_manager.anonymous_user = AnonymousUser
+    login_manager.anonymous_user = models.AnonymousUser
     app.wsgi_app = ProxyFix(app.wsgi_app)
     app.secret_key = settings.COOKIE_SECRET
 

redash/controllers.py

@@ -19,10 +19,9 @@ from flask_login import current_user, login_user, logout_user
 import sqlparse
 import events
 from permissions import require_permission
-from redash import settings, utils, __version__, statsd_client
 
-from redash import app, auth, api, redis_connection
-from redash import models
+from redash import redis_connection, statsd_client, models, settings, utils, __version__
+from redash.wsgi import app, auth, api
 
 import logging
 from tasks import QueryTask

redash/models.py

@@ -3,16 +3,46 @@ import hashlib
 import logging
 import time
 import datetime
-from flask.ext.peewee.utils import slugify
-from flask.ext.login import UserMixin, AnonymousUserMixin
 import itertools
-from passlib.apps import custom_app_context as pwd_context
+
 import peewee
+from passlib.apps import custom_app_context as pwd_context
 from playhouse.postgres_ext import ArrayField
-from redash import db, utils
+from flask.ext.login import UserMixin, AnonymousUserMixin
+
+from redash import utils, settings
 
 
-class BaseModel(db.Model):
+class Database(object):
+    def __init__(self):
+        self.database_config = dict(settings.DATABASE_CONFIG)
+        self.database_name = self.database_config.pop('name')
+        self.database = peewee.PostgresqlDatabase(self.database_name, **self.database_config)
+        self.app = None
+
+    def init_app(self, app):
+        self.app = app
+        self.register_handlers()
+
+    def connect_db(self):
+        self.database.connect()
+
+    def close_db(self, exc):
+        if not self.database.is_closed():
+            self.database.close()
+
+    def register_handlers(self):
+        self.app.before_request(self.connect_db)
+        self.app.teardown_request(self.close_db)
+
+
+db = Database()
+
+
+class BaseModel(peewee.Model):
+    class Meta:
+        database = db.database
+
     @classmethod
     def get_by_id(cls, model_id):
         return cls.get(cls.id == model_id)
@@ -383,11 +413,11 @@ class Dashboard(BaseModel):
 
     def save(self, *args, **kwargs):
         if not self.slug:
-            self.slug = slugify(self.name)
+            self.slug = utils.slugify(self.name)
 
             tries = 1
             while self.select().where(Dashboard.slug == self.slug).first() is not None:
-                self.slug = slugify(self.name) + "_{0}".format(tries)
+                self.slug = utils.slugify(self.name) + "_{0}".format(tries)
                 tries += 1
 
         super(Dashboard, self).save(*args, **kwargs)

redash/settings.py

@@ -5,9 +5,7 @@ import urlparse
 
 def parse_db_url(url):
     url_parts = urlparse.urlparse(url)
-    connection = {
-        'engine': 'peewee.PostgresqlDatabase',
-    }
+    connection = {'threadlocals': True}
 
     if url_parts.hostname and not url_parts.path:
         connection['name'] = url_parts.hostname

redash/tasks.py

@@ -1,12 +1,12 @@
 import time
 import datetime
 from celery.utils.log import get_task_logger
-import peewee
 import logging
 from celery.result import AsyncResult
 import redis
 from redash.data.query_runner import get_query_runner
-from redash import celery, redis_connection, models, statsd_client
+from redash import models, redis_connection, statsd_client
+from redash.worker import celery
 from redash.utils import gen_query_hash
 
 logger = get_task_logger(__name__)
@@ -77,7 +77,7 @@ class QueryTask(object):
 
     def to_dict(self):
         if self._async_result.status == 'STARTED':
-            updated_at = self._async_result.result['start_time']
+            updated_at = self._async_result.result.get('start_time', 0)
         else:
             updated_at = 0
 
@@ -169,7 +169,7 @@ def execute_query(self, query, data_source_id):
     # TODO: it is possible that storing the data will fail, and we will need to retry
     # while we already marked the job as done
     # Delete query_hash
-    redis_connection.delete('query_hash_job:%s', query_hash)
+    redis_connection.delete('query_hash_job:%s' % query_hash)
 
     if not error:
         query_result = models.QueryResult.store_result(data_source.id, query_hash, query, data, run_time, datetime.datetime.utcnow())

redash/utils.py

@@ -62,6 +62,10 @@ class SQLMetaData(object):
         return False
 
 
+def slugify(s):
+    return re.sub('[^a-z0-9_\-]+', '-', s.lower())
+
+
 def gen_query_hash(sql):
     """Returns hash of the given query after stripping all comments, line breaks and multiple
     spaces, and lower casing all text.

redash/worker.py

@@ -0,0 +1,21 @@
+from celery import Celery
+from datetime import timedelta
+from redash import settings
+
+
+celery = Celery('redash',
+                broker=settings.CELERY_BROKER,
+                include='redash.tasks')
+
+celery.conf.update(CELERY_RESULT_BACKEND=settings.CELERY_BACKEND,
+                   CELERYBEAT_SCHEDULE={
+                       'refresh_queries': {
+                           'task': 'redash.tasks.refresh_queries',
+                           'schedule': timedelta(seconds=30)
+                       },
+                   },
+                   CELERY_TIMEZONE='UTC')
+
+
+if __name__ == '__main__':
+    celery.start()

redash/wsgi.py

@@ -0,0 +1,32 @@
+import json
+from flask import Flask, make_response
+from flask.ext.restful import Api
+
+from redash import settings, utils
+from redash.models import db
+
+__version__ = '0.4.0'
+
+app = Flask(__name__,
+            template_folder=settings.STATIC_ASSETS_PATH,
+            static_folder=settings.STATIC_ASSETS_PATH,
+            static_path='/static')
+
+
+api = Api(app)
+
+# configure our database
+settings.DATABASE_CONFIG.update({'threadlocals': True})
+app.config['DATABASE'] = settings.DATABASE_CONFIG
+db.init_app(app)
+
+from redash.authentication import setup_authentication
+auth = setup_authentication(app)
+
+@api.representation('application/json')
+def json_representation(data, code, headers=None):
+    resp = make_response(json.dumps(data, cls=utils.JSONEncoder), code)
+    resp.headers.extend(headers or {})
+    return resp
+
+from redash import controllers

requirements.txt

@@ -5,11 +5,9 @@ Flask-Login==0.2.9
 passlib==1.6.2
 Jinja2==2.7.2
 MarkupSafe==0.18
-WTForms==1.0.5
 Werkzeug==0.9.4
 aniso8601==0.82
 blinker==1.3
-flask-peewee==0.6.5
 itsdangerous==0.23
 peewee==2.2.2
 psycopg2==2.5.1
@@ -20,7 +18,6 @@ requests==2.2.0
 six==1.5.2
 sqlparse==0.1.8
 wsgiref==0.1.2
-wtf-peewee==0.2.2
 Flask-Script==0.6.6
 honcho==0.5.0
 statsd==2.1.2

tests/__init__.py

@@ -1,28 +1,21 @@
 import logging
 from unittest import TestCase
-from redash import settings, db, app
-import redash.models
-
-# TODO: this isn't pretty...
+from redash import settings
 settings.DATABASE_CONFIG = {
     'name': 'circle_test',
-    'engine': 'peewee.PostgresqlDatabase',
     'threadlocals': True
 }
-app.config['DATABASE'] = settings.DATABASE_CONFIG
-db.load_database()
+
+from redash import models
 
 logging.getLogger('peewee').setLevel(logging.INFO)
 
-for model in redash.models.all_models:
-    model._meta.database = db.database
-
 
 class BaseTestCase(TestCase):
     def setUp(self):
-        redash.models.create_db(True, True)
-        redash.models.init_db()
+        models.create_db(True, True)
+        models.init_db()
 
     def tearDown(self):
-        db.close_db(None)
-        redash.models.create_db(False, True)
+        models.db.close_db(None)
+        models.create_db(False, True)

tests/test_controllers.py

@@ -8,7 +8,8 @@ from mock import patch
 from tests import BaseTestCase
 from tests.factories import dashboard_factory, widget_factory, visualization_factory, query_factory, \
     query_result_factory, user_factory, data_source_factory
-from redash import app, models, settings
+from redash import models, settings
+from redash.wsgi import app
 from redash.utils import json_dumps
 from redash.authentication import sign