From 4e7e0cdf479e4e112ef2f35b442aa70539b1748e Mon Sep 17 00:00:00 2001
From: jadmsaadaot <91914654+jadmsaadaot@users.noreply.github.com>
Date: Wed, 15 May 2024 09:47:41 -0700
Subject: [PATCH] Add the the consultation level when doing an first nation
 import (#2231)

* Add the the consultation level when doing an first nation import

* fix linting issues
---
 .../src/api/models/indigenous_work_queries.py | 29 +++++++++++++++++++
 epictrack-api/src/api/services/work.py        | 22 +++++++++-----
 2 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 epictrack-api/src/api/models/indigenous_work_queries.py

diff --git a/epictrack-api/src/api/models/indigenous_work_queries.py b/epictrack-api/src/api/models/indigenous_work_queries.py
new file mode 100644
index 000000000..cbfd8bc23
--- /dev/null
+++ b/epictrack-api/src/api/models/indigenous_work_queries.py
@@ -0,0 +1,29 @@
+# Copyright © 2019 Province of British Columbia
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Model to handle complex queries related to Indigenous Work."""
+
+from sqlalchemy.orm import joinedload
+
+from api.models import db, IndigenousWork, Work
+
+
+def find_all_by_project_id(project_id: int):
+    """Return all Indigenous Work for a project."""
+    query = (db.session.query(IndigenousWork)
+             .join(Work, IndigenousWork.work_id == Work.id)
+             .filter(Work.project_id == project_id)
+             .options(joinedload(IndigenousWork.work))
+             .distinct(Work.project_id, IndigenousWork.indigenous_nation_id)
+             )
+    return query.all()
diff --git a/epictrack-api/src/api/services/work.py b/epictrack-api/src/api/services/work.py
index 8773e7e4b..0ee3387fb 100644
--- a/epictrack-api/src/api/services/work.py
+++ b/epictrack-api/src/api/services/work.py
@@ -15,7 +15,6 @@
 from collections import defaultdict
 from datetime import datetime, timedelta
 from io import BytesIO
-from itertools import product
 from typing import Dict, List, Optional
 
 import pandas as pd
@@ -52,6 +51,7 @@
 from api.models.event_template import EventTemplateVisibilityEnum
 from api.models.indigenous_nation import IndigenousNation
 from api.models.indigenous_work import IndigenousWork
+from api.models.indigenous_work_queries import find_all_by_project_id
 from api.models.pagination_options import PaginationOptions
 from api.models.phase_code import PhaseVisibilityEnum
 from api.models.special_field import EntityEnum
@@ -834,13 +834,21 @@ def import_first_nations(cls, work_id: int, indigenous_nation_ids: [int]):
         ).update({"is_active": True})
         current_app.logger.info(f"Enabled {enabled_count} IndigenousWorks")
 
-        keys = ("work_id", "indigenous_nation_id")
-        indigenous_works = [
-            task_assignee
-            for i, j in product([work_id], indigenous_nation_ids)
-            if (task_assignee := dict(zip(keys, (i, j)))) not in existing_first_nations
+        work = Work.find_by_id(work_id)
+        nations_in_same_project = find_all_by_project_id(work.project_id)
+        selected_nations = [
+            nation for nation in nations_in_same_project
+            if nation.indigenous_nation_id in indigenous_nation_ids
+        ]
+        nations_to_insert = [
+            {
+                'work_id': work.id,
+                'indigenous_nation_id': nation.indigenous_nation_id,
+                'indigenous_consultation_level_id': nation.indigenous_consultation_level_id,
+            }
+            for nation in selected_nations
         ]
-        db.session.bulk_insert_mappings(IndigenousWork, mappings=indigenous_works)
+        db.session.bulk_insert_mappings(IndigenousWork, mappings=nations_to_insert)
         db.session.commit()
         return "Imported successfully"