Merge pull request #38 from internetofwater/mainstem

Return mainstem if matched
internetofwater · Jan 22, 2025 · a655021 · a655021
2 parents 3318815 + 01b0427
commit a655021
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 15 deletions.
diff --git 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
@@ -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