User surveys

marble_api/app.py

@@ -6,6 +6,7 @@
 
 _metadata = metadata.metadata("marble_api").json
 
+
 app = FastAPI(
     title=_metadata["name"],
     version=_metadata["version"],

marble_api/database/__init__.py

@@ -1,4 +1,7 @@
+from collections.abc import Callable
+
 from pymongo import AsyncMongoClient
+from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.database import AsyncDatabase
 
 from marble_api._config import config
@@ -18,3 +21,18 @@ def db(self) -> AsyncDatabase:
 
 
 client = Client(str(config.mongodb_uri), tz_aware=True)
+
+
+def collection(name: str) -> Callable:
+    """
+    Return a function that returns the named mongodb collection.
+
+    Using this ensures that a new collection object is created when needed instead
+    of creating a collection once when the module is loaded (which may drop connection
+    to the database if the mongodb instance goes offline temporarily.)
+    """
+
+    def _(name: str = name) -> AsyncCollection:
+        return client.db[name]
+
+    return _

marble_api/utils/models.py

@@ -1,48 +1,37 @@
-from copy import deepcopy
-from typing import Any
+import re
+from contextlib import suppress
+from typing import Annotated, Any
 
 import bson
-from pydantic import BaseModel, create_model
-from pydantic.fields import FieldInfo
+from bson.regex import Regex as BsonRegex
+from pydantic import (
+    AwareDatetime,
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    FieldSerializationInfo,
+    PydanticUndefinedAnnotation,
+    computed_field,
+    field_serializer,
+)
+from pydantic.json_schema import SkipJsonSchema
 
 
-def partial_model(model: type[BaseModel]) -> type[BaseModel]:
+def _convert_bson_regex(value: Any) -> str | re.Pattern:  # noqa: ANN401
     """
-    Make all fields in a BaseModel class optional.
+    Convert bson.regex.Regex to a python re.Pattern.
 
-    This makes each field's default None but does not update the annotation or
-    validations so explicitly setting the value to None may still raise a
-    validation error. Also, if a field has validate_default=True this will
-    make validate_default=False for the partial model to ensure that the new
-    (None) default value is not validated.
-
-    >>> class C(BaseModel):
-          a: int
-    >>> C(a=2).a
-    2
-    >>> C()  # validation error since a must be an integer
-    >>> @partial_model
-    ... class B(C): ...
-    >>> B().a  # is None
-    >>> B(a=5).a
-    5
-    >>> B(a=None)  # validation error since a must be an integer
-
-    Adapted from https://stackoverflow.com/a/76560886/5992438
+    MongoDB will return regular expressions as bson.regex.Regex which
+    pydantic doensn't know how to convert to a re.Pattern without help.
     """
+    if isinstance(value, BsonRegex):
+        return value.try_compile()
+    return value
 
-    def make_field_optional(field: FieldInfo) -> tuple[Any, FieldInfo]:
-        new_field = deepcopy(field)
-        new_field.validate_default = False
-        new_field.default = None
-        return new_field.annotation, new_field
 
-    return create_model(
-        model.__name__,
-        __base__=model,
-        __module__=model.__module__,
-        **{name: make_field_optional(info) for name, info in model.model_fields.items()},
-    )
+type PyObjectId = Annotated[str, BeforeValidator(str)]
+type Pattern = Annotated[re.Pattern, BeforeValidator(_convert_bson_regex)]
 
 
 def object_id(id_: str, error: Exception | None) -> bson.ObjectId:
@@ -56,3 +45,101 @@ def object_id(id_: str, error: Exception | None) -> bson.ObjectId:
     except bson.errors.InvalidId as err:
         if error is not None:
             raise error from err
+
+
+class PartialModel(BaseModel):
+    """Superclass for partial models (all fields are optional)."""
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs) -> None:
+        """
+        Make all models that this model a "partial model".
+
+        Partial models make all fields optional.
+
+        This sets all fields to have a default value of None but don't update the annotation
+        so that fields can't explicitly be set to None.
+        """
+        super().__pydantic_init_subclass__(**kwargs)
+        for field in cls.model_fields.values():
+            field.default = None
+
+        with suppress(PydanticUndefinedAnnotation):
+            cls.model_rebuild(force=True)
+
+
+class MarbleBaseModel(BaseModel):
+    """
+    Base model for all database models used in this repo.
+
+    Includes:
+        - an id stored as a PyObjectId in the database
+        - created field calculated from the id value
+        - updated field that should be updated by the route
+    """
+
+    id: SkipJsonSchema[PyObjectId | None] = Field(default=None, validation_alias="_id", exclude=True)
+    updated: SkipJsonSchema[AwareDatetime | None] = None  # updated should set by the route
+    model_config = ConfigDict(validate_by_name=True, arbitrary_types_allowed=True, validate_assignment=True)
+
+
+class MarbleBaseModelUpdate(PartialModel, MarbleBaseModel):
+    """
+    Base model for update models in this repo.
+
+    Ensures that subclasses of this model are partial models (all fields are optional)
+    """
+
+
+class MarbleBaseModelPublic(MarbleBaseModel):
+    """
+    Base model for public models (visible to the end user) in this repo.
+
+    Makes the id and updated fields visible and creates a computed field for model creation time.
+
+    Sets extra=allow so additional data can be added to models returned by route methods.
+    Does not set validate_assignment=True to allow flexibility when routes add/update data.
+    """
+
+    id: Annotated[str, BeforeValidator(str)] = Field(..., validation_alias="_id")
+    updated: AwareDatetime
+    model_config = ConfigDict(validate_by_name=True, arbitrary_types_allowed=True, extra="allow")
+
+    @computed_field
+    def created(self) -> AwareDatetime:
+        """Set the created time based on the object id."""
+        return object_id(self.id, None).generation_time
+
+
+class MarbleUserModel(MarbleBaseModel):
+    """
+    Base model for models that are associated with a user.
+
+    Requires that the user attribute is set when serialized.
+    """
+
+    user: SkipJsonSchema[str | None] = None  # user is set by the route after the model is first validated
+
+    @field_serializer("user")
+    def require_user_set(self, value: str, info: FieldSerializationInfo) -> SkipJsonSchema[str]:
+        """Require that the user name be set when the model is serialized."""
+        if not value:
+            raise ValueError(f"{info.field_name} must be set and non-empty")
+        return value
+
+
+class MarbleUserModelUpdate(MarbleBaseModelUpdate, MarbleUserModel):
+    """Base model for updating user associated models."""
+
+
+class MarbleUserModelPublic(MarbleBaseModelPublic, MarbleUserModel):
+    """Base model for public models associated with a user."""
+
+    user: str = Field(..., min_length=1)  # user is required to be set in the database
+
+    # NOTE: this must be redeclared so that the return annotation can be modified to
+    #       allow the user field to show up in the JSON schema.
+    @field_serializer("user")
+    def require_user_set(self, value: str, info: FieldSerializationInfo) -> str:
+        """Require that the user name be set when the model is serialized."""
+        return super().require_user_set(value, info)

