From b62ddfb7eff9d2ba8b619a079f486d676e0e2445 Mon Sep 17 00:00:00 2001 From: Andrew Suffield Date: Thu, 9 Feb 2012 13:48:45 +0000 Subject: [PATCH] Support table dropping and sqlalchemy Connection objects --- sqlaload/__init__.py | 2 +- sqlaload/schema.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/sqlaload/__init__.py b/sqlaload/__init__.py index 6e695cb..85f9521 100644 --- a/sqlaload/__init__.py +++ b/sqlaload/__init__.py @@ -1,6 +1,6 @@ from sqlaload.schema import connect -from sqlaload.schema import create_table, load_table, get_table +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 diff --git a/sqlaload/schema.py b/sqlaload/schema.py index 705adfe..342d5a5 100644 --- a/sqlaload/schema.py +++ b/sqlaload/schema.py @@ -40,6 +40,10 @@ def load_table(engine, table_name): return table def get_table(engine, table_name): + # Accept Connection objects here + if hasattr(engine, 'engine'): + engine = engine.engine + if table_name in TABLES[engine]: return TABLES[engine][table_name] if engine.has_table(table_name): @@ -47,6 +51,19 @@ def get_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 + if table_name in TABLES[engine]: + table = TABLES[engine][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) + def _guess_type(sample): if isinstance(sample, bool): return Boolean