Move mongodb connection adapter from common to network.storage.states…

…tore
scaleoutsystems · Dec 7, 2023 · 77bedca · 77bedca
1 parent 3d2c89a
commit 77bedca
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 14 deletions.
diff --git a/fedn/cli/run_cmd.py b/fedn/cli/run_cmd.py
@@ -9,7 +9,7 @@
 from fedn.network.dashboard.restservice import (decode_auth_token,
                                                 encode_auth_token)
 from fedn.network.reducer import Reducer
-from fedn.network.statestore.mongostatestore import MongoStateStore
+from fedn.network.storage.statestore.mongostatestore import MongoStateStore
 
 from .main import main
 

diff --git a/fedn/fedn/network/api/server.py b/fedn/fedn/network/api/server.py
@@ -4,7 +4,7 @@
                                 get_network_config, get_statestore_config)
 from fedn.network.api.interface import API
 from fedn.network.controller.control import Control
-from fedn.network.statestore.mongostatestore import MongoStateStore
+from fedn.network.storage.statestore.mongostatestore import MongoStateStore
 
 statestore_config = get_statestore_config()
 network_id = get_network_config()

diff --git a/fedn/fedn/network/combiner/modelservice.py b/fedn/fedn/network/combiner/modelservice.py
@@ -146,13 +146,13 @@ def Upload(self, request_iterator, context):
         for request in request_iterator:
             if request.status == fedn.ModelStatus.IN_PROGRESS:
                 self.models.get_ptr(request.id).write(request.data)
-                self.models.set_meta(request.id, fedn.ModelStatus.IN_PROGRESS)
+                self.models.set_model_metadata(request.id, fedn.ModelStatus.IN_PROGRESS)
 
             if request.status == fedn.ModelStatus.OK and not request.data:
                 result = fedn.ModelResponse(id=request.id, status=fedn.ModelStatus.OK,
                                             message="Got model successfully.")
                 # self.models_metadata.update({request.id: fedn.ModelStatus.OK})
-                self.models.set_meta(request.id, fedn.ModelStatus.OK)
+                self.models.set_model_metadata(request.id, fedn.ModelStatus.OK)
                 self.models.get_ptr(request.id).flush()
                 self.models.get_ptr(request.id).close()
                 return result
@@ -168,7 +168,7 @@ def Download(self, request, context):
         :rtype: :class:`fedn.common.net.grpc.fedn_pb2.ModelResponse`
         """
         try:
-            if self.models.get_meta(request.id) != fedn.ModelStatus.OK:
+            if self.models.get_model_metadata(request.id) != fedn.ModelStatus.OK:
                 print("Error file is not ready", flush=True)
                 yield fedn.ModelResponse(id=request.id, data=None, status=fedn.ModelStatus.FAILED)
         except Exception:

diff --git a/fedn/fedn/network/combiner/server.py b/fedn/fedn/network/combiner/server.py
@@ -17,7 +17,7 @@
 from fedn.network.combiner.connect import ConnectorCombiner, Status
 from fedn.network.combiner.modelservice import ModelService
 from fedn.network.combiner.round import RoundController
-from fedn.network.statestore.mongostatestore import MongoStateStore
+from fedn.network.storage.statestore.mongostatestore import MongoStateStore
 
 VALID_NAME_REGEX = '^[a-zA-Z0-9_-]*$'
 

diff --git a/fedn/fedn/network/dashboard/plots.py b/fedn/fedn/network/dashboard/plots.py
@@ -13,7 +13,7 @@
 from networkx.algorithms import community
 from plotly.subplots import make_subplots
 
-from fedn.common.storage.db.mongo import connect_to_mongodb
+from fedn.network.storage.statestore.mongostatestore import MongoStateStore
 
 
 class Plot:
@@ -24,8 +24,9 @@ class Plot:
     def __init__(self, statestore):
         try:
             statestore_config = statestore.get_config()
-            self.mdb = connect_to_mongodb(
-                statestore_config['mongo_config'], statestore_config['network_id'])
+            statestore = MongoStateStore(
+                statestore_config['network_id'], statestore_config['mongo_config'])
+            self.mdb = statestore.connect()
             self.status = self.mdb['control.status']
             self.round_time = self.mdb["control.round_time"]
             self.combiner_round_time = self.mdb["control.combiner_round_time"]

diff --git a/fedn/fedn/network/storage/statestore/mongostatestore.py b/fedn/fedn/network/storage/statestore/mongostatestore.py
@@ -27,7 +27,7 @@ def __init__(self, network_id, config):
         try:
             self.config = config
             self.network_id = network_id
-            self.mdb = self.connect(self.config, self.network_id)
+            self.mdb = self.connect()
 
             # FEDn network
             self.network = self.mdb["network"]
@@ -59,8 +59,7 @@ def __init__(self, network_id, config):
             self.clients = None
             raise
 
-    @classmethod
-    def connect(config, network_id):
+    def connect(self):
         """ Establish client connection to MongoDB.
 
         :param config: Dictionary containing connection strings and security credentials.
@@ -70,10 +69,10 @@ def connect(config, network_id):
         :return: MongoDB client pointing to the db corresponding to network_id
         """
         try:
-            mc = pymongo.MongoClient(**config)
+            mc = pymongo.MongoClient(**self.config)
             # This is so that we check that the connection is live
             mc.server_info()
-            mdb = mc[network_id]
+            mdb = mc[self.network_id]
             return mdb
         except Exception:
             raise