diff --git a/sqlaload/schema.py b/sqlaload/schema.py index 693d1c5..5a70b76 100644 --- a/sqlaload/schema.py +++ b/sqlaload/schema.py @@ -12,9 +12,6 @@ from migrate.versioning.util import construct_engine log = logging.getLogger(__name__) lock = RLock() -TABLES = defaultdict(dict) -INDEXES = dict() - def connect(url): """ Create an engine for the given database URL. """ kw = {} @@ -27,6 +24,8 @@ def connect(url): meta = MetaData() meta.bind = engine engine._metadata = meta + engine._tables = dict() + engine._indexes = dict() return engine def create_table(engine, table_name): @@ -36,14 +35,14 @@ def create_table(engine, table_name): col = Column('id', Integer, primary_key=True) table.append_column(col) table.create(engine) - TABLES[engine][table_name] = table + 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) - TABLES[engine][table_name] = table + engine._tables[table_name] = table return table def get_table(engine, table_name): @@ -52,8 +51,8 @@ def get_table(engine, table_name): engine = engine.engine with lock: - if table_name in TABLES[engine]: - return TABLES[engine][table_name] + if table_name in engine._tables: + return engine._tables[table_name] if engine.has_table(table_name): return load_table(engine, table_name) else: @@ -65,14 +64,14 @@ def drop_table(engine, table_name): engine = engine.engine with lock: - if table_name in TABLES[engine]: - table = TABLES[engine][table_name] + 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) - TABLES[engine].pop(table_name, None) + engine._tables.pop(table_name, None) def _guess_type(sample): if isinstance(sample, bool): @@ -113,14 +112,14 @@ 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] + 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 - INDEXES[name] = idx + engine._indexes[name] = idx return idx