From 01b0427fc8f906270bbef5941a7901e85faab92a Mon Sep 17 00:00:00 2001
From: Benjamin Webb <40066515+webb-ben@users.noreply.github.com>
Date: Wed, 22 Jan 2025 11:35:51 -0500
Subject: [PATCH] Return mainstem if matched

---
 src/nldi/api/plugins/FeaturePlugin.py  | 18 ++++++++++--------
 src/nldi/api/plugins/FlowlinePlugin.py | 13 ++++++-------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/src/nldi/api/plugins/FeaturePlugin.py b/src/nldi/api/plugins/FeaturePlugin.py
index d8b296d..b4a9107 100644
--- a/src/nldi/api/plugins/FeaturePlugin.py
+++ b/src/nldi/api/plugins/FeaturePlugin.py
@@ -181,15 +181,11 @@ def lookup_navigation(self, nav: str, srcinfo: Dict[str, Any]):
                 _return.append(self._sqlalchemy_to_feature(row, geojson, srcinfo))
         return _return
 
-    def _sqlalchemy_to_feature(self, feature, geojson, srcinfo) -> Dict[str, Any]:
-        try:
-            mainstem = feature.mainstem_lookup.uri
-        except AttributeError:
-            LOGGER.warning(f"Mainstem_lookup not found for {feature.identifier}")
-            mainstem = ""
+    def _sqlalchemy_to_feature(self, feature, geojson, srcinfo, mainstem="") -> Dict[str, Any]:
+        
         navigation = util.url_join(self.relative_url(srcinfo["source_suffix"]), feature.identifier, "navigation")
 
-        return {
+        _response = {
             "type": "Feature",
             "properties": {
                 "identifier": str(feature.identifier),
@@ -202,7 +198,13 @@ def _sqlalchemy_to_feature(self, feature, geojson, srcinfo) -> Dict[str, Any]:
                 "reachcode": feature.reachcode,
                 "measure": feature.measure,
                 "navigation": navigation,
-                # "mainstem": mainstem, #<< mainstem is not returned from current production API.
             },
             "geometry": json.loads(geojson),
         }
+
+        try:
+            _response["properties"]["mainstem"] = feature.mainstem_lookup.uri
+        except AttributeError:
+            LOGGER.warning(f"Mainstem_lookup not found for {feature.identifier}")
+
+        return _response
diff --git a/src/nldi/api/plugins/FlowlinePlugin.py b/src/nldi/api/plugins/FlowlinePlugin.py
index 3a27970..187c9ea 100644
--- a/src/nldi/api/plugins/FlowlinePlugin.py
+++ b/src/nldi/api/plugins/FlowlinePlugin.py
@@ -144,12 +144,6 @@ def _sqlalchemy_to_feature(self, item) -> Dict[str, Any]:
         """
         (feature, geojson) = item
 
-        try:
-            mainstem = item.mainstem_lookup.uri
-        except AttributeError:
-            LOGGER.info(f"Mainstem not found for {feature.nhdplus_comid}")
-            mainstem = ""
-
         navigation = util.url_join(self.relative_url, feature.nhdplus_comid, "navigation")
 
         _response = {
@@ -159,9 +153,14 @@ def _sqlalchemy_to_feature(self, item) -> Dict[str, Any]:
                 "source": "comid",
                 "sourceName": "NHDPlus comid",
                 "comid": str(feature.nhdplus_comid),
-                # "mainstem": mainstem,  ##< mainstem is not in current production reponse from labs.waterdata.usgs.gov
                 "navigation": navigation,
             },
             "geometry": json.loads(geojson),
         }
+
+        try:
+            _response["properties"]["mainstem"] = item.mainstem_lookup.uri
+        except AttributeError:
+            LOGGER.info(f"Mainstem not found for {feature.nhdplus_comid}")
+
         return _response