from sqlalchemy import (Column, ForeignKey, Integer, String, bindparam,
                        create_engine, insert, select)
from sqlalchemy.engine.url import URL
from sqlalchemy.orm import declarative_base, relationship

Base = declarative_base()


class ProductClass(Base):
    __tablename__ = "product_class"
    id = Column(Integer, primary_key=True)
    product_class_name = Column(String(length=255), nullable=False)


class Product(Base):
    __tablename__ = "product"
    id = Column(Integer, primary_key=True)
    product_name = Column(String(length=255), nullable=False)
    product_class_id = Column(Integer, ForeignKey("product_class.id"), nullable=False)
    product_class = relationship("ProductClass")


engine = create_engine(
    URL.create(
        drivername="postgresql+psycopg2",
        username="postgres",
        password="postgres",
        host="localhost",
        port=5432,
        database="test_db",
    ),
    future=True,
    echo=True,
)

Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)


with engine.begin() as conn:
    conn.execute(
        insert(ProductClass),
        [
            {"id": 1, "product_class_name": "soda"},
            {"id": 2, "product_class_name": "candy"},
        ],
    )

# we're supposed to get product_class_id
#     1 for "coca-cola"
#     2 for "chocolate"
insert_stmnt = insert(Product).values(
    product_class_id=select(ProductClass.id)
    .where(ProductClass.product_class_name == bindparam("product_class_name"))
    .scalar_subquery()
)
insert_rows = [
    {"product_class_name": "soda", "product_name": "coca-cola"},
    {"product_class_name": "candy", "product_name": "chocolate"},
]
with engine.begin() as conn:
    conn.execute(insert_stmnt, insert_rows)

# let's check the results
select_stmnt = select(Product.product_name, ProductClass.product_class_name).join(
    ProductClass
)
with engine.begin() as conn:
    result = conn.execute(select_stmnt)
    result_rows = [
        {"product_class_name": row.product_class_name, "product_name": row.product_name}
        for row in result
    ]

if insert_rows == result_rows:
    print("Everything is fine :)")
else:
    print("Something is seriously wrong!")
    print("Rows from database:")
    print(result_rows)