Monday, January 07, 2008

Cascading DROP TABLE with SQLAlchemy

A little quirk that can get you if you're using SQLAlchemy to create and drop your database is that PostgreSQL doesn't allow you to drop a table that has other tables referring to it via FOREIGN KEY constraints. PostgreSQL has a DROP TABLE ... CASCADE command that supports this (and drops all the dependent tables) but there's no easy way to use the DROP TABLE ... CASCADE statement. It's not too hard to make it available, though.

It turns out, however, that SQLAlchemy has a nice, pluggable database dialect system that is fairly simple to update. One part of this dialect system is a "SchemaDropper". So to cascade the DROP TABLE statements, I just created the following SchemaDropper (derived from the existing PostgreSQL PGSchemaDropper) and installed it as the default PostgreSQL dialect schemadropper. (Most of the code is copied from the base SchemaDropper class in sqlalchemy.sql.compiler)


from sqlalchemy.databases import postgres

class PGCascadeSchemaDropper(postgres.PGSchemaDropper):
def visit_table(self, table):
for column in table.columns:
if column.default is not None:
self.traverse_single(column.default)
self.append("\nDROP TABLE " +
self.preparer.format_table(table) +
" CASCADE")
self.execute()

postgres.dialect.schemadropper = PGCascadeSchemaDropper


And that's it!