From 34a65ee408159de435afa2b93698c08fa0c2550b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 29 Sep 2025 23:27:13 +0100 Subject: [PATCH] Handle messages properly --- .../storage/doc_embeddings/milvus/write.py | 15 ++++++----- .../storage/doc_embeddings/pinecone/write.py | 15 ++++++----- .../storage/doc_embeddings/qdrant/write.py | 15 ++++++----- .../storage/graph_embeddings/milvus/write.py | 15 ++++++----- .../graph_embeddings/pinecone/write.py | 15 ++++++----- .../storage/graph_embeddings/qdrant/write.py | 15 ++++++----- .../storage/triples/cassandra/write.py | 27 ++++++++++--------- .../storage/triples/falkordb/write.py | 19 ++++++------- .../storage/triples/memgraph/write.py | 19 ++++++------- .../trustgraph/storage/triples/neo4j/write.py | 19 ++++++------- 10 files changed, 92 insertions(+), 82 deletions(-) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py index f723d5e1..b2fddeb9 100755 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py @@ -95,16 +95,17 @@ class Processor(DocumentEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -119,17 +120,17 @@ class Processor(DocumentEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for document embeddings""" try: - self.vecstore.delete_collection(message.user, message.collection) + self.vecstore.delete_collection(request.user, request.collection) # Send success response response = StorageManagementResponse( error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index 29092748..b0a318f7 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -200,16 +200,17 @@ class Processor(DocumentEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -224,10 +225,10 @@ class Processor(DocumentEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for document embeddings""" try: - index_name = f"d-{message.user}-{message.collection}" + index_name = f"d-{request.user}-{request.collection}" if self.pinecone.has_index(index_name): self.pinecone.delete_index(index_name) @@ -240,7 +241,7 @@ class Processor(DocumentEmbeddingsStoreService): error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index f53de0c5..3a2f8f56 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -143,16 +143,17 @@ class Processor(DocumentEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -167,10 +168,10 @@ class Processor(DocumentEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for document embeddings""" try: - collection_name = f"d_{message.user}_{message.collection}" + collection_name = f"d_{request.user}_{request.collection}" if self.qdrant.collection_exists(collection_name): self.qdrant.delete_collection(collection_name) @@ -183,7 +184,7 @@ class Processor(DocumentEmbeddingsStoreService): error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py index a93bc9a9..3058b452 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py @@ -91,16 +91,17 @@ class Processor(GraphEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -115,17 +116,17 @@ class Processor(GraphEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for graph embeddings""" try: - self.vecstore.delete_collection(message.user, message.collection) + self.vecstore.delete_collection(request.user, request.collection) # Send success response response = StorageManagementResponse( error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index cc845737..02d61a89 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -199,16 +199,17 @@ class Processor(GraphEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -223,10 +224,10 @@ class Processor(GraphEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for graph embeddings""" try: - index_name = f"t-{message.user}-{message.collection}" + index_name = f"t-{request.user}-{request.collection}" if self.pinecone.has_index(index_name): self.pinecone.delete_index(index_name) @@ -239,7 +240,7 @@ class Processor(GraphEmbeddingsStoreService): error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index e3eb0362..c7144871 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -150,16 +150,17 @@ class Processor(GraphEmbeddingsStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -174,10 +175,10 @@ class Processor(GraphEmbeddingsStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for graph embeddings""" try: - collection_name = f"t_{message.user}_{message.collection}" + collection_name = f"t_{request.user}_{request.collection}" if self.qdrant.collection_exists(collection_name): self.qdrant.delete_collection(collection_name) @@ -190,7 +191,7 @@ class Processor(GraphEmbeddingsStoreService): error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index 86cc1dc7..4ae915cb 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -125,16 +125,17 @@ class Processor(TriplesStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -149,31 +150,31 @@ class Processor(TriplesStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete all data for a specific collection from the unified triples table""" try: # Create or reuse connection for this user's keyspace - if self.table is None or self.table != message.user: + if self.table is None or self.table != request.user: self.tg = None try: if self.cassandra_username and self.cassandra_password: self.tg = KnowledgeGraph( hosts=self.cassandra_host, - keyspace=message.user, + keyspace=request.user, username=self.cassandra_username, password=self.cassandra_password ) else: self.tg = KnowledgeGraph( hosts=self.cassandra_host, - keyspace=message.user, + keyspace=request.user, ) except Exception as e: - logger.error(f"Failed to connect to Cassandra for user {message.user}: {e}") + logger.error(f"Failed to connect to Cassandra for user {request.user}: {e}") raise - self.table = message.user + self.table = request.user # Delete all triples for this collection from the unified table # In the unified table schema, collection is the partition key @@ -183,8 +184,8 @@ class Processor(TriplesStoreService): """ try: - self.tg.session.execute(delete_cql, (message.collection,)) - logger.info(f"Deleted all triples for collection {message.collection} from keyspace {message.user}") + self.tg.session.execute(delete_cql, (request.collection,)) + logger.info(f"Deleted all triples for collection {request.collection} from keyspace {request.user}") except Exception as e: logger.error(f"Failed to delete collection data: {e}") raise @@ -194,7 +195,7 @@ class Processor(TriplesStoreService): error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py b/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py index fe410da8..8687271a 100755 --- a/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/falkordb/write.py @@ -193,16 +193,17 @@ class Processor(TriplesStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -217,28 +218,28 @@ class Processor(TriplesStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete the collection for FalkorDB triples""" try: # Delete all nodes and literals for this user/collection node_result = self.io.query( "MATCH (n:Node {user: $user, collection: $collection}) DETACH DELETE n", - params={"user": message.user, "collection": message.collection} + params={"user": request.user, "collection": request.collection} ) literal_result = self.io.query( "MATCH (n:Literal {user: $user, collection: $collection}) DETACH DELETE n", - params={"user": message.user, "collection": message.collection} + params={"user": request.user, "collection": request.collection} ) - logger.info(f"Deleted {node_result.nodes_deleted} nodes and {literal_result.nodes_deleted} literals for collection {message.user}/{message.collection}") + logger.info(f"Deleted {node_result.nodes_deleted} nodes and {literal_result.nodes_deleted} literals for collection {request.user}/{request.collection}") # Send success response response = StorageManagementResponse( error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py index 40aaa6e7..24083166 100755 --- a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py @@ -325,16 +325,17 @@ class Processor(TriplesStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -349,7 +350,7 @@ class Processor(TriplesStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete all data for a specific collection""" try: with self.io.session(database=self.db) as session: @@ -357,7 +358,7 @@ class Processor(TriplesStoreService): node_result = session.run( "MATCH (n:Node {user: $user, collection: $collection}) " "DETACH DELETE n", - user=message.user, collection=message.collection + user=request.user, collection=request.collection ) nodes_deleted = node_result.consume().counters.nodes_deleted @@ -365,20 +366,20 @@ class Processor(TriplesStoreService): literal_result = session.run( "MATCH (n:Literal {user: $user, collection: $collection}) " "DETACH DELETE n", - user=message.user, collection=message.collection + user=request.user, collection=request.collection ) literals_deleted = literal_result.consume().counters.nodes_deleted # Note: Relationships are automatically deleted with DETACH DELETE - logger.info(f"Deleted {nodes_deleted} nodes and {literals_deleted} literals for {message.user}/{message.collection}") + logger.info(f"Deleted {nodes_deleted} nodes and {literals_deleted} literals for {request.user}/{request.collection}") # Send success response response = StorageManagementResponse( error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}") diff --git a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py index 23e601ef..4b6732b5 100755 --- a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py @@ -276,16 +276,17 @@ class Processor(TriplesStoreService): async def on_storage_management(self, message, consumer, flow): """Handle storage management requests""" - logger.info(f"Storage management request: {message.operation} for {message.user}/{message.collection}") + request = message.value() + logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if message.operation == "delete-collection": - await self.handle_delete_collection(message) + if request.operation == "delete-collection": + await self.handle_delete_collection(request) else: response = StorageManagementResponse( error=Error( type="invalid_operation", - message=f"Unknown operation: {message.operation}" + message=f"Unknown operation: {request.operation}" ) ) await self.storage_response_producer.send(response) @@ -300,7 +301,7 @@ class Processor(TriplesStoreService): ) await self.storage_response_producer.send(response) - async def handle_delete_collection(self, message): + async def handle_delete_collection(self, request): """Delete all data for a specific collection""" try: with self.io.session(database=self.db) as session: @@ -308,7 +309,7 @@ class Processor(TriplesStoreService): node_result = session.run( "MATCH (n:Node {user: $user, collection: $collection}) " "DETACH DELETE n", - user=message.user, collection=message.collection + user=request.user, collection=request.collection ) nodes_deleted = node_result.consume().counters.nodes_deleted @@ -316,20 +317,20 @@ class Processor(TriplesStoreService): literal_result = session.run( "MATCH (n:Literal {user: $user, collection: $collection}) " "DETACH DELETE n", - user=message.user, collection=message.collection + user=request.user, collection=request.collection ) literals_deleted = literal_result.consume().counters.nodes_deleted # Note: Relationships are automatically deleted with DETACH DELETE - logger.info(f"Deleted {nodes_deleted} nodes and {literals_deleted} literals for {message.user}/{message.collection}") + logger.info(f"Deleted {nodes_deleted} nodes and {literals_deleted} literals for {request.user}/{request.collection}") # Send success response response = StorageManagementResponse( error=None # No error means success ) await self.storage_response_producer.send(response) - logger.info(f"Successfully deleted collection {message.user}/{message.collection}") + logger.info(f"Successfully deleted collection {request.user}/{request.collection}") except Exception as e: logger.error(f"Failed to delete collection: {e}")