Merge branch 'refs/heads/master' into HEAD

dataset/__init__.py

@@ -0,0 +1,12 @@
+from sqlaload.schema import connect
+from sqlaload.schema import create_table, load_table, get_table, drop_table
+from sqlaload.schema import create_column
+from sqlaload.write import add_row, update_row
+from sqlaload.write import upsert, update, delete
+from sqlaload.query import distinct, resultiter, all, find_one, find, query
+from sqlaload.db import create
+
+# shut up useless SA warning:
+import warnings
+warnings.filterwarnings('ignore', 'Unicode type received non-unicode bind param value.')
+

dataset/query.py

@@ -0,0 +1,79 @@
+import logging
+from itertools import count
+
+from sqlalchemy.sql import expression, and_
+from sqlaload.schema import _ensure_columns, get_table
+
+log = logging.getLogger(__name__)
+
+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 find_one(engine, table, **kw):
+    table = get_table(engine, table)
+    res = list(find(engine, table, _limit=1, **kw))
+    if not len(res):
+        return None
+    return res[0]
+
+def find(engine, table, _limit=None, _step=5000, _offset=0,
+         order_by='id', **kw):
+    table = get_table(engine, table)
+    _ensure_columns(engine, table, kw)
+    order_by = [table.c[order_by].asc()]
+
+    qargs = []
+    try:
+        for col, val in kw.items():
+            qargs.append(table.c[col]==val)
+    except KeyError:
+        return
+
+    for i in count():
+        qoffset = _offset + (_step * i)
+        qlimit = _step
+        if _limit is not None:
+            qlimit = min(_limit-(_step*i), _step)
+        if qlimit <= 0:
+            break
+        q = table.select(whereclause=and_(*qargs), limit=qlimit,
+                offset=qoffset, order_by=order_by)
+        #print q
+        rows = list(resultiter(engine.execute(q)))
+        if not len(rows):
+            return 
+        for row in rows:
+            yield row
+
+def query(engine, query):
+    for res in resultiter(engine.execute(query)):
+        yield res
+
+def distinct(engine, table, *columns, **kw):
+    table = get_table(engine, table)
+    qargs = []
+    try:
+        columns = [table.c[c] for c in columns]
+        for col, val in kw.items():
+            qargs.append(table.c[col]==val)
+    except KeyError:
+        return []
+
+    q = expression.select(columns, distinct=True,
+            whereclause=and_(*qargs),
+            order_by=[c.asc() for c in columns])
+    return list(resultiter(engine.execute(q)))
+
+def all(engine, table):
+    return find(engine, table)
+
+
+
+

dataset/schema.py

@@ -0,0 +1,130 @@
+import logging
+from datetime import datetime
+from collections import defaultdict
+from threading import RLock
+
+from sqlalchemy import create_engine
+from sqlalchemy import Integer, UnicodeText, Float, DateTime, Boolean
+from sqlalchemy.schema import Table, MetaData, Column, Index
+from sqlalchemy.sql import and_, expression
+from migrate.versioning.util import construct_engine
+
+log = logging.getLogger(__name__)
+lock = RLock()
+
+def connect(url):
+    """ Create an engine for the given database URL. """
+    kw = {}
+    if url.startswith('postgres'):
+        #kw['pool_size'] = 5
+        from sqlalchemy.pool import NullPool
+        kw['poolclass'] = NullPool
+    engine = create_engine(url, **kw)
+    engine = construct_engine(engine)
+    meta = MetaData()
+    meta.bind = engine
+    engine._metadata = meta
+    engine._tables = dict()
+    engine._indexes = dict()
+    return engine
+
+def create_table(engine, table_name):
+    with lock:
+        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)
+        engine._tables[table_name] = table
+        return table
+
+def load_table(engine, table_name):
+    with lock:
+        log.debug("Loading table: %s on %r" % (table_name, engine))
+        table = Table(table_name, engine._metadata, autoload=True)
+        engine._tables[table_name] = table
+        return table
+
+def get_table(engine, table_name):
+    if isinstance(table_name, Table):
+        return table_name
+
+    # Accept Connection objects here
+    if hasattr(engine, 'engine'):
+        engine = engine.engine
+    
+    with lock:
+        if table_name in engine._tables:
+            return engine._tables[table_name]
+        if engine.has_table(table_name):
+            return load_table(engine, table_name)
+        else:
+            return create_table(engine, table_name)
+
+def drop_table(engine, table_name):
+    # Accept Connection objects here
+    if hasattr(engine, 'engine'):
+        engine = engine.engine
+
+    with lock:
+        if table_name in engine._tables:
+            table = engine._tables[table_name]
+        elif engine.has_table(table_name):
+            table = Table(table_name, engine._metadata)
+        else:
+            return
+        table.drop(engine)
+        engine._tables.pop(table_name, None)
+
+def _guess_type(sample):
+    if isinstance(sample, bool):
+        return Boolean
+    elif 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))
+        create_column(engine, table, column, _type)
+
+def _args_to_clause(table, args):
+    clauses = []
+    for k, v in args.items():
+        clauses.append(table.c[k] == v)
+    return and_(*clauses)
+
+def create_column(engine, table, name, type):
+    table = get_table(engine, table)    
+    with lock:
+        if name not in table.columns.keys():
+            col = Column(name, type)
+            col.create(table, connection=engine)
+
+def create_index(engine, table, columns, name=None):
+    table = get_table(engine, table)    
+    with lock:
+        if not name:
+            sig = abs(hash('||'.join(columns)))
+            name = 'ix_%s_%s' % (table.name, sig)
+        if name in engine._indexes:
+            return engine._indexes[name]
+        try:
+            columns = [table.c[c] for c in columns]
+            idx = Index(name, *columns)
+            idx.create(engine)
+        except:
+            idx = None
+        engine._indexes[name] = idx
+        return idx
+

dataset/write.py

@@ -0,0 +1,64 @@
+import logging
+
+from sqlaload.schema import _ensure_columns, _args_to_clause
+from sqlaload.schema import create_index, get_table
+
+log = logging.getLogger(__name__)
+
+def add_row(engine, table, row, ensure=True, 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. """
+    table = get_table(engine, table)
+    if ensure:
+        _ensure_columns(engine, table, row, types=types)
+    engine.execute(table.insert(row))
+
+def update_row(engine, table, row, unique, ensure=True, types={}):
+    if not len(unique):
+        return False
+    table = get_table(engine, table)
+    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: '%s' filter column does not exist: %s", table.name, ke)
+        return False
+
+def upsert(engine, table, row, unique, ensure=True, types={}):
+    table = get_table(engine, table)
+    if ensure:
+        create_index(engine, table, unique)
+
+    if not update_row(engine, table, row, unique, ensure=ensure, types=types):
+        add_row(engine, table, row, ensure=ensure, types=types)
+
+def update(engine, table, criteria, values, ensure=True, types={}):
+    table = get_table(engine, table)
+    if ensure:
+        _ensure_columns(engine, table, values, types=types)
+    q = table.update().values(values)
+    for column, value in criteria.items():
+        q = q.where(table.c[column]==value)
+    engine.execute(q)
+
+def delete(engine, table, **kw):
+    table = get_table(engine, table)
+    _ensure_columns(engine, table, kw)
+
+    qargs = []
+    try:
+        for col, val in kw.items():
+            qargs.append(table.c[col]==val)
+    except KeyError:
+        return
+
+    q = table.delete()
+    for k, v in kw.items():
+        q= q.where(table.c[k]==v)
+    engine.execute(q)
+