How to create ForeignKeyConstraint in SQLAlchemy

I'm trying to temporarily drop a foreign key constraint on a table to run efficient inserts, using SQLAlchemy.

I've written the following context manager to handle dropping and recreating:

from contextlib import contextmanager
from typing import Iterator

from sqlalchemy import ForeignKeyConstraint
from sqlalchemy.orm import Session
from sqlalchemy.sql.ddl import CreateConstraint, DropConstraint


@contextmanager
def drop_constraint(
    session: Session, constraint: ForeignKeyConstraint
) -> Iterator[None]:
    session.execute(DropConstraint(constraint))
    yield
    session.execute(CreateConstraint(constraint))

Dropping the constraint works well, but I get the following error when re-creating it on exit:

ArgumentError: Executable SQL or text() construct expected,
got <sqlalchemy.sql.ddl.CreateConstraint object at ...>.

Any idea how I can create this constraint on the fly with pure SQLAlchemy 2.0?

Solution

Looks like I got confused by class names. Replacing CreateConstraint by AddConstraint in the provided code will work as expected.