Add workspace field to GraphQL Dataset (#4318)

* add minimal gql workspace definitions * graphql workspace test * lint imports * lint import
voxel51 · Apr 29, 2024 · bdfc97f · bdfc97f
1 parent f073ba4
commit bdfc97f
Show file tree

Hide file tree

Showing 9 changed files with 388 additions and 202 deletions.
diff --git a/app/packages/relay/src/queries/__generated__/datasetQuery.graphql.ts b/app/packages/relay/src/queries/__generated__/datasetQuery.graphql.ts
diff --git a/app/packages/relay/src/queries/dataset.ts b/app/packages/relay/src/queries/dataset.ts
@@ -2,10 +2,11 @@ import { graphql } from "relay-runtime";
 
 export default graphql`
   query datasetQuery(
-    $savedViewSlug: String
+    $extendedView: BSONArray!
     $name: String!
+    $savedViewSlug: String
     $view: BSONArray!
-    $extendedView: BSONArray!
+    $workspaceSlug: String
   ) {
     config {
       colorBy
@@ -71,6 +72,13 @@ export default graphql`
           }
         }
       }
+
+      workspace(slug: $workspaceSlug) {
+        id
+        child
+        slug
+      }
+
       ...datasetFragment
     }
     ...savedViewsFragment

diff --git a/app/schema.graphql b/app/schema.graphql
@@ -240,6 +240,7 @@ type Dataset {
   frameIndexes: [Index!]
   sampleIndexes: [Index!]
   stages(slug: String = null, view: BSONArray = null): BSONArray
+  workspace(slug: String): Workspace
 }
 
 type DatasetAppConfig {
@@ -739,8 +740,8 @@ input SidebarGroupInput {
 
 enum SidebarMode {
   all
-  disabled
   best
+  disabled
   fast
 }
 
@@ -833,6 +834,19 @@ type WildcardProjection {
   inclusion: Boolean!
 }
 
+type Workspace {
+  id: ID!
+  color: String
+  child: BSON!
+  datasetId: ID!
+  createdAt: datetime
+  description: String
+  name: String
+  lastModifiedAt: datetime
+  lastLoadedAt: datetime
+  slug: String
+}
+
 scalar date
 
 scalar datetime
diff --git a/fiftyone/server/query.py b/fiftyone/server/query.py
@@ -47,6 +47,7 @@
 from fiftyone.server.scalars import BSON, BSONArray, JSON
 from fiftyone.server.stage_definitions import stage_definitions
 from fiftyone.server.utils import from_dict
+from fiftyone.server.workspace import Workspace
 
 
 ID = gql.scalar(
@@ -291,6 +292,20 @@ async def estimated_frame_count(
                 self.frame_collection_name
             ].estimated_document_count()
 
+    @gql.field
+    async def workspace(
+        self, slug: t.Optional[str], info: Info
+    ) -> t.Optional[Workspace]:
+        if slug:
+            doc = await info.context.db["workspaces"].find_one({"slug": slug})
+
+            if doc:
+                doc["id"] = doc.pop("_id")
+                doc["dataset_id"] = doc.pop("_dataset_id")
+                return from_dict(Workspace, doc)
+
+        return None
+
     @staticmethod
     def modifier(doc: dict) -> dict:
         doc["id"] = doc.pop("_id")

diff --git a/fiftyone/server/workspace.py b/fiftyone/server/workspace.py
@@ -0,0 +1,29 @@
+"""
+FiftyOne Server workspace.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+
+from datetime import datetime
+import typing as t
+
+import strawberry as gql
+
+from fiftyone.server.scalars import BSON
+
+
+@gql.type
+class Workspace:
+    id: gql.ID
+
+    color: t.Optional[str]
+    child: BSON
+    dataset_id: gql.ID
+    created_at: t.Optional[datetime]
+    description: t.Optional[str]
+    name: t.Optional[str]
+    last_modified_at: t.Optional[datetime]
+    last_loaded_at: t.Optional[datetime]
+    slug: t.Optional[str]
diff --git a/tests/unittests/lightning_tests.py b/tests/unittests/lightning_tests.py
@@ -17,14 +17,14 @@
 import fiftyone as fo
 
 from fiftyone.server.constants import SCALAR_OVERRIDES
-from fiftyone.server.context import get_context
 from fiftyone.server.lightning import (
     lightning_resolver,
     LightningInput,
     LightningPathInput,
 )
 
 from decorators import drop_async_dataset
+from utils.graphql import execute
 
 
 @gql.type
@@ -711,17 +711,17 @@ def _add_samples(dataset: fo.Dataset, lower: t.Dict, upper: t.Dict):
 async def _execute(
     query: str, dataset: fo.Dataset, field: fo.Field, keys: t.Set[str]
 ):
-    return await schema.execute(
+    return await execute(
+        schema,
         query,
-        variable_values={
+        {
             "input": asdict(
                 LightningInput(
                     dataset=dataset.name,
                     paths=_get_paths(dataset, field, keys),
                 )
             )
         },
-        context_value=get_context(use_global_db_client=False),
     )
 
 

diff --git a/tests/unittests/server_aggregations_tests.py b/tests/unittests/server_aggregations_tests.py
@@ -6,7 +6,6 @@
 |
 """
 
-import typing as t
 import unittest
 
 import strawberry as gql
@@ -15,11 +14,11 @@
 import fiftyone as fo
 
 from fiftyone.server.constants import SCALAR_OVERRIDES
-from fiftyone.server.context import get_context
 from fiftyone.server.aggregate import AggregateQuery
 from fiftyone.server.aggregations import aggregate_resolver
 
 from decorators import drop_async_dataset
+from utils.graphql import execute
 
 
 @gql.type
@@ -55,7 +54,8 @@ async def test_group_mode_sidebar_counts(self, dataset: fo.Dataset):
             }
         """
 
-        result = await _execute(
+        result = await execute(
+            schema,
             query,
             {
                 "form": {
@@ -127,7 +127,8 @@ async def test_group_mode_histogram_counts(self, dataset: fo.Dataset):
             }
         """
 
-        result = await _execute(
+        result = await execute(
+            schema,
             query,
             {
                 "dataset": dataset.name,
@@ -173,11 +174,3 @@ def _add_samples(dataset: fo.Dataset):
             ),
         ]
     )
-
-
-async def _execute(query: str, variables: t.Dict):
-    return await schema.execute(
-        query,
-        variable_values=variables,
-        context_value=get_context(use_global_db_client=False),
-    )
diff --git a/tests/unittests/server_workspace_tests.py b/tests/unittests/server_workspace_tests.py
@@ -0,0 +1,72 @@
+"""
+FiftyOne Server workspace tests.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+
+import unittest
+
+import strawberry as gql
+from strawberry.schema.config import StrawberryConfig
+
+import fiftyone as fo
+
+from fiftyone.server.constants import SCALAR_OVERRIDES
+from fiftyone.server.query import Dataset
+
+from decorators import drop_async_dataset
+from utils.graphql import execute
+
+
+@gql.type
+class DatasetQuery(Dataset):
+    dataset: Dataset = gql.field(resolver=Dataset.resolver)
+
+
+schema = gql.Schema(
+    query=DatasetQuery,
+    scalar_overrides=SCALAR_OVERRIDES,
+    config=StrawberryConfig(auto_camel_case=False),
+)
+
+
+class TestGroupModeSidebarCounts(unittest.IsolatedAsyncioTestCase):
+    @drop_async_dataset
+    async def test_workspace_field(self, dataset: fo.Dataset):
+        histograms_panel = fo.Panel(type="Histograms")
+        name = "histograms"
+        workspace = fo.Space(children=[histograms_panel])
+        dataset.save_workspace(name, workspace)
+
+        query = """
+            query Query($name: String!, $slug: String!) {
+                dataset(name: $name) {
+                    workspace(slug: $slug) {
+                        child
+                        name
+                        slug
+                    }
+                }
+            }
+        """
+
+        result = await execute(
+            schema,
+            query,
+            {"name": dataset.name, "slug": name},
+        )
+
+        self.assertEqual(
+            result.data,
+            {
+                "dataset": {
+                    "workspace": {
+                        "child": workspace.to_dict(),
+                        "name": "histograms",
+                        "slug": name,
+                    }
+                }
+            },
+        )
diff --git a/tests/unittests/utils/graphql.py b/tests/unittests/utils/graphql.py
@@ -0,0 +1,23 @@
+import typing as t
+
+import strawberry as gql
+
+from fiftyone.server.context import get_context
+
+
+async def execute(schema: gql.Schema, query: str, variables: t.Dict):
+    """Execute a test GraphQL query.
+
+    Args:
+        schema: a :class:`strawberry.Schema`
+        query: a GraphQL query string
+        variables: a variables dictionary
+
+    Returns:
+        an :class:`strawberry.types.execution.ExecutionResult`
+    """
+    return await schema.execute(
+        query,
+        variable_values=variables,
+        context_value=get_context(use_global_db_client=False),
+    )