Georgiana.b
Georgiana.b

Reputation: 359

Dump PostgreSQL database schema with Python

I'm looking for a way to dump a PostgreSQL database schema using Python. Ideally, the result would be either a .sql dump or any other format that can later be used by SQLAlchemy to create a new database from that schema.

I made this experiment with SQLAlchemy's MetaData class:

from sqlalchemy import MetaData, create_engine

engine = create_engine(source_database_url)
test_engine = create_engine(test_database_url)
metadata = MetaData()
metadata.reflect(engine)
metadata.create_all(test_engine)

This does exactly what I want except for exporting the schema for later use. So, considering that SQLAlchemy can successfully reflect a schema and create another table based on it, I'm hoping there is a way to programmatically export it in the process.

I am aware that I could summon pg_dump and pg_restore from the code but I would like to avoid external dependencies and the troubles that come with them.

Is there a way to achieve this with SQLAlchemy or other Python libraries?

Upvotes: 5

Views: 5315

Answers (1)

Georgiana.b
Georgiana.b

Reputation: 359

After digging further I found this helpful section from SQLAlchemy docs. So, starting from that code snippet I created the following class:

import io
from sqlalchemy import create_engine, MetaData
from sqlalchemy.orm import sessionmaker

class SchemaDump(object):
    def __init__(self, db_url, schema_file_path):
        self.db_url = db_url
        self.schema_file_path = schema_file_path
        self.buf = io.BytesIO()

    def dump_shema(self):
        engine = create_engine(self.db_url)
        metadata = MetaData()
        metadata.reflect(engine)

        def dump(sql, *multiparams, **params):
            f = sql.compile(dialect=engine.dialect)
            self.buf.write(str(f).encode('utf-8'))
            self.buf.write(b';\n')

        new_engine = create_engine(self.db_url, strategy='mock', executor=dump)
        metadata.create_all(new_engine, checkfirst=True)

        with io.open(self.schema_file_path, 'wb+') as schema:
            schema.write(self.buf.getvalue())

This is still quite sketchy but the main idea is to capture the raw SQL statements returned by sql.compile(dialect=engine.dialect) in buf and write them to the file.

I wrote an equally sketchy class to restore the database from the .sql dump created by the class above:

import io
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

class RestoreSchema(object):
    def __init__(self, db_url, schema_file_path):
        self.db_url = db_url
        self.schema_file_path = schema_file_path

    def restore_schema(self):
        raw_schema = ''
        with io.open(self.schema_file_path) as sql_schema:
            raw_schema = sql_schema.read()
        engine = create_engine(self.db_url)
        Session = sessionmaker(bind=engine)
        session = Session()
        conn = engine.connect()
        transaction = conn.begin()
        try:
            conn.execute(raw_schema)
            transaction.commit()
        except Exception as e:
            transaction.rollback()
            raise e
        finally:
            session.close()

You still have to worry about whether the tables already exist etc. but the effect is exactly the same as in the code snippet from my question.

Upvotes: 9

Related Questions