spliting models

app/main.py

@@ -1,10 +1,10 @@
 from alembic import command
 from alembic.config import Config
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, Depends, HTTPException
 from sqlmodel import SQLModel, Session, select
 from typing import List
-from .database import engine
-from .models import Author, Book
+from .database import engine, get_session
+from .models import Author, AuthorBase, Book, BookBase
 
 alembic_cfg = Config("alembic.ini")
 app = FastAPI(
@@ -38,37 +38,33 @@ async def hello_world():
 
 # Create an author
 @app.post("/authors/", response_model=Author, tags=["authors"])
-def create_author(author: Author):
-    with Session(engine) as session:
-        session.add(author)
-        session.commit()
-        session.refresh(author)
-        return author
-
-# Read authors
-@app.get("/authors/", response_model=List[Author], tags=["authors"])
-def read_authors():
-    with Session(engine) as session:
-        authors = session.exec(select(Author)).all()
-        return authors
-
-# Update an author
-@app.put("/authors/{author_id}", response_model=Author, tags=["authors"])
-def update_author(author_id: int, author: Author):
-    with Session(engine) as session:
-        db_author = session.get(Author, author_id)
-        if not db_author:
-            raise HTTPException(status_code=404, detail="Author not found")
-        db_author.name = author.name
+def create_author(author: AuthorBase, session: Session = Depends(get_session)):
+    db_author = Author(name=author.name)
     session.add(db_author)
     session.commit()
     session.refresh(db_author)
     return db_author
 
+# Read authors
+@app.get("/authors/", response_model=List[Author], tags=["authors"])
+def read_authors(session: Session = Depends(get_session)):
+    authors = session.exec(select(Author)).all()
+    return authors
+
+# Update an author
+@app.put("/authors/{author_id}", response_model=Author, tags=["authors"])
+def update_author(author_id: int, author: AuthorBase, session: Session = Depends(get_session)):
+    db_author = session.get(Author, author_id)
+    if not db_author:
+        raise HTTPException(status_code=404, detail="Author not found")
+    db_author.name = author.name
+    session.commit()
+    session.refresh(db_author)
+    return db_author
+
 # Delete an author
-@app.delete("/authors/{author_id}", tags=["authors"])
-def delete_author(author_id: int):
-    with Session(engine) as session:
+@app.delete("/authors/{author_id}", response_model=AuthorBase, tags=["authors"])
+def delete_author(author_id: int, session: Session = Depends(get_session)):
     db_author = session.get(Author, author_id)
     if not db_author:
         raise HTTPException(status_code=404, detail="Author not found")
@@ -78,8 +74,7 @@ def delete_author(author_id: int):
 
 # Create a book
 @app.post("/books/", response_model=Book, tags=["books"])
-def create_book(book: Book):
-    with Session(engine) as session:
+def create_book(book: BookBase, session: Session = Depends(get_session)):
     session.add(book)
     session.commit()
     session.refresh(book)
@@ -87,29 +82,25 @@ def create_book(book: Book):
 
 # Read books
 @app.get("/books/", response_model=List[Book], tags=["books"])
-def read_books():
-    with Session(engine) as session:
+def read_books(session: Session = Depends(get_session)):
     books = session.exec(select(Book)).all()
     return books
 
 # Update a book
 @app.put("/books/{book_id}", response_model=Book, tags=["books"])
-def update_book(book_id: int, book: Book):
-    with Session(engine) as session:
+def update_book(book_id: int, book: Book, session: Session = Depends(get_session)):
     db_book = session.get(Book, book_id)
     if not db_book:
         raise HTTPException(status_code=404, detail="Book not found")
     db_book.title = book.title
     db_book.authors = book.authors
-        session.add(db_book)
     session.commit()
     session.refresh(db_book)
     return db_book
 
 # Delete a book
 @app.delete("/books/{book_id}", tags=["books"])
-def delete_book(book_id: int):
-    with Session(engine) as session:
+def delete_book(book_id: int, session: Session = Depends(get_session)):
     db_book = session.get(Book, book_id)
     if not db_book:
         raise HTTPException(status_code=404, detail="Book not found")

app/models.py

@@ -6,16 +6,36 @@ class AuthorBookLink(SQLModel, table=True):
     author_id: int | None = Field(default=None, foreign_key="author.id", primary_key=True)
     book_id: int | None = Field(default=None, foreign_key="book.id", primary_key=True)
 
-# Author model
-class Author(SQLModel, table=True):
-    id: int | None = Field(primary_key=True, index=True)
+# Author DTO model
+class AuthorBase(SQLModel):
     name: str
 
+    class Config: # pyright: ignore
+        json_schema_extra = {
+            "example": {
+                "name": "author_name",
+            }
+        }
+
+# Author DB model
+class Author(AuthorBase, table=True):
+    id: int | None = Field(default=None, primary_key=True, index=True)
     books: List["Book"] = Relationship(back_populates="authors", link_model=AuthorBookLink)
 
-# Book model
-class Book(SQLModel, table=True):
-    id: int | None = Field(primary_key=True, index=True)
+# Book DTO model
+class BookBase(SQLModel):
     title: str
     description: str
+
+    class Config: # pyright: ignore
+        json_schema_extra = {
+            "example": {
+                "title": "book_title",
+                "description": "book_description",
+            }
+        }
+
+# Book DB model
+class Book(BookBase, table=True):
+    id: int | None = Field(default=None, primary_key=True, index=True)
     authors: List[Author] = Relationship(back_populates="books", link_model=AuthorBookLink)