core utility functions from webstore

.gitignore

@@ -0,0 +1 @@
+*.pyc

sqlaload/__init__.py

@@ -1,3 +1,113 @@
+import logging
+from datetime import datetime
+
+from sqlalchemy import create_engine
+from sqlalchemy import Integer, UnicodeText, Float, DateTime
+from sqlalchemy.schema import Table, MetaData, Column
+from sqlalchemy.sql import and_, expression
+from migrate.versioning.util import construct_engine
+
+log = logging.getLogger(__name__)
+
+def connect(url):
+    """ Create an engine for the given database URL. """
+    engine = create_engine(url)
+    engine = construct_engine(engine)
+    meta = MetaData()
+    meta.bind = engine
+    engine._metadata = meta
+    return engine
+
+def create_table(engine, table_name):
+    log.debug("Creating table: %s on %r" % (table_name, engine))
+    table = Table(table_name, engine._metadata)
+    col = Column('id', Integer, primary_key=True)
+    table.append_column(col)
+    table.create(engine)
+    return table
+
+def load_table(engine, table_name):
+    return Table(table_name, engine._metadata, autoload=True)
+
+def get_table(engine, table_name):
+    if engine.has_table(table_name):
+        return load_table(engine, table_name)
+    else:
+        return create_table(engine, table_name)
+
+def _guess_type(sample):
+    if isinstance(sample, int):
+        return Integer
+    elif isinstance(sample, float):
+        return Float
+    elif isinstance(sample, datetime):
+        return DateTime
+    return UnicodeText
+
+def _ensure_columns(engine, table, row, types={}):
+    columns = set(row.keys()) - set(table.columns.keys())
+    for column in columns:
+        if column in types:
+            _type = types[column]
+        else:
+            _type = _guess_type(row[column])
+        log.debug("Creating column: %s (%s) on %r" % (column, 
+            _type, table.name))
+        col = Column(column, _type)
+        col.create(table, connection=engine)
+
+def _args_to_clause(table, args):
+    clauses = []
+    for k, v in args.items():
+        clauses.append(table.c[k] == v)
+    return and_(*clauses)
+
+
+def add_row(engine, table, row, ensure=False, types={}):
+    """ Add a row (type: dict). If ``ensure`` is set, any of 
+    the keys of the row are not table columns, they will be type
+    guessed and created. """
+    if ensure:
+        _ensure_columns(engine, table, row, types=types)
+    engine.execute(table.insert(row))
+
+
+def update_row(engine, table, row, unique, ensure=False, types={}):
+    if not len(unique):
+        return False
+    clause = dict([(u, row.get(u)) for u in unique])
+    if ensure:
+        _ensure_columns(engine, table, row, types=types)
+    try:
+        stmt = table.update(_args_to_clause(table, clause), row)
+        rp = engine.execute(stmt)
+        return rp.rowcount > 0
+    except KeyError, ke:
+        log.warn("UPDATE: filter column does not exist: %s" % ke)
+        return False
+
+def upsert(engine, table, row, unique, ensure=False, types={}):
+    if not update_row(engine, table, row, unique, ensure=ensure, types=types):
+        add_row(engine, table, row, ensure=ensure, types=types)
+
+def resultiter(rp):
+    """ SQLAlchemy ResultProxies are not iterable to get a 
+    list of dictionaries. This is to wrap them. """
+    keys = rp.keys()
+    while True:
+        row = rp.fetchone()
+        if row is None:
+            break
+        yield dict(zip(keys, row))
+
+def distinct(engine, table, *columns):
+    columns = [table.c[c] for c in columns]
+    q = expression.select(columns, distinct=True)
+    return resultiter(engine.execute(q))
+
+def all(engine, table):
+    q = table.select()
+    return resultiter(engine.execute(q))
 
 
 

test.py

@@ -0,0 +1,75 @@
+import unittest
+from datetime import datetime
+
+from sqlaload import *
+
+TEST_DATA = [
+    {
+        'date': datetime(2011, 01, 01),
+        'temperature': 1,
+        'place': 'Galway'
+    },
+    {
+        'date': datetime(2011, 01, 02),
+        'temperature': -1,
+        'place': 'Galway'
+    },
+    {
+        'date': datetime(2011, 01, 03),
+        'temperature': 0,
+        'place': 'Galway'
+    },
+    {
+        'date': datetime(2011, 01, 01),
+        'temperature': 6,
+        'place': 'Berkeley'
+    },
+    {
+        'date': datetime(2011, 01, 02),
+        'temperature': 8,
+        'place': 'Berkeley'
+    },
+    {
+        'date': datetime(2011, 01, 03),
+        'temperature': 5,
+        'place': 'Berkeley'
+    }
+    ]
+
+class TableTestCase(unittest.TestCase):
+    
+    def setUp(self):
+        self.engine = connect('sqlite:///:memory:')
+
+    def test_create_table(self):
+        table = create_table(self.engine, 'foo')
+        assert table.exists()
+        assert len(table.columns)==1, table.columns
+
+
+class QueryTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.engine = connect('sqlite:///:memory:')
+        self.table = create_table(self.engine, 'weather')
+        for entry in TEST_DATA:
+            upsert(self.engine, self.table, entry, ['place', 'date'],
+                    ensure=True)
+    
+    def test_all(self):
+        x = list(all(self.engine, self.table))
+        assert len(x)==len(TEST_DATA), len(x)
+    
+    def test_distinct(self):
+        x = list(distinct(self.engine, self.table, 'place'))
+        assert len(x)==2, x
+        p = [i['place'] for i in x]
+        assert 'Berkeley' in p, p
+        assert 'Galway' in p, p
+
+    def test_all(self):
+        x = list(all(self.engine, self.table))
+        assert len(x)==len(TEST_DATA), len(x)
+
+if __name__ == '__main__':
+    unittest.main()