Refactor and fix constraint for fhir

app/aidbox/operations.py

@@ -8,14 +8,17 @@
     assemble,
     constraint_check,
     extract,
-    extract_questionnaire_instance,
     get_questionnaire_context,
     populate,
     resolve_expression,
 )
-from ..sdc.utils import parameter_to_env
+from ..sdc.utils import parameter_to_env, validate_context
 from ..utils import get_extract_services
 from .utils import AidboxSdcRequest, aidbox_operation, get_user_sdk_client, prepare_args
+from ..sdc.exception import (
+    ConstraintCheckOperationOutcome,
+    MissingParamOperationOutcome,
+)
 
 
 @aidbox_operation(["GET"], ["Questionnaire", {"name": "id"}, "$assemble"])
@@ -28,7 +31,9 @@ async def assemble_op(request: AidboxSdcRequest):
     )
 
     assembled_questionnaire_lazy = await assemble(request.fhir_client, questionnaire)
-    assembled_questionnaire = json.loads(json.dumps(assembled_questionnaire_lazy, default=list))
+    assembled_questionnaire = json.loads(
+        json.dumps(assembled_questionnaire_lazy, default=list)
+    )
     if request.is_fhir:
         assembled_questionnaire = await from_first_class_extension(
             assembled_questionnaire, request.aidbox_client
@@ -41,30 +46,30 @@ async def assemble_op(request: AidboxSdcRequest):
 async def constraint_check_operation(request: AidboxSdcRequest):
     env = parameter_to_env(request.resource)
 
-    questionnaire = (
+    fce_questionnaire = (
         await to_first_class_extension(env["Questionnaire"], request.aidbox_client)
         if request.is_fhir
         else env["Questionnaire"]
     )
-    as_root = questionnaire.get("runOnBehalfOfRoot")
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
     client = client if as_root else get_user_sdk_client(request.request, request.client)
 
-    return web.json_response(await constraint_check(client, env))
+    return web.json_response(await constraint_check(client, fce_questionnaire, env))
 
 
 @aidbox_operation(["POST"], ["Questionnaire", "$context"])
 @prepare_args
 async def get_questionnaire_context_operation(request: AidboxSdcRequest):
     env = parameter_to_env(request.resource)
 
-    questionnaire = (
+    fce_questionnaire = (
         await to_first_class_extension(env["Questionnaire"], request.aidbox_client)
         if request.is_fhir
         else env["Questionnaire"]
     )
-    as_root = questionnaire.get("runOnBehalfOfRoot")
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
     client = client if as_root else get_user_sdk_client(request.request, request.client)
-    result = await get_questionnaire_context(client, env)
+    result = await get_questionnaire_context(client, fce_questionnaire, env)
 
     return web.json_response(result)
 
@@ -74,39 +79,50 @@ async def get_questionnaire_context_operation(request: AidboxSdcRequest):
 async def extract_questionnaire_operation(request: AidboxSdcRequest):
     resource = request.resource
 
-    as_root = False
     if resource["resourceType"] == "QuestionnaireResponse":
         env = {}
-        questionnaire_response = resource
-        questionnaire = (
+        env_questionnaire_response = resource
+        fce_questionnaire = (
             await request.aidbox_client.resources("Questionnaire")
             .search(_id=resource["questionnaire"])
             .get()
         )
-        as_root = questionnaire.get("runOnBehalfOfRoot")
+        env_questionnaire = (
+            await from_first_class_extension(fce_questionnaire, request.aidbox_client)
+            if request.is_fhir
+            else fce_questionnaire
+        )
     elif resource["resourceType"] == "Parameters":
         env = parameter_to_env(request.resource)
-        questionnaire = (
+        if "Questionnaire" not in env:
+            raise MissingParamOperationOutcome("`Questionnaire` parameter is required")
+        if "QuestionnaireResponse" not in env:
+            raise MissingParamOperationOutcome(
+                "`QuestionnaireResponse` parameter is required"
+            )
+
+        fce_questionnaire = (
             await to_first_class_extension(env["Questionnaire"], request.aidbox_client)
             if request.is_fhir
             else env["Questionnaire"]
         )
-        questionnaire_response = env.get("QuestionnaireResponse")
-        as_root = questionnaire.get("runOnBehalfOfRoot")
+        env_questionnaire = env["Questionnaire"]
+        env_questionnaire_response = env["QuestionnaireResponse"]
 
     mappings = [
         await request.aidbox_client.resources("Mapping").search(_id=m["id"]).get()
-        for m in questionnaire.get("mapping", [])
+        for m in fce_questionnaire.get("mapping", [])
     ]
 
     context = {
-        "Questionnaire": questionnaire,
-        "QuestionnaireResponse": questionnaire_response,
+        "Questionnaire": env_questionnaire,
+        "QuestionnaireResponse": env_questionnaire_response,
         **env,
     }
 
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
     client = client if as_root else get_user_sdk_client(request.request, request.client)
-    await constraint_check(client, context)
+    await constraint_check(client, fce_questionnaire, context)
     extraction_result = await extract(
         client, mappings, context, get_extract_services(request.request["app"])
     )
@@ -117,75 +133,127 @@ async def extract_questionnaire_operation(request: AidboxSdcRequest):
 @prepare_args
 async def extract_questionnaire_instance_operation(request: AidboxSdcRequest):
     resource = request.resource
-    questionnaire = (
+    fce_questionnaire = (
         await request.aidbox_client.resources("Questionnaire")
         .search(_id=request.route_params["id"])
         .get()
     )
-    as_root = questionnaire.get("runOnBehalfOfRoot")
+    env_questionnaire = (
+        await from_first_class_extension(fce_questionnaire, request.aidbox_client)
+        if request.is_fhir
+        else fce_questionnaire
+    )
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
     extract_client = (
-        request.client if as_root else get_user_sdk_client(request.request, request.client)
+        request.client
+        if as_root
+        else get_user_sdk_client(request.request, request.client)
     )
     return web.json_response(
         await extract_questionnaire_instance(
             request.aidbox_client,
             extract_client,
-            questionnaire,
+            fce_questionnaire,
+            env_questionnaire,
             resource,
             get_extract_services(request.request["app"]),
         )
     )
 
 
+async def extract_questionnaire_instance(
+    aidbox_client,
+    extract_client,
+    fce_questionnaire,
+    env_questionnaire,
+    resource,
+    extract_services,
+):
+    if resource["resourceType"] == "QuestionnaireResponse":
+        env = {}
+        env_questionnaire_response = extract_client.resource(
+            "QuestionnaireResponse", **resource
+        )
+    elif resource["resourceType"] == "Parameters":
+        env = parameter_to_env(resource)
+        if "QuestionnaireResponse" not in env:
+            raise MissingParamOperationOutcome(
+                "`QuestionnaireResponse` parameter is required"
+            )
+
+        env_questionnaire_response = env["QuestionnaireResponse"]
+    else:
+        raise MissingParamOperationOutcome(
+            "Either `QuestionnaireResponse` resource or Parameters containing  QuestionnaireResponse are required",
+        )
+
+    if "launchContext" in fce_questionnaire:
+        validate_context(fce_questionnaire["launchContext"], env)
+    context = {
+        "QuestionnaireResponse": env_questionnaire_response,
+        "Questionnaire": env_questionnaire,
+        **env,
+    }
+    mappings = [
+        await aidbox_client.resources("Mapping").search(_id=m["id"]).get()
+        for m in fce_questionnaire.get("mapping", [])
+    ]
+    await constraint_check(extract_client, fce_questionnaire, context)
+
+    return await extract(extract_client, mappings, context, extract_services)
+
+
 @aidbox_operation(["POST"], ["Questionnaire", "$populate"])
 @prepare_args
 async def populate_questionnaire(request: AidboxSdcRequest):
     env = parameter_to_env(request.resource)
 
     if "Questionnaire" not in env:
-        # TODO: return OperationOutcome
-        return web.json_response(
-            {
-                "error": "bad_request",
-                "error_description": "`Questionnaire` parameter is required",
-            },
-            status=422,
-        )
+        raise MissingParamOperationOutcome("`Questionnaire` parameter is required")
 
-    questionnaire = (
+    fce_questionnaire = (
         await to_first_class_extension(env["Questionnaire"], request.aidbox_client)
         if request.is_fhir
         else env["Questionnaire"]
     )
-    as_root = questionnaire.get("runOnBehalfOfRoot")
-    client = request.client if as_root else get_user_sdk_client(request.request, request.client)
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
+    client = (
+        request.client
+        if as_root
+        else get_user_sdk_client(request.request, request.client)
+    )
 
-    populated_resource = await populate(client, questionnaire, env)
+    fce_populated_qr = await populate(client, fce_questionnaire, env)
     if request.is_fhir:
-        populated_resource = await from_first_class_extension(
-            populated_resource, request.aidbox_client
+        fce_populated_qr = await from_first_class_extension(
+            fce_populated_qr, request.aidbox_client
         )
-    return web.json_response(populated_resource)
+    return web.json_response(fce_populated_qr)
 
 
 @aidbox_operation(["POST"], ["Questionnaire", {"name": "id"}, "$populate"])
 @prepare_args
 async def populate_questionnaire_instance(request: AidboxSdcRequest):
-    questionnaire = (
+    fce_questionnaire = (
         await request.aidbox_client.resources("Questionnaire")
         .search(_id=request.route_params["id"])
         .get()
     )
     env = parameter_to_env(request.resource)
-    as_root = questionnaire.get("runOnBehalfOfRoot")
+    env["Questionnaire"] = (
+        await from_first_class_extension(fce_questionnaire, request.aidbox_client)
+        if request.is_fhir
+        else fce_questionnaire
+    )
+    as_root = fce_questionnaire.get("runOnBehalfOfRoot")
     client = client if as_root else get_user_sdk_client(request.request, request.client)
 
-    populated_resource = await populate(client, questionnaire, env)
+    fce_populated_qr = await populate(client, fce_questionnaire, env)
     if request.is_fhir:
-        populated_resource = await from_first_class_extension(
-            populated_resource, request.aidbox_client
+        fce_populated_qr = await from_first_class_extension(
+            fce_populated_qr, request.aidbox_client
         )
-    return web.json_response(populated_resource)
+    return web.json_response(fce_populated_qr)
 
 
 @aidbox_operation(["POST"], ["Questionnaire", "$resolve-expression"], public=True)

app/fhir_server/operations.py

@@ -37,11 +37,12 @@ async def assemble_handler(request: web.BaseRequest):
 @routes.post("/QuestionnaireResponse/$constraint-check")
 async def constraint_check_handler(request: web.BaseRequest):
     env = parameter_to_env(await request.json())
+    # TODO: I believe there's a bug, it should be in FHIR format
     env["Questionnaire"] = to_first_class_extension(env["Questionnaire"])
     env["QuestionnaireResponse"] = to_first_class_extension(env["QuestionnaireResponse"])
     client = request.app["client"]
 
-    return web.json_response(await constraint_check(client, env))
+    return web.json_response(await constraint_check(client, env["Questionnaire"], env))
 
 
 @routes.post("/Questionnaire/$context")
@@ -50,7 +51,13 @@ async def get_questionnaire_context_handler(request: web.BaseRequest):
     env = parameter_to_env(await request.json())
     client = request.app["client"]
 
-    return web.json_response(await get_questionnaire_context(client, env))
+    return web.json_response(
+        await get_questionnaire_context(
+            client, 
+            to_first_class_extension(env["Questionnaire"]),
+            env
+        )
+    )
 
 
 @routes.post("/Questionnaire/$extract")
@@ -93,14 +100,14 @@ async def extract_questionnaire_handler(request: web.BaseRequest):
             ]
         )
         jute_templates.append(json.loads(template_string))
-
+    # TODO: I believe there's a bug, it should be in FHIR
     context = {
         "Questionnaire": to_first_class_extension(questionnaire),
         "QuestionnaireResponse": to_first_class_extension(questionnaire_response),
         **env,
     }
 
-    await constraint_check(client, context)
+    await constraint_check(client, to_first_class_extension(questionnaire), context)
     extraction_result = await extract(
         client, jute_templates, context, get_extract_services(request.app)
     )
@@ -142,11 +149,12 @@ async def extract_questionnaire_instance_operation(request: web.BaseRequest):
 
     if resource["resourceType"] == "QuestionnaireResponse":
         questionnaire_response = client.resource("QuestionnaireResponse", **resource)
+        # TODO: I believe there's a bug, it should be in FHIR
         context = {
             "Questionnaire": to_first_class_extension(questionnaire),
             "QuestionnaireResponse": questionnaire_response,
         }
-        await constraint_check(client, context)
+        await constraint_check(client, questionnaire, context)
         return web.json_response(
             await extract(client, jute_templates, context, get_extract_services(request.app))
         )
@@ -176,7 +184,7 @@ async def extract_questionnaire_instance_operation(request: web.BaseRequest):
             "Questionnaire": questionnaire,
             **env,
         }
-        await constraint_check(client, context)
+        await constraint_check(client, questionnaire, context)
         return web.json_response(
             await extract(client, jute_templates, context, get_extract_services(request.app))
         )