[TIL] insert시 ForeignKey에 제한 걸기: DB Trigger 사용

Django - python - postgresql을 이용하여 작업을 진행하던 중, 발생했던 문제점을 해결한 기록이다.

 

 

문제상황

아래처럼 새로운 테이블(모델)을 생성하고자 했다.

class IWantNewModel:
    old_model = ForeignKey(
        "models.OldModel",
        on_delete=CASCADE,
		...
        )

이 때, FK가 되는 old_model의 타입을 특정 타입이 되도록 제한하고 싶었다. DB 자체에서 제한을 걸면 나중에 서비스단 로직에서 굳이 validation을 복잡하게 할 필요가 없기 때문이다. 한층 더 추상화되는 점은 있다만, 그래도 편리할 수 있다.

 

Django에서는 이런 경우(FK가 아닌 경우에는) CheckConstraint 조건을 주어서 DB에서 제한을 걸도록 한다.

- https://docs.djangoproject.com/en/4.2/ref/models/constraints/#checkconstraint

 

그런데 FK의 속성을 제한하는 기능은 제공되지 않는다고한다.

- https://stackoverflow.com/questions/60254433/error-using-checkconstraint-in-model-meta-along-with-django-genericforeignkey

 

다른 글에서도 찾았었는데, 특정 FK 컬럼과 관련된 외래 모델의 속성을 제한하는 기능은 지원하지 않는다고 한다.

 

 

해결방법

 

DB Trigger를 이용하면 된다! DB Trigger란 테이블에 데이터를 INSERT, UPDATE, DELETE 등 DML(Data Manipulation Language) 작업이 있을 때마다 자동으로 특정 명령을 실행하게 하는 방법이다. 찾아보면 여러 좋은 참조 자료들이 많다.

https://limkydev.tistory.com/154

 

- Transaction에 묶여서 ROLLBACK도 그대로 적용된다.

- 주로 통계자료, 무결성 확보 등에 사용된다.

 

새로운 테이블을 만드는 migration을 적용한 뒤, 추가적으로 제한 조건을 걸기 위해서 empty migration 파일을 만들었다. 그리고 대략적으로 아래와 같은 방식의 코드를 적용하고 실행하면 된다. rollback을 위해 drop하는 코드도 함께 들어있다. 모델명, 테이블명은 정확하지 않다. SQL 문법상 NEW, END IF; 등 궁금한 점이 있다면 위 참조 링크나 다른 곳에서 찾아보면 된다.

def create_check_function(apps, schema_editor):
    sql = """
        CREATE OR REPLACE FUNCTION check_old_model_id_validity() RETURNS TRIGGER AS $$
        BEGIN
            IF NEW.old_model_id IS NOT NULL THEN
                IF NOT EXISTS (
                    SELECT 1 FROM model WHERE id = NEW.old_model_id AND type = '{}'
                ) THEN
                    RAISE EXCEPTION 'Invalid old_model_id';
                END IF;
            END IF;
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """.format(
        OldModel.Type.SPECIFIC_TYPE
    )

    schema_editor.execute(sql)


def drop_check_function(apps, schema_editor):
    sql = "DROP FUNCTION IF EXISTS check_old_model_id_validity"
    schema_editor.execute(sql)


def add_constraints(apps, schema_editor):
    sql = (
        "ALTER TABLE IWANTNEWMODEL "
        "ADD CONSTRAINT old_model_id_foreign_key "
        "FOREIGN KEY (old_model_id) "
        "REFERENCES model (id)"
    )
    schema_editor.execute(sql)

    sql = (
        "CREATE TRIGGER old_model_id_function_constraint "
        "BEFORE INSERT OR UPDATE ON IWANTNEWMODEL "
        "FOR EACH ROW EXECUTE FUNCTION check_old_model_id_validity()"
    )
    schema_editor.execute(sql)


def drop_constraints(apps, schema_editor):
    sql = "DROP TRIGGER IF EXISTS old_model_id_function_constraint ON IWANTNEWMODEL"
    schema_editor.execute(sql)

    sql = "ALTER TABLE IWANTNEWMODEL " "DROP CONSTRAINT old_model_id_foreign_key"
    schema_editor.execute(sql)


class Migration(migrations.Migration):

    dependencies = [
        (...),
    ]

    operations = [
        migrations.RunPython(create_check_function, reverse_code=drop_check_function),
        migrations.RunPython(add_constraints, reverse_code=drop_constraints),
    ]