autoincrement a composite primary key

테스트 코드가 없는 프로젝트에 테스트 코드를 작성하려고 하니 만만치 않았는데, 여러 문제 중 하나가 ORM을 만들어놓고 create table 등의 raw query로 테이블 관리를 한다는 것이었다. 나름의 이유가 있어서 그렇게 한다고 하지만... ORM 모델, 개발 환경 테이블, 운영 환경 테이블이 제각각이였다. 당장 ORM 모델을 적용하기엔 무리라서 일단 테스트 코드에서 문제 없이 코드가 돌아가도록 수정할 필요가 있었다.

• 복합 기본키를 구성하는 컬럼에 auto_increment를 바로 적용할 수 없고 alter table로 수정돼서 적용된다. create_all()을 실행하고나서 alter table을 또 적용해주는 게 어렵다고 판단했다. CustomAutoIncrementColumn을 생성해서 적용했다.

    if os.getenv("FLASK_ENV") == "test":
        Base.metadata.create_all(
            bind=engine, tables=[table.__table__ for table in TO_BE_CREATED_TABLES]
        )

• DateTime 컬럼에 server_default 속성에 FetchedValue()가 적용되는데, 테스트 코드에서 테이블을 생성하기 때문에 테이블의 기본 값을 가져와서 적용할 수 없는 것 같다. 이 부분은 func.now()를 적용해서 현재 시간을 기본 값으로 넣어줬다.

# 실행하면
CREATE TABLE log (
  log_data datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
  log_id bigint(20) unsigned NOT NULL AUTO_INCREMENT,
 
  PRIMARY KEY (log_data,log_id),
  KEY log_id (log_id)
)

# DDL alter table이 있음
-- auto-generated definition
create table log
(
    log_data datetime default CURRENT_TIMESTAMP not null,
    log_id   bigint unsigned,
    primary key (log_data, log_id)
);

# key 지정을 위해 index 적용함. : sqlalchemy (1075, 'Incorrect table definition; there can be only one auto column and it must be defined as a key')
create index log_id
    on log (log_id);

alter table log
    modify log_id bigint unsigned auto_increment;

---

Column을 상속받은 custom column에서 환경에 따라 분기했다.

import os
from sqlalchemy import Column, func
from sqlalchemy.sql.schema import FetchedValue
from app.main.utils.model import get_random_id


class CustomAutoIncrementColumn(Column):
    def __init__(self, *args, **kwargs):
        if os.getenv("FLASK_ENV") != "test":
            kwargs.setdefault("index", True)
            kwargs.setdefault("autoincrement", True)
        else:
            kwargs.setdefault("default", get_random_id())

        super().__init__(*args, **kwargs)


class CustomDateTimeNotNullableColumn(Column):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault("nullable", False)

        if os.getenv("FLASK_ENV") != "test":
            kwargs.setdefault(
                "server_default",
                FetchedValue(),
            )
        else:
            kwargs.setdefault("default", func.now())
            kwargs.setdefault(
                "server_default",
                func.now(),
            )

        super().__init__(*args, **kwargs)


class CustomDateTimeNullableColumn(Column):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault("nullable", True)

        if os.getenv("FLASK_ENV") != "test":
            kwargs.setdefault(
                "server_default",
                FetchedValue(),
            )
        else:
            kwargs.setdefault("default", func.now())

        super().__init__(*args, **kwargs)


class CustomIntegerColumn(Column):
    def __init__(self, *args, **kwargs):
        if os.getenv("FLASK_ENV") != "test":
            kwargs.setdefault(
                "server_default",
                FetchedValue(),
            )
            kwargs.setdefault("nullable", False)
        else:
            kwargs.setdefault("default", 0)

        super().__init__(*args, **kwargs)


class CustomStringColumn(Column):
    def __init__(self, *args, **kwargs):
        if os.getenv("FLASK_ENV") != "test":
            kwargs.setdefault(
                "server_default",
                FetchedValue(),
            )
            kwargs.setdefault("nullable", False)
        else:
            kwargs.setdefault("default", "N")

        super().__init__(*args, **kwargs)

---

model.py

class Log(Base):
    __tablename__ = "log"
    __table_args__ = (
        PrimaryKeyConstraint("log_date", "log_id", name="log_pk"),
    )
    
    log_date = CustomDateTimeNotNullableColumn(DateTime)
    log_id = CustomAutoIncrementColumn(BIGINT(unsigned=True))

결론

이미 만들어진 코드, 구조를 수정하는건 쉽지 않다
처음부터 ORM 모델 + alembic같은 db migration 도구를 사용하자...
테스트 코드부터 짜자