How do I change the schema for both a table and a foreign key?

Question

I have the following simplified database access layer and two tables:

class DataAccessLayer():
    def __init__(self):
        conn_string = "mysql+mysqlconnector://root:root@localhost/"
        self.engine = create_engine(conn_string)
        Base.metadata.create_all(self.engine)
        Session = sessionmaker()
        Session.configure(bind=self.engine)
        self.session = Session()


class MatchesATP(Base):
    __tablename__ = "matches_atp"
    __table_args__ = {"schema": "belgarath", "extend_existing": True}

    ID_M = Column(Integer, primary_key=True)
    ID_T_M = Column(Integer, ForeignKey("oncourt.tours_atp.ID_T"))


class TournamentsATP(Base):
    __tablename__ = "tours_atp"
    __table_args__ = {"schema": "oncourt", "extend_existing": True}

    ID_T = Column(Integer, primary_key=True)
    NAME_T = Column(String(255))

I want to be able to switch the schema names for the two tables to test databases as follows:

belgarath to belgarath_test

oncourt to oncourt_test

I've tried adding:

self.session.connection(execution_options={"schema_translate_map": {"belgarath": belgarath, "oncourt": oncourt}})

To the bottom of DataAccessLayer and then initialising the class with two variables as follows:

def __init__(self, belgarath, oncourt):

However, when I build the following query:

dal = DataAccessLayer("belgarath_test", "oncourt_test")
query = dal.session.query(MatchesATP)
print(query)

I get the following SQL:

SELECT belgarath.matches_atp.`ID_M` AS `belgarath_matches_atp_ID_M`, belgarath.matches_atp.`ID_T_M` AS `belgarath_matches_atp_ID_T_M`
FROM belgarath.matches_atp

This is still referencing the belgarath table.

I also can't figure out a way of changing the schema of the foreign key of oncourt.tours_atp.ID_T at the same time as the tables.

Are there individual solutions or a combined solution to my issues?

Answer 1

It only took me 18 months to figure this out. Turns out I needed to add the schema_translate_map to an engine and then create the session with this engine:

from sqlalchemy import create_engine


engine = create_engine(conn_str, echo=False)
schema_engine = engine.execution_options(schema_translate_map={<old_schema_name>: <new_schema_name>})
NewSession = sessionmaker(bind=schema_engine)
session = NewSession()

All ready to roll...

Answer 2

You might wanna decorate your subclassed Base declarative model with the @declared_attr decorator.

Try this--

In a base class for your models, say __init__.py...

from sqlalchemy.ext.declarative import declarative_base, declared_attr


SCHEMA_MAIN = 'belgarath'  # figure out how you want to retrieve this
SCHEMA_TEST = 'belgarath_test'

class _Base(object):

    @declared_attr
    def __table_args__(cls):
        return {'schema': SCHEMA_MAIN}

...
Base = declarative_base(cls=_Base)
Base.metadata.schema = SCHEMA_MAIN

Now that you have a Base that subclasses _Base with the main schema already defined, all your other models will subclass Base and do the following:

from . import Base, declared_attr, SCHEMA_TEST

class TestModel(Base):
    
    @declared_attr
    def __table_args__(cls):
        return {'schema': SCHEMA_TEST}

Changing a schema for a foreign key could look like this:

class TournamentsATP(Base):
    __tablename__ = "tours_atp"
    __table_args__ = {"schema": "oncourt", "extend_existing": True}

    ID_T = Column(Integer, primary_key=True)
    NAME_T = Column(String(255))

    match_id = Column('match_id', Integer, ForeignKey(f'{__table_args__.get("schema")}.matches_atp.id'))

Where match_id is a foreign key to matches_atp.id by using the __table_args[schema] element defined at the class level via @declared_attr.

Answer 3

Assuming your goal is to:

1. have dev/test/prod schemas on a single mysql host
2. allow your ORM classes to be flexible enough to be used in three different environments without modification

Then John has you most of the way to one type of solution. You could use @declared_attr to dynamically generate __table_args__ as he has suggested.

You could also consider using something like flask-sqlalchemy that comes with a built-in solution for this:

    import os

    DB_ENV = os.getenv(DB_ENV)

    
    SQLALCHEMY_BINDS = {
        'belgarath': 'mysql+mysqlconnector://root:root@localhost/belgarath{}'.format(DB_ENV),
        'oncourt': 'mysql+mysqlconnector://root:root@localhost/oncourt{}'.format(DB_ENV)
    }
    
    class MatchesATP(Base):
        __bind_key__ = "belgarath"
        ID_M = Column(Integer, primary_key=True)
        ID_T_M = Column(Integer, ForeignKey("oncourt.tours_atp.ID_T"))
    
    
    class TournamentsATP(Base):
        __bind_key__ = "oncourt"
        ID_T = Column(Integer, primary_key=True)
        NAME_T = Column(String(255))

Basically this method allows you to create a link to a schema (a bind key), and that schema is defined at run-time via the connection string. More information at the flask-sqlalchemy link.