Handle messages properly

This commit is contained in:
Cyber MacGeddon 2025-09-29 23:27:13 +01:00
parent f46b8bb008
commit 34a65ee408
10 changed files with 92 additions and 82 deletions

View file

@ -95,16 +95,17 @@ class Processor(DocumentEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -119,17 +120,17 @@ class Processor(DocumentEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for document embeddings"""
try: try:
self.vecstore.delete_collection(message.user, message.collection) self.vecstore.delete_collection(request.user, request.collection)
# Send success response # Send success response
response = StorageManagementResponse( response = StorageManagementResponse(
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -200,16 +200,17 @@ class Processor(DocumentEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -224,10 +225,10 @@ class Processor(DocumentEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for document embeddings"""
try: try:
index_name = f"d-{message.user}-{message.collection}" index_name = f"d-{request.user}-{request.collection}"
if self.pinecone.has_index(index_name): if self.pinecone.has_index(index_name):
self.pinecone.delete_index(index_name) self.pinecone.delete_index(index_name)
@ -240,7 +241,7 @@ class Processor(DocumentEmbeddingsStoreService):
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -143,16 +143,17 @@ class Processor(DocumentEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -167,10 +168,10 @@ class Processor(DocumentEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for document embeddings"""
try: try:
collection_name = f"d_{message.user}_{message.collection}" collection_name = f"d_{request.user}_{request.collection}"
if self.qdrant.collection_exists(collection_name): if self.qdrant.collection_exists(collection_name):
self.qdrant.delete_collection(collection_name) self.qdrant.delete_collection(collection_name)
@ -183,7 +184,7 @@ class Processor(DocumentEmbeddingsStoreService):
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -91,16 +91,17 @@ class Processor(GraphEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -115,17 +116,17 @@ class Processor(GraphEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for graph embeddings"""
try: try:
self.vecstore.delete_collection(message.user, message.collection) self.vecstore.delete_collection(request.user, request.collection)
# Send success response # Send success response
response = StorageManagementResponse( response = StorageManagementResponse(
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -199,16 +199,17 @@ class Processor(GraphEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -223,10 +224,10 @@ class Processor(GraphEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for graph embeddings"""
try: try:
index_name = f"t-{message.user}-{message.collection}" index_name = f"t-{request.user}-{request.collection}"
if self.pinecone.has_index(index_name): if self.pinecone.has_index(index_name):
self.pinecone.delete_index(index_name) self.pinecone.delete_index(index_name)
@ -239,7 +240,7 @@ class Processor(GraphEmbeddingsStoreService):
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -150,16 +150,17 @@ class Processor(GraphEmbeddingsStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -174,10 +175,10 @@ class Processor(GraphEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for graph embeddings"""
try: try:
collection_name = f"t_{message.user}_{message.collection}" collection_name = f"t_{request.user}_{request.collection}"
if self.qdrant.collection_exists(collection_name): if self.qdrant.collection_exists(collection_name):
self.qdrant.delete_collection(collection_name) self.qdrant.delete_collection(collection_name)
@ -190,7 +191,7 @@ class Processor(GraphEmbeddingsStoreService):
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -125,16 +125,17 @@ class Processor(TriplesStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -149,31 +150,31 @@ class Processor(TriplesStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete all data for a specific collection from the unified triples table"""
try: try:
# Create or reuse connection for this user's keyspace # 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 self.tg = None
try: try:
if self.cassandra_username and self.cassandra_password: if self.cassandra_username and self.cassandra_password:
self.tg = KnowledgeGraph( self.tg = KnowledgeGraph(
hosts=self.cassandra_host, hosts=self.cassandra_host,
keyspace=message.user, keyspace=request.user,
username=self.cassandra_username, username=self.cassandra_username,
password=self.cassandra_password password=self.cassandra_password
) )
else: else:
self.tg = KnowledgeGraph( self.tg = KnowledgeGraph(
hosts=self.cassandra_host, hosts=self.cassandra_host,
keyspace=message.user, keyspace=request.user,
) )
except Exception as e: 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 raise
self.table = message.user self.table = request.user
# Delete all triples for this collection from the unified table # Delete all triples for this collection from the unified table
# In the unified table schema, collection is the partition key # In the unified table schema, collection is the partition key
@ -183,8 +184,8 @@ class Processor(TriplesStoreService):
""" """
try: try:
self.tg.session.execute(delete_cql, (message.collection,)) self.tg.session.execute(delete_cql, (request.collection,))
logger.info(f"Deleted all triples for collection {message.collection} from keyspace {message.user}") logger.info(f"Deleted all triples for collection {request.collection} from keyspace {request.user}")
except Exception as e: except Exception as e:
logger.error(f"Failed to delete collection data: {e}") logger.error(f"Failed to delete collection data: {e}")
raise raise
@ -194,7 +195,7 @@ class Processor(TriplesStoreService):
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -193,16 +193,17 @@ class Processor(TriplesStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -217,28 +218,28 @@ class Processor(TriplesStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete the collection for FalkorDB triples"""
try: try:
# Delete all nodes and literals for this user/collection # Delete all nodes and literals for this user/collection
node_result = self.io.query( node_result = self.io.query(
"MATCH (n:Node {user: $user, collection: $collection}) DETACH DELETE n", "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( literal_result = self.io.query(
"MATCH (n:Literal {user: $user, collection: $collection}) DETACH DELETE n", "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 # Send success response
response = StorageManagementResponse( response = StorageManagementResponse(
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -325,16 +325,17 @@ class Processor(TriplesStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -349,7 +350,7 @@ class Processor(TriplesStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete all data for a specific collection"""
try: try:
with self.io.session(database=self.db) as session: with self.io.session(database=self.db) as session:
@ -357,7 +358,7 @@ class Processor(TriplesStoreService):
node_result = session.run( node_result = session.run(
"MATCH (n:Node {user: $user, collection: $collection}) " "MATCH (n:Node {user: $user, collection: $collection}) "
"DETACH DELETE n", "DETACH DELETE n",
user=message.user, collection=message.collection user=request.user, collection=request.collection
) )
nodes_deleted = node_result.consume().counters.nodes_deleted nodes_deleted = node_result.consume().counters.nodes_deleted
@ -365,20 +366,20 @@ class Processor(TriplesStoreService):
literal_result = session.run( literal_result = session.run(
"MATCH (n:Literal {user: $user, collection: $collection}) " "MATCH (n:Literal {user: $user, collection: $collection}) "
"DETACH DELETE n", "DETACH DELETE n",
user=message.user, collection=message.collection user=request.user, collection=request.collection
) )
literals_deleted = literal_result.consume().counters.nodes_deleted literals_deleted = literal_result.consume().counters.nodes_deleted
# Note: Relationships are automatically deleted with DETACH DELETE # 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 # Send success response
response = StorageManagementResponse( response = StorageManagementResponse(
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")

View file

@ -276,16 +276,17 @@ class Processor(TriplesStoreService):
async def on_storage_management(self, message, consumer, flow): async def on_storage_management(self, message, consumer, flow):
"""Handle storage management requests""" """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: try:
if message.operation == "delete-collection": if request.operation == "delete-collection":
await self.handle_delete_collection(message) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
error=Error( error=Error(
type="invalid_operation", type="invalid_operation",
message=f"Unknown operation: {message.operation}" message=f"Unknown operation: {request.operation}"
) )
) )
await self.storage_response_producer.send(response) await self.storage_response_producer.send(response)
@ -300,7 +301,7 @@ class Processor(TriplesStoreService):
) )
await self.storage_response_producer.send(response) 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""" """Delete all data for a specific collection"""
try: try:
with self.io.session(database=self.db) as session: with self.io.session(database=self.db) as session:
@ -308,7 +309,7 @@ class Processor(TriplesStoreService):
node_result = session.run( node_result = session.run(
"MATCH (n:Node {user: $user, collection: $collection}) " "MATCH (n:Node {user: $user, collection: $collection}) "
"DETACH DELETE n", "DETACH DELETE n",
user=message.user, collection=message.collection user=request.user, collection=request.collection
) )
nodes_deleted = node_result.consume().counters.nodes_deleted nodes_deleted = node_result.consume().counters.nodes_deleted
@ -316,20 +317,20 @@ class Processor(TriplesStoreService):
literal_result = session.run( literal_result = session.run(
"MATCH (n:Literal {user: $user, collection: $collection}) " "MATCH (n:Literal {user: $user, collection: $collection}) "
"DETACH DELETE n", "DETACH DELETE n",
user=message.user, collection=message.collection user=request.user, collection=request.collection
) )
literals_deleted = literal_result.consume().counters.nodes_deleted literals_deleted = literal_result.consume().counters.nodes_deleted
# Note: Relationships are automatically deleted with DETACH DELETE # 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 # Send success response
response = StorageManagementResponse( response = StorageManagementResponse(
error=None # No error means success error=None # No error means success
) )
await self.storage_response_producer.send(response) 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: except Exception as e:
logger.error(f"Failed to delete collection: {e}") logger.error(f"Failed to delete collection: {e}")