from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import Index
from sqlalchemy import Integer
from sqlalchemy import String
from sqlalchemy import UniqueConstraint
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session

Base = declarative_base()


class MyTable(Base):
    __tablename__ = "my_table"

    id = Column(Integer, primary_key=True)
    name = Column(String)
    col1 = Column(String, nullable=True)
    col2 = Column(String, nullable=True)

    # fails:
    if True:
        ix_col1_col2 = Index(
            "ix_col1_col2",
            name,
            unique=True,
        )

    # works
    if False:
        ix_col1_col2 = UniqueConstraint(
            name,
            name="ix_col1_col2"
        )


e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)
Base.metadata.drop_all(e)
Base.metadata.create_all(e)

data = [
    {"name": "foo", "col1": "qux", "col2": None},
    {"name": "bar", "col1": None, "col2": "quux"},
]
column_names = data[0].keys()

stmt = insert(MyTable).values(data)
upsert = stmt.on_conflict_do_update(
    constraint=MyTable.ix_col1_col2,
    set_={
        key: val for key, val in stmt.excluded.items() if key in column_names
    },
)

with Session(e) as sess:
    sess.execute(upsert)