minor fix merge issue

chuongmep · Sep 10, 2024 · caeda6c · caeda6c
1 parent 1318d66
commit caeda6c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/APSToolkitPython/src/aps_toolkit/AECDataModel.py b/APSToolkitPython/src/aps_toolkit/AECDataModel.py
@@ -257,7 +257,7 @@ def get_element_projects_by_parameters(self, projectId: str, parameters: list[st
 
         return data_df
 
-    def get_elements_by_projects(self, projectId: str, cursor: str = None) -> pd.DataFrame:
+    def get_elements_by_projects(self, projectId: str, cursor: str = None, is_recursive: bool = False) -> pd.DataFrame:
         # Define the GraphQL query with the cursor as a dynamic variable
         data = {
             "query": """
@@ -305,11 +305,16 @@ def get_elements_by_projects(self, projectId: str, cursor: str = None) -> pd.Dat
         data_df = pd.DataFrame()
         for i, row in df.iterrows():
             props_dict = {}
+            id = row['id']
             single_df = pd.json_normalize(row['properties.results'])
             for i in range(len(single_df)):
                 props_dict[single_df['name'][i]] = single_df['value'][i]
+            props_dict['id'] = id
             single_df = pd.DataFrame(props_dict, index=[0], dtype="object")
             data_df = pd.concat([data_df, single_df], axis=0)
-
+        # recursive call to get all data if cursor is not None
+        cursor = result['data']['elementsByProject']['pagination']['cursor']
+        if cursor and is_recursive:
+            # merge by id
+            data_df = pd.concat([data_df, self.get_elements_by_projects(projectId, cursor)], axis=0)
         return data_df
-
diff --git a/APSToolkitPython/src/test/test_aec_data_model.py b/APSToolkitPython/src/test/test_aec_data_model.py
@@ -49,11 +49,7 @@ def test_get_element_projects_by_parameters(self):
                                                                          "Family Name", "Type Name", "Comments"])
         self.assertIsNotNone(result)
 
-    def test_get_element_projects_by_parameters2(self):
-        result = self.aec_data_model.get_element_projects_by_parameters2(self.project_id)
-        self.assertIsNotNone(result)
-
-    def test_get_element_projects_by_parameters3(self):
+    def test_get_elements_by_projects(self):
         cursor = "YWRjdXJzfjB-NTB-NTA"
         # cursor = ""
         result = self.aec_data_model.get_elements_by_projects(self.project_id, cursor)