Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Overriding SQLAlchemy Columns for Testing

I'm trying to override a SQLAlchemy column to have different attributes in my test DB. Basically, I have a composite primary key in all of my MySQL tables: an autoncrement BIGINT id and a DATETIME create_time. For unit testing, I want to use SQLite, but my research has led me to believe that autoincrement is not possible for this use case. Since I don't use my create_time column for anything other than creating a rolling partition in MySQL, I want to override the two columns in my unit tests so they are compatible with SQLite. Note: I am using the Base = declarative_base() route, and all my models inherit Base and PrimaryKeyMixin.

In my models.py, I have created a mixin PrimaryKeyMixin that all tables inherit. All tables also inherit TABLE_ARGS for the __table_args__ component.

TABLE_ARGS = {
    'mysql_engine': 'InnoDB',
    'mysql_charset': 'utf8'
}

class PrimaryKeyMixin:
    id = Column(BIGINT(unsigned=True), primary_key=True, autoincrement=True)
    create_time = Column(DATETIME, primary_key=True, default=func.now())

In my unit tests, I do the following:

from myapp.models import TABLE_ARGS, PrimaryKeyMixin, Column, INT, DATETIME, func
TABLE_ARGS['sqlite_autoincrement'] = True
PrimaryKeyMixin.id = Column(INT(unsigned=True), primary_key=True, autoincrement=True)
PrimaryKeyMixin.create_date = Column(DATETIME, default=func.now())

TEST_DB = '/tmp/myapp_test.db'
from myapp import db
from myapp.models import Base
db.ENGINE = create_engine(
    'sqlite:///%s' % TEST_DB,
    encoding='utf-8',
)
Base.metadata.drop_all(db.ENGINE)
Base.metadata.create_all(db.ENGINE)

However, when I add echo=True to the ENGINE, it still is creating a PRIMARY_KEY (id, create_time) for each table.

like image 544
Matt Stern Avatar asked May 20 '26 08:05

Matt Stern


1 Answers

Did it like this:

Base.metadata = MetaData()

def copy_table(model):
    Table(model.__tablename__, Base.metadata, *change_columns(model.__table__.columns), sqlite_autoincrement=True)

def change_columns(columns):
    new_cols = [column.copy() for column in columns if column.name != 'id' and column.name !='create_time']
    new_cols.append(Column('id', INT(unsigned=True), primary_key=True, autoincrement=True))
    new_cols.append(Column('create_time', DATETIME, default=func.now(), autoincrement=True))
    return new_cols

for model in PrimaryKeyMixin.__subclasses__():
    copy_table(model)
like image 173
Matt Stern Avatar answered May 22 '26 21:05

Matt Stern