marble_api/utils/routes.py

@@ -1,13 +1,20 @@
-from typing import Any
+import datetime
+from typing import Annotated, Any, TypedDict
 
 import pymongo
 from bson import ObjectId
-from fastapi import Request
+from fastapi import HTTPException, Query, Request, Response, status
+from pydantic import BaseModel
 from pymongo.asynchronous.collection import AsyncCollection
 from stac_pydantic.links import Links
 
+from marble_api.database import client
+from marble_api.utils.models import MarbleBaseModel
 
-async def paginated_query(
+type MongoRecord = dict[str, Any]
+
+
+async def _paginated_query(
     collection: AsyncCollection,
     limit: int,
     request: Request,
@@ -16,7 +23,7 @@ async def paginated_query(
     before: ObjectId | None = None,
     ascending: bool = True,
     **selector,
-) -> tuple[list[dict[str, Any]], Links]:
+) -> tuple[list[MongoRecord], Links]:
     """Return the result of a paginated query."""
     reverse_it = False
     sort_order = pymongo.ASCENDING if ascending else pymongo.DESCENDING
@@ -82,3 +89,133 @@ async def paginated_query(
             }
         )
     return data, links
+
+
+def _build_selector(
+    id_: ObjectId | None, user: str | None, additional_selector: dict[str, Any] | None
+) -> dict[str, Any]:
+    selector = additional_selector or {}
+    if id_ is not None:
+        selector["_id"] = id_
+    if user is not None:
+        selector["user"] = user
+    return selector
+
+
+async def post_record(collection: AsyncCollection, user: str | None, model: MarbleBaseModel) -> MongoRecord:
+    """Create a record in the database collection based on the given model and return the created model."""
+    if user is not None:
+        model.user = user
+    model.updated = datetime.datetime.now(tz=datetime.timezone.utc)
+    new_model = model.model_dump(by_alias=True)
+    result = await collection.insert_one(new_model)
+    new_model["id"] = str(result.inserted_id)
+    return new_model
+
+
+async def patch_record(
+    collection: AsyncCollection,
+    id_: ObjectId | None,
+    user: str | None,
+    data: Any,  # not MarbleBaseModelUpdate because patch can operate on individual fields  # noqa: ANN401
+    validation_class: MarbleBaseModel,
+    allow_update_user: bool = False,
+    additional_selector: dict[str, Any] | None = None,
+    set_prefix: str | None = None,
+) -> MongoRecord:
+    """Update fields of a record in the database collection and return the updated record."""
+    if isinstance(data, BaseModel):
+        updated_fields = data.model_dump(exclude_unset=True, by_alias=True)
+    else:
+        updated_fields = data  # allows updating partial records that are not defined by a pydantic model
+    if set_prefix is not None:
+        updated_fields = {f"{set_prefix}{field}": value for field, value in updated_fields.items()}
+    selector = _build_selector(id_, user, additional_selector)
+    if not allow_update_user and "user" in updated_fields:
+        # do not allow updating user field
+        raise HTTPException(status_code=403, detail="Forbidden")
+    # updated timestamps are handled automatically
+    if updated_fields:
+        updated_fields["updated"] = datetime.datetime.now(tz=datetime.timezone.utc)
+        async with client.start_session() as session:
+            async with await session.start_transaction():
+                result = await collection.find_one_and_update(
+                    selector, {"$set": updated_fields}, return_document=pymongo.ReturnDocument.AFTER
+                )
+                if result is not None:
+                    validation_class(**result)
+                    return result
+    else:
+        if (result := await collection.find_one(selector)) is not None:
+            return result
+
+    raise HTTPException(status_code=404, detail="Not found")
+
+
+async def get_record(
+    collection: AsyncCollection,
+    id_: ObjectId | None,
+    user: str | None,
+    additional_selector: dict[str, Any] | None = None,
+) -> MongoRecord:
+    """Get a record from the database collection with the given id."""
+    selector = _build_selector(id_, user, additional_selector)
+    if (result := await collection.find_one(selector)) is not None:
+        return result
+
+    raise HTTPException(status_code=404, detail="Not found")
+
+
+class RecordsResponse(TypedDict):
+    """Describes the return value for the get_records function."""
+
+    records: list[MongoRecord]
+    links: Links
+
+
+async def get_records(
+    collection: AsyncCollection,
+    request: Request,
+    user: str | None = None,
+    after: ObjectId | None = None,
+    before: ObjectId | None = None,
+    limit: Annotated[int, Query(le=100, gt=0)] = 10,
+    sort_by: str = "id",
+    ascending: bool = True,
+    additional_selector: dict[str, Any] | None = None,
+    records_key: str = "records",
+) -> RecordsResponse:
+    """Get records from the database collection, paginated according to the after, before, sort_by, and ascending params."""
+    # note: created is not a field, it is based on the timesamp used to create the value of _id
+    if sort_by in ("id", "created"):
+        sort_by = "_id"
+
+    selector = _build_selector(None, user, additional_selector)
+
+    records, links = await _paginated_query(
+        collection=collection,
+        limit=limit,
+        request=request,
+        sort_by=sort_by,
+        after=after,
+        before=before,
+        ascending=ascending,
+        **selector,
+    )
+    return {records_key: records, "links": links}
+
+
+async def delete_record(
+    collection: AsyncCollection,
+    id_: ObjectId | None,
+    user: str | None,
+    additional_selector: dict[str, Any] | None = None,
+) -> Response:
+    """Delete a data request from the database collection with the given id."""
+    selector = _build_selector(id_, user, additional_selector)
+
+    result = await collection.delete_one(selector)
+    if result.deleted_count == 1:
+        return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+    raise HTTPException(status_code=404, detail="Not found")

marble_api/versions/v1/__init__.py

@@ -3,6 +3,7 @@
 from marble_api.utils.auth import authenticate_magpie_admin, authenticate_magpie_user
 from marble_api.versions.v1.data_request.routes import admin_router as data_request_admin_router
 from marble_api.versions.v1.data_request.routes import user_router as data_request_user_router
+from marble_api.versions.v1.survey.routes import survey_admin_router, survey_response_admin_router, survey_user_router
 
 router = APIRouter(prefix="/v1")
 
@@ -12,5 +13,9 @@
 user_router.include_router(data_request_user_router)
 admin_router.include_router(data_request_admin_router)
 
+user_router.include_router(survey_user_router)
+admin_router.include_router(survey_admin_router)
+admin_router.include_router(survey_response_admin_router)
+
 router.include_router(user_router)
 router.include_router(admin_router)