From 03e546e9f0793d351fe75b9aba0290b6249d4bf7 Mon Sep 17 00:00:00 2001
From: Ife Lawal <il693@nyu.edu>
Date: Mon, 26 Apr 2021 14:00:04 -0400
Subject: [PATCH] Pydantic documentation Fixes #1829 (#1871)

- Improved the `sqlalchemy_to_pydantic` function to accept additional schema fields on top of the SQLAlchemy model fields
- Added the solves and solved_by_me fields to the Swagger documentation (Closes #1829)
---
 CTFd/api/v1/challenges.py      | 4 +++-
 CTFd/api/v1/helpers/schemas.py | 8 ++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/CTFd/api/v1/challenges.py b/CTFd/api/v1/challenges.py
index 979313d8..1dbe7f58 100644
--- a/CTFd/api/v1/challenges.py
+++ b/CTFd/api/v1/challenges.py
@@ -62,7 +62,9 @@ challenges_namespace = Namespace(
     "challenges", description="Endpoint to retrieve Challenges"
 )
 
-ChallengeModel = sqlalchemy_to_pydantic(Challenges)
+ChallengeModel = sqlalchemy_to_pydantic(
+    Challenges, include={"solves": int, "solved_by_me": bool}
+)
 TransientChallengeModel = sqlalchemy_to_pydantic(Challenges, exclude=["id"])
 
 
diff --git a/CTFd/api/v1/helpers/schemas.py b/CTFd/api/v1/helpers/schemas.py
index 51345c19..cd3ed4f6 100644
--- a/CTFd/api/v1/helpers/schemas.py
+++ b/CTFd/api/v1/helpers/schemas.py
@@ -1,4 +1,4 @@
-from typing import Container, Type
+from typing import Container, Dict, Type
 
 from pydantic import BaseModel, create_model
 from sqlalchemy.inspection import inspect
@@ -6,7 +6,7 @@ from sqlalchemy.orm.properties import ColumnProperty
 
 
 def sqlalchemy_to_pydantic(
-    db_model: Type, *, exclude: Container[str] = None
+    db_model: Type, *, include: Dict[str, type] = None, exclude: Container[str] = None
 ) -> Type[BaseModel]:
     """
     Mostly copied from https://github.com/tiangolo/pydantic-sqlalchemy
@@ -27,6 +27,10 @@ def sqlalchemy_to_pydantic(
                 if column.default is None and not column.nullable:
                     default = ...
                 fields[name] = (python_type, default)
+    if bool(include):
+        for name, python_type in include.items():
+            default = None
+            fields[name] = (python_type, default)
     pydantic_model = create_model(
         db_model.__name__, **fields  # type: ignore
     )