diff --git a/trustgraph-base/trustgraph/base/publisher.py b/trustgraph-base/trustgraph/base/publisher.py index bc302599..654588a7 100644 --- a/trustgraph-base/trustgraph/base/publisher.py +++ b/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): diff --git a/trustgraph-base/trustgraph/base/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py index 66c1a4fe..dfc7f791 100644 --- a/trustgraph-base/trustgraph/base/subscriber.py +++ b/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): diff --git a/trustgraph-cli/scripts/tg-load-kg-core b/trustgraph-cli/scripts/tg-load-kg-core index 46811bd3..8eb5b0db 100755 --- a/trustgraph-cli/scripts/tg-load-kg-core +++ b/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' ) diff --git a/trustgraph-cli/scripts/tg-save-kg-core b/trustgraph-cli/scripts/tg-save-kg-core index 298f2e84..4c8e5ffa 100755 --- a/trustgraph-cli/scripts/tg-save-kg-core +++ b/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)' diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py index 46ef805e..1f459081 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py index c27f0c60..70e78c87 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py index 84aa927c..9b59a0ed 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py +++ b/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