From f8b9d9802e8800d8e17cc8c65041bcaa71f0b114 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Fri, 8 Jul 2022 17:32:56 +0300
Subject: [PATCH] Add support for organization.get_entities_count

---
 apiserver/apimodels/organization.py         |  9 ++++
 apiserver/schema/services/organization.conf | 53 ++++++++++++++++++++-
 apiserver/services/organization.py          | 35 +++++++++++++-
 3 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/apiserver/apimodels/organization.py b/apiserver/apimodels/organization.py
index a7b3f6f..424b567 100644
--- a/apiserver/apimodels/organization.py
+++ b/apiserver/apimodels/organization.py
@@ -1,5 +1,7 @@
 from jsonmodels import fields, models
 
+from apiserver.apimodels import DictField
+
 
 class Filter(models.Base):
     tags = fields.ListField([str])
@@ -9,3 +11,10 @@ class Filter(models.Base):
 class TagsRequest(models.Base):
     include_system = fields.BoolField(default=False)
     filter = fields.EmbeddedField(Filter)
+
+
+class EntitiesCountRequest(models.Base):
+    projects = DictField()
+    tasks = DictField()
+    models = DictField()
+    pipelines = DictField()
diff --git a/apiserver/schema/services/organization.conf b/apiserver/schema/services/organization.conf
index 1625bbe..48bbc06 100644
--- a/apiserver/schema/services/organization.conf
+++ b/apiserver/schema/services/organization.conf
@@ -102,4 +102,55 @@ get_user_companies {
             }
         }
     }
-}
\ No newline at end of file
+}
+get_entities_count {
+    "999.0": {
+        description: "Get counts for the company entities according to the passed search criteria"
+        request {
+            type: object
+            properties {
+                projects {
+                    type: object
+                    additionalProperties: true
+                    description: Search criteria for projects
+                }
+                tasks {
+                    type: object
+                    additionalProperties: true
+                    description: Search criteria for experiments
+                }
+                models {
+                    type: object
+                    additionalProperties: true
+                    description: Search criteria for models
+                }
+                pipelines {
+                    type: object
+                    additionalProperties: true
+                    description: Search criteria for pipelines
+                }
+            }
+        }
+        response {
+            type: object
+            properties {
+                projects {
+                    type: integer
+                    description: The number of projects matching the criteria
+                }
+                tasks {
+                    type: integer
+                    description: The number of experiments matching the criteria
+                }
+                models {
+                    type: integer
+                    description: The number of models matching the criteria
+                }
+                pipelines {
+                    type: integer
+                    description: The number of pipelines matching the criteria
+                }
+            }
+        }
+    }
+}
diff --git a/apiserver/services/organization.py b/apiserver/services/organization.py
index 7bc82a9..33da38e 100644
--- a/apiserver/services/organization.py
+++ b/apiserver/services/organization.py
@@ -1,9 +1,15 @@
 from collections import defaultdict
 from operator import itemgetter
+from typing import Mapping, Type
 
-from apiserver.apimodels.organization import TagsRequest
+from mongoengine import Q
+
+from apiserver.apimodels.organization import TagsRequest, EntitiesCountRequest
 from apiserver.bll.organization import OrgBLL, Tags
-from apiserver.database.model import User
+from apiserver.database.model import User, AttributedDocument, EntityVisibility
+from apiserver.database.model.model import Model
+from apiserver.database.model.project import Project
+from apiserver.database.model.task.task import Task
 from apiserver.service_repo import endpoint, APICall
 from apiserver.services.utils import get_tags_filter_dictionary, sort_tags_response
 
@@ -41,3 +47,28 @@ def get_user_companies(call: APICall, company_id: str, _):
             }
         ]
     }
+
+
+@endpoint("organization.get_entities_count", request_data_model=EntitiesCountRequest)
+def get_entities_count(call: APICall, company, _):
+    entity_classes: Mapping[str, Type[AttributedDocument]] = {
+        "projects": Project,
+        "tasks": Task,
+        "models": Model,
+        "pipelines": Project,
+    }
+    ret = {}
+    for field, entity_cls in entity_classes.items():
+        data = call.data.get(field)
+        if data is None:
+            continue
+
+        query = Q()
+        if entity_cls in (Project, Task) and not data.get("search_hidden"):
+            query &= Q(system_tags__ne=EntityVisibility.hidden.value)
+
+        ret[field] = entity_cls.get_count(
+            company=company, query_dict=data, query=query, allow_public=True,
+        )
+
+    call.result.data = ret