Update code in endpoints/pen.post.py

endpoints/pen.post.py

@@ -12,70 +12,70 @@ app = FastAPI()
 
 # SQLAlchemy setup
 Base = declarative_base()
-engine = create_engine('sqlite:///database.db', echo=True)
-Session = sessionmaker(bind=engine)
+engine = create_engine('sqlite:///database.db', connect_args={'check_same_thread': False})
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+# Pydantic model (schema)
+class PenSchema(BaseModel):
+    name: str
+    color: str
+    price: float
+
+    class Config:
+        orm_mode = True
 
 # SQLAlchemy model
 class Pen(Base):
     __tablename__ = 'pens'
 
-    id = Column(Integer, primary_key=True)
+    id = Column(Integer, primary_key=True, index=True)
     name = Column(String)
     color = Column(String)
-    brand = Column(String)
+    price = Column(Integer)
 
-    def __repr__(self):
-        return f"Pen(id={self.id}, name='{self.name}', color='{self.color}', brand='{self.brand}')"
+Base.metadata.create_all(bind=engine)
 
-Base.metadata.create_all(engine)
-
-# Pydantic model
-class PenSchema(BaseModel):
-    name: str
-    color: str
-    brand: str
-
-    class Config:
-        orm_mode = True
-
-# Endpoint to save a pen to the database
-@app.post('/pen')
-def save_pen(pen: PenSchema):
-    session = Session()
+# Dependency
+def get_db():
+    db = SessionLocal()
     try:
-        pen_obj = Pen(name=pen.name, color=pen.color, brand=pen.brand)
-        session.add(pen_obj)
-        session.commit()
-        return {"message": "Pen saved successfully"}
-    except Exception as e:
-        session.rollback()
-        raise HTTPException(status_code=500, detail=str(e))
+        yield db
     finally:
-        session.close()
+        db.close()
+
+@app.post('/pen', status_code=201)
+def create_pen(pen: PenSchema, db=get_db()):
+    new_pen = Pen(name=pen.name, color=pen.color, price=pen.price)
+    db.add(new_pen)
+    db.commit()
+    db.refresh(new_pen)
+    return new_pen
 ```
 
 Here's a breakdown of the code:
 
-1. We import the necessary modules from FastAPI, Pydantic, and SQLAlchemy.
-2. We set up SQLAlchemy by creating a declarative base, engine, and session maker.
-3. We define the `Pen` model using SQLAlchemy's `declarative_base` and `Column` classes.
-4. We create a Pydantic model `PenSchema` that will be used for data validation and serialization/deserialization.
-5. We define the `/pen` endpoint using the `@app.post` decorator from FastAPI.
-6. Inside the `save_pen` function, we create a new SQLAlchemy session and try to save the pen data to the database.
-7. If the operation is successful, we commit the changes and return a success message.
-8. If an exception occurs, we roll back the changes and raise an `HTTPException` with a 500 status code and the error message.
-9. Finally, we close the session in the `finally` block.
+1. We import the necessary modules and create a FastAPI instance.
+2. We set up SQLAlchemy with a SQLite database and define a base class for our models.
+3. We define a Pydantic model (`PenSchema`) to validate and parse the incoming data.
+4. We define a SQLAlchemy model (`Pen`) that maps to the `pens` table in the database.
+5. We create the `pens` table in the database using `Base.metadata.create_all(bind=engine)`.
+6. We define a dependency function `get_db()` that creates a new database session for each request and closes it after the request is processed.
+7. We define the `/pen` endpoint with the `POST` method and the `status_code=201` (Created) response.
+8. Inside the endpoint function, we create a new `Pen` instance using the data from the `PenSchema` model.
+9. We add the new `Pen` instance to the database session, commit the changes, and refresh the instance to get the generated ID.
+10. Finally, we return the new `Pen` instance.
 
-To use this endpoint, you can send a POST request to `/pen` with a JSON payload containing the pen data, like this:
+To use this endpoint, you can send a POST request to `/pen` with a JSON payload containing the pen details, like this:
+
+```
+POST /pen
+Content-Type: application/json
 
-```json
 {
-    "name": "Pilot G2",
-    "color": "Blue",
-    "brand": "Pilot"
+    "name": "Pilot FriXion",
+    "color": "Black",
+    "price": 2.99
 }
 ```
 
-The endpoint will validate the data using the `PenSchema` model, create a new `Pen` instance, and save it to the database.
-
-Note that you'll need to have SQLAlchemy and a compatible database driver installed for this code to work. In this example, we're using SQLite, but you can easily switch to a different database by modifying the engine URL.
+The endpoint will create a new pen in the database and return the created pen object with the generated ID.