Merge pull request #14 from MetaCell/feature/CELE-39

CELE-39 Enhance REST API with filters

applications/visualizer/api/openapi.json

@@ -143,11 +143,92 @@
         ]
       }
     },
+    "/api/cells/search": {
+      "get": {
+        "operationId": "search_cells",
+        "summary": "Search Cells",
+        "parameters": [
+          {
+            "in": "query",
+            "name": "name",
+            "schema": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Name"
+            },
+            "required": false
+          },
+          {
+            "in": "query",
+            "name": "dataset_ids",
+            "schema": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Dataset Ids"
+            },
+            "required": false
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "OK",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "items": {
+                    "$ref": "#/components/schemas/Neuron"
+                  },
+                  "title": "Response",
+                  "type": "array"
+                }
+              }
+            }
+          }
+        },
+        "tags": [
+          "neurons"
+        ]
+      }
+    },
     "/api/cells": {
       "get": {
         "operationId": "get_all_cells",
         "summary": "Get All Cells",
         "parameters": [
+          {
+            "in": "query",
+            "name": "dataset_ids",
+            "schema": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Dataset Ids"
+            },
+            "required": false
+          },
           {
             "in": "query",
             "name": "page",
@@ -400,6 +481,13 @@
             "title": "Name",
             "type": "string"
           },
+          "datasetIds": {
+            "items": {
+              "type": "string"
+            },
+            "title": "Datasetids",
+            "type": "array"
+          },
           "nclass": {
             "maxLength": 30,
             "title": "Nclass",
@@ -433,6 +521,7 @@
         },
         "required": [
           "name",
+          "datasetIds",
           "nclass",
           "neurotransmitter",
           "type"

applications/visualizer/api/openapi.yaml

@@ -93,6 +93,11 @@ components:
       type: object
     Neuron:
       properties:
+        datasetIds:
+          items:
+            type: string
+          title: Datasetids
+          type: array
         embryonic:
           default: false
           title: Embryonic
@@ -122,6 +127,7 @@ components:
           type: string
       required:
       - name
+      - datasetIds
       - nclass
       - neurotransmitter
       - type
@@ -153,6 +159,16 @@ paths:
       description: Returns all the cells (neurons) from the DB
       operationId: get_all_cells
       parameters:
+      - in: query
+        name: dataset_ids
+        required: false
+        schema:
+          anyOf:
+          - items:
+              type: string
+            type: array
+          - type: 'null'
+          title: Dataset Ids
       - in: query
         name: page
         required: false
@@ -171,6 +187,41 @@ paths:
       summary: Get All Cells
       tags:
       - neurons
+  /api/cells/search:
+    get:
+      operationId: search_cells
+      parameters:
+      - in: query
+        name: name
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Name
+      - in: query
+        name: dataset_ids
+        required: false
+        schema:
+          anyOf:
+          - items:
+              type: string
+            type: array
+          - type: 'null'
+          title: Dataset Ids
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                items:
+                  $ref: '#/components/schemas/Neuron'
+                title: Response
+                type: array
+          description: OK
+      summary: Search Cells
+      tags:
+      - neurons
   /api/connections:
     get:
       description: Gets the connections of a dedicated Dataset

applications/visualizer/backend/api/api.py

@@ -1,10 +1,13 @@
+from collections import defaultdict
 from typing import Optional
 
 from django.http import HttpResponse
 from ninja import NinjaAPI, Router, Schema, Query
 from ninja.pagination import paginate, PageNumberPagination
 from django.shortcuts import aget_object_or_404
-from django.db.models import Q
+from django.db.models import Q, F, Value, CharField, Func, OuterRef
+from django.db.models.manager import BaseManager
+from django.db.models.functions import Coalesce, Concat
 
 from .schemas import Dataset, Neuron, Connection
 from .models import (
@@ -82,34 +85,94 @@ async def get_dataset(request, dataset: str):
     return await aget_object_or_404(DatasetModel, id=dataset)
 
 
+def annotate_neurons_w_dataset_ids(neurons: BaseManager[NeuronModel]) -> None:
+    """Queries the datasets ids for each neuron."""
+    neuron_names = neurons.values_list("name", flat=True).distinct()
+    pre = (
+        ConnectionModel.objects.filter(pre__in=neuron_names)
+        .values_list("pre", "dataset")
+        .distinct()
+    )
+    post = (
+        ConnectionModel.objects.filter(post__in=neuron_names)
+        .values_list("post", "dataset")
+        .distinct()
+    )
+
+    # Filter out repeated dataset ids
+    neurons_dataset_ids = defaultdict(set)
+    for neuron, dataset in pre.union(post):
+        neurons_dataset_ids[neuron].add(dataset)
+
+    for neuron in neurons:
+        neuron.dataset_ids = neurons_dataset_ids[neuron.name]  # type: ignore
+
+
+def neurons_from_datasets(
+    neurons: BaseManager[NeuronModel], dataset_ids: list[str]
+) -> BaseManager[NeuronModel]:
+    """Filters neurons belonging to specific datasets."""
+    return neurons.filter(
+        Q(
+            name__in=ConnectionModel.objects.filter(
+                dataset__id__in=dataset_ids
+            ).values_list("pre", flat=True)
+        )
+        | Q(
+            name__in=ConnectionModel.objects.filter(
+                dataset__id__in=dataset_ids
+            ).values_list("post", flat=True)
+        )
+    )
+
+
 @api.get(
     "/datasets/{dataset}/neurons",
     response={200: list[Neuron], 404: ErrorMessage},
     tags=["datasets"],
 )
-async def get_dataset_neurons(request, dataset: str):
+def get_dataset_neurons(request, dataset: str):
     """Returns all the neurons of a dedicated dataset"""
-    return await to_list(
-        NeuronModel.objects.filter(
-            Q(
-                name__in=ConnectionModel.objects.filter(
-                    dataset__id=dataset
-                ).values_list("pre", flat=True)
-            )
-            | Q(
-                name__in=ConnectionModel.objects.filter(
-                    dataset__id=dataset
-                ).values_list("post", flat=True)
-            )
-        )
-    )
+    neurons = neurons_from_datasets(NeuronModel.objects, [dataset])
+    annotate_neurons_w_dataset_ids(neurons)
+    return neurons
+
+
+@api.get("/cells/search", response=list[Neuron], tags=["neurons"])
+def search_cells(
+    request,
+    name: Optional[str] = Query(None),
+    dataset_ids: Optional[list[str]] = Query(None),
+):
+    neurons = NeuronModel.objects
+
+    if name:
+        neurons = neurons.filter(name__istartswith=name)
+
+    if dataset_ids:
+        neurons = neurons_from_datasets(neurons, dataset_ids)
+    else:
+        neurons = neurons.all()
+
+    annotate_neurons_w_dataset_ids(neurons)
+
+    return neurons
 
 
 @api.get("/cells", response=list[Neuron], tags=["neurons"])
-@paginate(PageNumberPagination, page_size=50)
-def get_all_cells(request):
+@paginate(PageNumberPagination, page_size=50)  # BUG: this is not being applied
+def get_all_cells(request, dataset_ids: Optional[list[str]] = Query(None)):
     """Returns all the cells (neurons) from the DB"""
-    return NeuronModel.objects.all()
+    neurons = NeuronModel.objects
+
+    if dataset_ids:
+        neurons = neurons_from_datasets(neurons, dataset_ids)
+    else:
+        neurons = neurons.all()
+
+    annotate_neurons_w_dataset_ids(neurons)
+
+    return neurons
 
 
 # # @api.post("/connections", response=list[Connection], tags=["connectivity"])

applications/visualizer/backend/api/schemas.py

@@ -37,6 +37,7 @@ class Meta:
 
 class Neuron(ModelSchema, BilingualSchema):
     name: str
+    dataset_ids: list[str]
 
     class Meta:
         model = NeuronModel