From 66bfbf734b946b363f410efb57872e628cc03cd8 Mon Sep 17 00:00:00 2001
From: Mike Korobov <kmike84@gmail.com>
Date: Sat, 25 Jun 2011 02:19:13 +0600
Subject: [PATCH] PostgreSQL support. Fix #1.

---
 .hgignore           |  1 +
 qsstats/__init__.py | 24 +++++++++++++++++++-----
 qsstats/utils.py    | 11 ++++++++++-
 runtests.py         | 10 +++++++++-
 setup.py            |  5 +++++
 5 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/.hgignore b/.hgignore
index 37020b8..876161c 100644
--- a/.hgignore
+++ b/.hgignore
@@ -27,3 +27,4 @@ build/
 dist/
 .build/
 MANIFEST
+django_qsstats_magic.egg-info
diff --git a/qsstats/__init__.py b/qsstats/__init__.py
index 56f3acf..9394192 100644
--- a/qsstats/__init__.py
+++ b/qsstats/__init__.py
@@ -24,6 +24,18 @@ def __init__(self, qs=None, date_field=None, aggregate=None, today=None):
         self.aggregate = aggregate or Count('id')
         self.today = today or self.update_today()
 
+    def _guess_engine(self):
+        if hasattr(self.qs, 'db'): # django 1.2+
+            engine_name = settings.DATABASES[self.qs.db]['ENGINE']
+        else:
+            engine_name = settings.DATABASE_ENGINE
+        if 'mysql' in engine_name:
+            return 'mysql'
+        if 'postg' in engine_name: #postgres, postgis
+            return 'postgresql'
+        if 'sqlite' in engine_name:
+            return 'sqlite'
+
     # Aggregates for a specific period of time
 
     def for_interval(self, interval, dt, date_field=None, aggregate=None):
@@ -44,15 +56,14 @@ def __getattr__(self, name):
             return partial(self.this_interval, name[5:])
         raise AttributeError
 
-
     def time_series(self, start, end=None, interval='days',
-                    date_field=None, aggregate=None, engine='mysql'):
+                    date_field=None, aggregate=None, engine=None):
         ''' Aggregate over time intervals '''
         end = end or self.today
         args = [start, end, interval, date_field, aggregate]
+        engine = engine or self._guess_engine()
         sid = transaction.savepoint()
         try:
-            #TODO: engine should be guessed
             return self._fast_time_series(*(args+[engine]))
         except (QuerySetStatsError, DatabaseError,):
             transaction.savepoint_rollback(sid)
@@ -75,10 +86,11 @@ def _slow_time_series(self, start, end, interval='days',
         return stat_list
 
     def _fast_time_series(self, start, end, interval='days',
-                          date_field=None, aggregate=None, engine='mysql'):
+                          date_field=None, aggregate=None, engine=None):
         ''' Aggregate over time intervals using just 1 sql query '''
         date_field = date_field or self.date_field
         aggregate = aggregate or self.aggregate
+        engine = engine or self._guess_engine()
 
         start, _ = get_bounds(start, interval.rstrip('s'))
         _, end = get_bounds(end, interval.rstrip('s'))
@@ -89,7 +101,9 @@ def _fast_time_series(self, start, end, interval='days',
                         filter(**kwargs).order_by().values('d').\
                         annotate(agg=aggregate)
 
-        data = dict((parse(item['d'], yearfirst=True), item['agg']) for item in aggregate_data)
+        def to_dt(d): # leave dates as-is
+            return parse(d, yearfirst=True) if isinstance(d, basestring) else d
+        data = dict((to_dt(item['d']), item['agg']) for item in aggregate_data)
 
         stat_list = []
         dt = start
diff --git a/qsstats/utils.py b/qsstats/utils.py
index 738672a..ba2a532 100644
--- a/qsstats/utils.py
+++ b/qsstats/utils.py
@@ -53,13 +53,22 @@ def get_interval_sql(date_field, interval, engine):
             'weeks': "DATE_FORMAT(DATE_SUB(`"+date_field+"`, INTERVAL(WEEKDAY(`"+date_field+"`)) DAY), '%%Y-%%m-%%d')",
             'months': "DATE_FORMAT(`" + date_field +"`, '%%Y-%%m-01')",
             'years': "DATE_FORMAT(`" + date_field +"`, '%%Y-01-01')",
+        },
+        'postgresql': {
+            'minutes': "date_trunc('minute', %s)" % date_field,
+            'hours': "date_trunc('hour', %s)" % date_field,
+            'days': "date_trunc('day', %s)" % date_field,
+            'weeks': "date_trunc('week', %s)" % date_field,
+            'months': "date_trunc('month', %s)" % date_field,
+            'years': "date_trunc('year', %s)" % date_field,
         }
     }
 
     try:
         engine_sql = SQL[engine]
     except KeyError:
-        raise UnsupportedEngine('%s DB engine is not supported' % engine)
+        msg = '%s DB engine is not supported. Supported engines are: %s' % (engine, ", ".join(SQL.keys()))
+        raise UnsupportedEngine(msg)
 
     try:
         return engine_sql[interval]
diff --git a/runtests.py b/runtests.py
index 895ac32..4610710 100755
--- a/runtests.py
+++ b/runtests.py
@@ -3,10 +3,18 @@
 
 from django.conf import settings
 from django.core.management import call_command
+import sys
+
+engine = sys.argv[1]
 
 settings.configure(
     INSTALLED_APPS=('qsstats', 'django.contrib.auth', 'django.contrib.contenttypes'),
-    DATABASE_ENGINE = 'sqlite3',
+    DATABASES = {
+        'default': {
+            'ENGINE': 'django.db.backends.' + engine,
+            'NAME': 'test'
+        }
+    }
 )
 
 if __name__ == "__main__":
diff --git a/setup.py b/setup.py
index 3f42347..222f6b3 100755
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,11 @@
 #!/usr/bin/env python
 from distutils.core import setup
 
+for cmd in ('egg_info', 'develop'):
+    import sys
+    if cmd in sys.argv:
+        from setuptools import setup
+
 setup(
     name='django-qsstats-magic',
     version='0.5.2',