From df80995d2e44d7f277c98477033fbfe15b8cc165 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 30 Sep 2025 12:36:23 +0100 Subject: [PATCH] Remove implicit collection creation --- .../trustgraph/direct/cassandra_kg.py | 34 ++++++++++++----- .../query/doc_embeddings/qdrant/service.py | 34 ++++++----------- .../query/graph_embeddings/qdrant/service.py | 33 ++++++----------- .../storage/doc_embeddings/qdrant/write.py | 37 ++++++++++++++++++- .../storage/graph_embeddings/qdrant/write.py | 37 ++++++++++++++++++- .../storage/objects/cassandra/write.py | 29 ++++++++++++++- 6 files changed, 149 insertions(+), 55 deletions(-) diff --git a/trustgraph-flow/trustgraph/direct/cassandra_kg.py b/trustgraph-flow/trustgraph/direct/cassandra_kg.py index f44c8975..a43a346a 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra_kg.py +++ b/trustgraph-flow/trustgraph/direct/cassandra_kg.py @@ -29,6 +29,7 @@ class KnowledgeGraph: self.po_table = "triples_p" self.object_table = "triples_o" self.collection_table = "triples_collection" # For SPO queries and deletion + self.collection_metadata_table = "collection_metadata" # For tracking which collections exist if username and password: ssl_context = SSLContext(PROTOCOL_TLSv1_2) @@ -116,7 +117,17 @@ class KnowledgeGraph: ); """); - logger.info("Optimized multi-table schema initialized (4 tables)") + # Table 5: Collection metadata tracking + # Tracks which collections exist without polluting triple data + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.collection_metadata_table} ( + collection text, + created_at timestamp, + PRIMARY KEY (collection) + ); + """); + + logger.info("Optimized multi-table schema initialized (5 tables)") def prepare_statements(self): """Prepare statements for optimal performance""" @@ -248,10 +259,10 @@ class KnowledgeGraph: ) def collection_exists(self, collection): - """Check if collection exists by querying triples_collection table""" + """Check if collection exists by querying collection_metadata table""" try: result = self.session.execute( - f"SELECT collection FROM {self.collection_table} WHERE collection = %s LIMIT 1", + f"SELECT collection FROM {self.collection_metadata_table} WHERE collection = %s LIMIT 1", (collection,) ) return bool(list(result)) @@ -260,15 +271,14 @@ class KnowledgeGraph: return False def create_collection(self, collection): - """Create collection by inserting marker row in triples_collection table""" + """Create collection by inserting metadata row""" try: - # Insert a system marker triple to establish collection exists - # This won't interfere with data as it uses a reserved system namespace + import datetime self.session.execute( - f"INSERT INTO {self.collection_table} (collection, s, p, o) VALUES (%s, %s, %s, %s)", - (collection, "__system__", "__collection_created__", "__marker__") + f"INSERT INTO {self.collection_metadata_table} (collection, created_at) VALUES (%s, %s)", + (collection, datetime.datetime.now()) ) - logger.info(f"Created collection marker for {collection}") + logger.info(f"Created collection metadata for {collection}") except Exception as e: logger.error(f"Error creating collection: {e}") raise e @@ -324,6 +334,12 @@ class KnowledgeGraph: if count % 100 != 0: self.session.execute(batch) + # Step 3: Delete collection metadata + self.session.execute( + f"DELETE FROM {self.collection_metadata_table} WHERE collection = %s", + (collection,) + ) + logger.info(f"Deleted {count} triples from collection {collection}") def close(self): diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index 0183bd58..dd79315e 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -39,20 +39,9 @@ class Processor(DocumentEmbeddingsQueryService): self.qdrant = QdrantClient(url=store_uri, api_key=api_key) - def ensure_collection_exists(self, collection, dim): - """Ensure collection exists, create if it doesn't""" - if not self.qdrant.collection_exists(collection): - try: - self.qdrant.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.COSINE - ), - ) - logger.info(f"Created collection: {collection}") - except Exception as e: - logger.error(f"Qdrant collection creation failed: {e}") - raise e + def collection_exists(self, collection): + """Check if collection exists (no implicit creation)""" + return self.qdrant.collection_exists(collection) async def query_document_embeddings(self, msg): @@ -60,15 +49,16 @@ class Processor(DocumentEmbeddingsQueryService): chunks = [] + collection = ( + "d_" + msg.user + "_" + msg.collection + ) + + # Check if collection exists - return empty if not + if not self.collection_exists(collection): + logger.info(f"Collection {collection} does not exist, returning empty results") + return [] + for vec in msg.vectors: - - dim = len(vec) - collection = ( - "d_" + msg.user + "_" + msg.collection - ) - - self.ensure_collection_exists(collection, dim) - search_result = self.qdrant.query_points( collection_name=collection, query=vec, diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index dbc44b92..459ec2a0 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -39,20 +39,9 @@ class Processor(GraphEmbeddingsQueryService): self.qdrant = QdrantClient(url=store_uri, api_key=api_key) - def ensure_collection_exists(self, collection, dim): - """Ensure collection exists, create if it doesn't""" - if not self.qdrant.collection_exists(collection): - try: - self.qdrant.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.COSINE - ), - ) - logger.info(f"Created collection: {collection}") - except Exception as e: - logger.error(f"Qdrant collection creation failed: {e}") - raise e + def collection_exists(self, collection): + """Check if collection exists (no implicit creation)""" + return self.qdrant.collection_exists(collection) def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -67,15 +56,17 @@ class Processor(GraphEmbeddingsQueryService): entity_set = set() entities = [] + collection = ( + "t_" + msg.user + "_" + msg.collection + ) + + # Check if collection exists - return empty if not + if not self.collection_exists(collection): + logger.info(f"Collection {collection} does not exist, returning empty results") + return [] + for vec in msg.vectors: - dim = len(vec) - collection = ( - "t_" + msg.user + "_" + msg.collection - ) - - self.ensure_collection_exists(collection, dim) - # Heuristic hack, get (2*limit), so that we have more chance # of getting (limit) entities search_result = self.qdrant.query_points( diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 2a608648..8898b8e9 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -140,7 +140,9 @@ class Processor(DocumentEmbeddingsStoreService): logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if request.operation == "delete-collection": + if request.operation == "create-collection": + await self.handle_create_collection(request) + elif request.operation == "delete-collection": await self.handle_delete_collection(request) else: response = StorageManagementResponse( @@ -161,6 +163,39 @@ class Processor(DocumentEmbeddingsStoreService): ) await self.storage_response_producer.send(response) + async def handle_create_collection(self, request): + """Create a Qdrant collection for document embeddings""" + try: + collection_name = f"d_{request.user}_{request.collection}" + + if self.qdrant.collection_exists(collection_name): + logger.info(f"Qdrant collection {collection_name} already exists") + else: + # Create collection with default dimension (will be recreated with correct dim on first write if needed) + # Using a placeholder dimension - actual dimension determined by first embedding + self.qdrant.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=384, # Default dimension, common for many models + distance=Distance.COSINE + ) + ) + logger.info(f"Created Qdrant collection: {collection_name}") + + # Send success response + response = StorageManagementResponse(error=None) + await self.storage_response_producer.send(response) + + except Exception as e: + logger.error(f"Failed to create collection: {e}", exc_info=True) + response = StorageManagementResponse( + error=Error( + type="creation_error", + message=str(e) + ) + ) + await self.storage_response_producer.send(response) + async def handle_delete_collection(self, request): """Delete the collection for document embeddings""" try: diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 24abc572..0d1b7674 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -147,7 +147,9 @@ class Processor(GraphEmbeddingsStoreService): logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") try: - if request.operation == "delete-collection": + if request.operation == "create-collection": + await self.handle_create_collection(request) + elif request.operation == "delete-collection": await self.handle_delete_collection(request) else: response = StorageManagementResponse( @@ -168,6 +170,39 @@ class Processor(GraphEmbeddingsStoreService): ) await self.storage_response_producer.send(response) + async def handle_create_collection(self, request): + """Create a Qdrant collection for graph embeddings""" + try: + collection_name = f"t_{request.user}_{request.collection}" + + if self.qdrant.collection_exists(collection_name): + logger.info(f"Qdrant collection {collection_name} already exists") + else: + # Create collection with default dimension (will be recreated with correct dim on first write if needed) + # Using a placeholder dimension - actual dimension determined by first embedding + self.qdrant.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=384, # Default dimension, common for many models + distance=Distance.COSINE + ) + ) + logger.info(f"Created Qdrant collection: {collection_name}") + + # Send success response + response = StorageManagementResponse(error=None) + await self.storage_response_producer.send(response) + + except Exception as e: + logger.error(f"Failed to create collection: {e}", exc_info=True) + response = StorageManagementResponse( + error=Error( + type="creation_error", + message=str(e) + ) + ) + await self.storage_response_producer.send(response) + async def handle_delete_collection(self, request): """Delete the collection for graph embeddings""" try: diff --git a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py index 2350c161..809d0261 100644 --- a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py @@ -434,7 +434,16 @@ class Processor(FlowProcessor): logger.info(f"Received storage management request: {msg.operation} for {msg.user}/{msg.collection}") try: - if msg.operation == "delete-collection": + if msg.operation == "create-collection": + await self.create_collection(msg.user, msg.collection) + + # Send success response + response = StorageManagementResponse( + error=None # No error means success + ) + await self.storage_response_producer.send(response) + logger.info(f"Successfully created collection {msg.user}/{msg.collection}") + elif msg.operation == "delete-collection": await self.delete_collection(msg.user, msg.collection) # Send success response @@ -467,6 +476,24 @@ class Processor(FlowProcessor): ) await self.storage_response_producer.send(response) + async def create_collection(self, user: str, collection: str): + """Create/verify collection exists in Cassandra object store""" + # Connect if not already connected + self.connect_cassandra() + + # Sanitize names for safety + safe_keyspace = self.sanitize_name(user) + + # Ensure keyspace exists + if safe_keyspace not in self.known_keyspaces: + self.ensure_keyspace(safe_keyspace) + self.known_keyspaces.add(safe_keyspace) + + # For Cassandra objects, collection is just a property in rows + # No need to create separate tables per collection + # Just mark that we've seen this collection + logger.info(f"Collection {collection} ready for user {user} (using keyspace {safe_keyspace})") + async def delete_collection(self, user: str, collection: str): """Delete all data for a specific collection""" # Connect if not already connected