From 83208fa093c24fa79f6ba5df567b414f5d9afde3 Mon Sep 17 00:00:00 2001
From: Marc Aymerich <glicerinu@gmail.com>
Date: Sun, 3 May 2015 19:07:16 +0000
Subject: [PATCH] Added tasks app

---
 orchestra/contrib/tasks/bin/orchestra-beat |  5 ++--
 orchestra/utils/db.py                      | 32 ++++++++++------------
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/orchestra/contrib/tasks/bin/orchestra-beat b/orchestra/contrib/tasks/bin/orchestra-beat
index 7ef9b5a3..d1ac3f13 100755
--- a/orchestra/contrib/tasks/bin/orchestra-beat
+++ b/orchestra/contrib/tasks/bin/orchestra-beat
@@ -41,11 +41,12 @@ def get_tasks(manage):
         sys.stderr.write("I am unable to connect to the database\n")
         sys.exit(1)
     script, settings_file = sys.argv[:2]
+    enabled = 1 if 'sqlite' in settings['ENGINE'] else True
     query = (
         "SELECT c.minute, c.hour, c.day_of_week, c.day_of_month, c.month_of_year, p.id "
         "FROM djcelery_periodictask as p, djcelery_crontabschedule as c "
-        "WHERE p.crontab_id = c.id AND p.enabled = True"
-    )
+        "WHERE p.crontab_id = c.id AND p.enabled = {}"
+    ).format(enabled)
     tasks = db.run_query(conn, query)
     conn.close()
     return tasks
diff --git a/orchestra/utils/db.py b/orchestra/utils/db.py
index 9daf5a7e..f2222efd 100644
--- a/orchestra/utils/db.py
+++ b/orchestra/utils/db.py
@@ -1,5 +1,3 @@
-import ast
-
 from django import db
 
 
@@ -19,25 +17,23 @@ def close_connection(execute):
 
 def get_settings(settings_file):
     """ get db settings from settings.py file without importing """
-    settings = {}
-    with open(settings_file, 'r') as handler:
-        body = ast.parse(handler.read()).body
-        for var in body:
-            targets = getattr(var, 'targets', None)
-            if targets and targets[0].id == 'DATABASES':
-                keys = var.value.values[0].keys
-                values = var.value.values[0].values
-                for key, value in zip(keys, values):
-                    if key.s == 'ENGINE':
-                        if value.s not in ('django.db.backends.sqlite3', 'django.db.backends.postgresql_psycopg2'):
-                            raise ValueError("%s engine not supported." % value)
-                    settings[key.s] = getattr(value, 's', None)
-                return settings
+    settings = {'__file__': settings_file}
+    with open(settings_file) as f:
+        __file__ = 'rata'
+        exec(f.read(), settings)
+    settings = settings['DATABASES']['default']
+    if settings['ENGINE'] not in  ('django.db.backends.sqlite3', 'django.db.backends.postgresql_psycopg2'):
+        raise ValueError("%s engine not supported." % settings['ENGINE'])
+    return settings
 
 
 def get_connection(settings):
-    import psycopg2
-    conn = psycopg2.connect("dbname='{NAME}' user='{USER}' host='{HOST}' password='{PASSWORD}'".format(**settings))
+    if settings['ENGINE'] == 'django.db.backends.sqlite3':
+        import sqlite3
+        return sqlite3.connect(settings['NAME'])
+    elif settings['ENGINE'] == 'django.db.backends.postgresql_psycopg2':
+        import psycopg2
+        return psycopg2.connect("dbname='{NAME}' user='{USER}' host='{HOST}' password='{PASSWORD}'".format(**settings))
     return conn