load/save

trustgraph-base/trustgraph/base/publisher.py

@@ -15,17 +15,22 @@ class Publisher:
         self.q = asyncio.Queue(maxsize=max_size)
         self.chunking_enabled = chunking_enabled
         self.running = True
+        self.task = None
 
     async def start(self):
         self.task = asyncio.create_task(self.run())
 
     async def stop(self):
         self.running = False
-        await self.task
+
+        if self.task:
+            await self.task
 
     async def join(self):
         await self.stop()
-        await self.task
+
+        if self.task:
+            await self.task
 
     async def run(self):
 

trustgraph-base/trustgraph/base/subscriber.py

@@ -19,6 +19,7 @@ class Subscriber:
         self.lock = asyncio.Lock()
         self.running = True
         self.metrics = metrics
+        self.task = None
 
     def __del__(self):
         self.running = False
@@ -28,11 +29,15 @@ class Subscriber:
 
     async def stop(self):
         self.running = False
-        await self.task
+
+        if self.task:
+            await self.task
 
     async def join(self):
         await self.stop()
-        await self.task
+
+        if self.task:
+            await self.task
 
     async def run(self):
 

trustgraph-cli/scripts/tg-load-kg-core

@@ -200,7 +200,7 @@ async def run(running, **args):
     ge_q = asyncio.Queue(maxsize=10)
     t_q = asyncio.Queue(maxsize=10)
 
-    flow_id = args["flow"]
+    flow_id = args["flow_id"]
     url = args["url"]
 
     load_task = asyncio.create_task(
@@ -217,7 +217,7 @@ async def run(running, **args):
         load_ge(
             running = running,
             queue = ge_q,
-            url = f"{url}api/v1/{flow_id}/import/graph-embeddings"
+            url = f"{url}api/v1/flow/{flow_id}/import/graph-embeddings"
         )
     )
 
@@ -225,7 +225,7 @@ async def run(running, **args):
         load_triples(
             running = running,
             queue = t_q,
-            url = f"{url}api/v1/{flow_id}/import/triples"
+            url = f"{url}api/v1/flow/{flow_id}/import/triples"
         )
     )
 
@@ -265,7 +265,6 @@ async def main(running):
 
     parser.add_argument(
         '-f', '--flow-id',
-        # Make it mandatory, difficult to over-write an existing file
         default="0000",
         help=f'Flow ID'
     )

trustgraph-cli/scripts/tg-save-kg-core

@@ -27,7 +27,7 @@ async def fetch_ge(running, queue, user, collection, url):
 
     async with aiohttp.ClientSession() as session:
 
-        async with session.ws_connect(f"{url}stream/graph-embeddings") as ws:
+        async with session.ws_connect(url) as ws:
 
             while running.get():
 
@@ -74,7 +74,7 @@ async def fetch_triples(running, queue, user, collection, url):
 
     async with aiohttp.ClientSession() as session:
 
-        async with session.ws_connect(f"{url}stream/triples") as ws:
+        async with session.ws_connect(url) as ws:
 
             while running.get():
 
@@ -160,11 +160,14 @@ async def run(running, **args):
 
     q = asyncio.Queue()
 
+    flow_id = args["flow_id"]
+    url = args["url"]
+
     ge_task = asyncio.create_task(
         fetch_ge(
             running=running,
             queue=q, user=args["user"], collection=args["collection"],
-            url=args["url"] + "api/v1/"
+            url = f"{url}api/v1/flow/{flow_id}/import/graph-embeddings"
         )
     )
 
@@ -172,7 +175,7 @@ async def run(running, **args):
         fetch_triples(
             running=running, queue=q,
             user=args["user"], collection=args["collection"],
-            url=args["url"] + "api/v1/"
+            url = f"{url}api/v1/flow/{flow_id}/import/triples"
         )
     )
 
@@ -224,6 +227,12 @@ async def main(running):
         help=f'Output format (default: msgpack)',
     )
 
+    parser.add_argument(
+        '-f', '--flow-id',
+        default="0000",
+        help=f'Flow ID'
+    )
+
     parser.add_argument(
         '--user',
         help=f'User ID to filter on (default: no filter)'

trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py

@@ -24,7 +24,10 @@ class DocumentEmbeddingsImport:
 
     async def destroy(self):
         self.running.stop()
-        await self.ws.close()
+
+        if self.ws:
+            await self.ws.close()
+
         await self.publisher.stop()
 
     async def receive(self, msg):
@@ -54,6 +57,8 @@ class DocumentEmbeddingsImport:
         while self.running.get():
             await asyncio.sleep(0.5)
 
-        await self.ws.close()
+        if self.ws:
+            await self.ws.close()
+
         self.ws = None
 

trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py

@@ -24,7 +24,10 @@ class GraphEmbeddingsImport:
 
     async def destroy(self):
         self.running.stop()
-        await self.ws.close()
+
+        if self.ws:
+            await self.ws.close()
+
         await self.publisher.stop()
 
     async def receive(self, msg):
@@ -54,6 +57,8 @@ class GraphEmbeddingsImport:
         while self.running.get():
             await asyncio.sleep(0.5)
 
-        await self.ws.close()
+        if self.ws:
+            await self.ws.close()
+
         self.ws = None
 

trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py

@@ -24,7 +24,10 @@ class TriplesImport:
 
     async def destroy(self):
         self.running.stop()
-        await self.ws.close()
+
+        if self.ws:
+            await self.ws.close()
+
         await self.publisher.stop()
 
     async def receive(self, msg):
@@ -48,6 +51,8 @@ class TriplesImport:
         while self.running.get():
             await asyncio.sleep(0.5)
 
-        await self.ws.close()
+        if self.ws:
+            await self.ws.close()
+
         self.ws = None