|
| 1 | +import os |
1 | 2 | from typing import Optional |
2 | 3 |
|
3 | 4 | from sqlalchemy import Column, String, UniqueConstraint |
| 5 | +from sqlalchemy.exc import IntegrityError |
4 | 6 | from sqlmodel import Field, Session, SQLModel, create_engine, select |
5 | 7 |
|
6 | 8 |
|
7 | 9 | class Hero(SQLModel, table=True): |
8 | 10 | id: Optional[int] = Field(default=None, primary_key=True) |
9 | | - # Single-column unique using sa_column for full control (e.g., explicit SQL type and nullability) |
| 11 | + # Single-column unique via sa_column |
10 | 12 | email: str = Field(sa_column=Column(String(255), unique=True, nullable=False)) |
11 | 13 | name: str |
12 | 14 | secret_name: str |
13 | 15 | age: Optional[int] = None |
14 | 16 |
|
15 | | - # Composite (multi-column) unique constraint using the idiomatic SQLAlchemy approach |
| 17 | + # Composite unique constraint |
16 | 18 | __table_args__ = ( |
17 | 19 | UniqueConstraint("name", "secret_name", name="uq_hero_name_secret"), |
18 | 20 | ) |
19 | 21 |
|
20 | 22 |
|
21 | 23 | sqlite_file_name = "database_unique.db" |
22 | 24 | sqlite_url = f"sqlite:///{sqlite_file_name}" |
23 | | -engine = create_engine(sqlite_url, echo=True) |
| 25 | +engine = create_engine(sqlite_url) |
24 | 26 |
|
25 | 27 |
|
26 | 28 | def create_db_and_tables() -> None: |
| 29 | + # Reset DB for demo |
| 30 | + if os.path.exists(sqlite_file_name): |
| 31 | + os.remove(sqlite_file_name) |
27 | 32 | SQLModel.metadata.create_all(engine) |
28 | 33 |
|
29 | 34 |
|
30 | 35 | def create_heroes() -> None: |
31 | | - hero_1 = Hero( email="[email protected]", name="Ted Lasso", secret_name="Coach") |
32 | | - hero_2 = Hero( email="[email protected]", name="Roy Kent", secret_name="Roy") |
33 | | - hero_3 = Hero( |
34 | | - email="[email protected]", name="Keeley Jones", secret_name="Keeley" |
35 | | - ) |
36 | | - |
37 | 36 | with Session(engine) as session: |
38 | | - session.add(hero_1) |
39 | | - session.add(hero_2) |
40 | | - session.add(hero_3) |
| 37 | + hero_1 = Hero( email="[email protected]", name="Ted Lasso", secret_name="Coach") |
| 38 | + hero_2 = Hero( email="[email protected]", name="Roy Kent", secret_name="Roy") |
| 39 | + hero_3 = Hero( email="[email protected]", name="Keeley Jones", secret_name="Keeley") |
| 40 | + |
| 41 | + print( "Adding Hero 1: Ted Lasso ([email protected])") |
| 42 | + print( "Adding Hero 2: Roy Kent ([email protected])") |
| 43 | + print( "Adding Hero 3: Keeley Jones ([email protected])") |
| 44 | + session.add_all([hero_1, hero_2, hero_3]) |
41 | 45 | session.commit() |
| 46 | + print("Inserted 3 heroes.\n") |
| 47 | + |
| 48 | + # Duplicate (name, secret_name) should fail (different email) |
| 49 | + hero_4 = Hero( email="[email protected]", name="Roy Kent", secret_name="Roy") |
| 50 | + try: |
| 51 | + print("Attempting to insert a duplicate (name, secret_name) ...") |
| 52 | + session.add(hero_4) |
| 53 | + session.commit() |
| 54 | + except IntegrityError as e: |
| 55 | + session.rollback() |
| 56 | + print("Composite unique constraint enforced:", str(e.orig)) |
42 | 57 |
|
43 | 58 |
|
44 | 59 | def select_heroes() -> None: |
45 | 60 | with Session(engine) as session: |
| 61 | + print("\nSelecting by email (unique column):") |
46 | 62 | statement = select( Hero). where( Hero. email == "[email protected]") |
47 | 63 | hero_1 = session.exec(statement).one() |
48 | 64 | print("Hero 1:", hero_1) |
49 | 65 |
|
| 66 | + print("\nSelecting by composite key (name, secret_name):") |
50 | 67 | statement = select(Hero).where( |
51 | 68 | (Hero.name == "Roy Kent") & (Hero.secret_name == "Roy") |
52 | 69 | ) |
|
0 commit comments