add get folders

chuongmep · Sep 10, 2024 · 26d510c · 26d510c
1 parent 7758d5e
commit 26d510c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/APSToolkitPython/src/aps_toolkit/AECDataModel.py b/APSToolkitPython/src/aps_toolkit/AECDataModel.py
@@ -74,4 +74,25 @@ def get_projects(self, hub_id: str) -> pd.DataFrame:
         }
         result = self.execute_query(data)
         projects = result['data']['projects']['results']
-        return pd.json_normalize(projects)
+        return pd.json_normalize(projects)
+
+    def get_folders(self, project_id: str) -> pd.DataFrame:
+        data = {
+            "query": """
+                query GetFolders($projectId: ID!) {
+                  foldersByProject(projectId: $projectId) {
+                    results {
+                      id
+                      name
+                      objectCount
+                    }
+                  }
+                }
+            """,
+            "variables": {
+                "projectId": project_id
+            }
+        }
+        result = self.execute_query_variables(data['query'], data['variables'])
+        folders = result['data']['foldersByProject']['results']
+        return pd.json_normalize(folders)
diff --git a/APSToolkitPython/src/test/test_aec_data_model.py b/APSToolkitPython/src/test/test_aec_data_model.py
@@ -30,6 +30,10 @@ def test_get_projects(self):
         result = self.aec_data_model.get_projects(self.hub_id)
         self.assertIsNotNone(result)
 
+    def test_get_folders(self):
+        result = self.aec_data_model.get_folders(self.project_id)
+        self.assertIsNotNone(result)
+
     def test_version_group_by_project(self):
         query = """
             query GetElementGroupsByProject($projectId: ID!) {