Write indexes.

sqlaload/schema.py

@@ -4,13 +4,14 @@ from collections import defaultdict
 
 from sqlalchemy import create_engine
 from sqlalchemy import Integer, UnicodeText, Float, DateTime, Boolean
-from sqlalchemy.schema import Table, MetaData, Column
+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__)
 
 TABLES = defaultdict(dict)
+INDEXES = dict()
 
 def connect(url):
     """ Create an engine for the given database URL. """
@@ -97,3 +98,18 @@ def create_column(engine, table, name, type):
         col = Column(name, type)
         col.create(table, connection=engine)
 
+def create_index(engine, table, columns, name=None):
+    if not name:
+        sig = abs(hash('||'.join(columns)))
+        name = 'ix_%s_%s' % (table.name, sig)
+    if name in INDEXES:
+        return INDEXES[name]
+    try:
+        columns = [table.c[c] for c in columns]
+        idx = Index(name, *columns)
+        idx.create(engine)
+    except:
+        idx = None
+    INDEXES[name] = idx
+    return idx
+

sqlaload/write.py

@@ -1,6 +1,7 @@
 import logging
 
 from sqlaload.schema import _ensure_columns, _args_to_clause
+from sqlaload.schema import create_index
 
 log = logging.getLogger(__name__)
 
@@ -27,6 +28,9 @@ def update_row(engine, table, row, unique, ensure=True, types={}):
         return False
 
 def upsert(engine, table, row, unique, ensure=True, types={}):
+    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)