Remove implicit collection creation

This commit is contained in:
Cyber MacGeddon 2025-09-30 12:36:23 +01:00
parent 9ec52d7c65
commit df80995d2e
6 changed files with 149 additions and 55 deletions

View file

@ -29,6 +29,7 @@ class KnowledgeGraph:
self.po_table = "triples_p" self.po_table = "triples_p"
self.object_table = "triples_o" self.object_table = "triples_o"
self.collection_table = "triples_collection" # For SPO queries and deletion 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: if username and password:
ssl_context = SSLContext(PROTOCOL_TLSv1_2) 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): def prepare_statements(self):
"""Prepare statements for optimal performance""" """Prepare statements for optimal performance"""
@ -248,10 +259,10 @@ class KnowledgeGraph:
) )
def collection_exists(self, collection): def collection_exists(self, collection):
"""Check if collection exists by querying triples_collection table""" """Check if collection exists by querying collection_metadata table"""
try: try:
result = self.session.execute( 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,) (collection,)
) )
return bool(list(result)) return bool(list(result))
@ -260,15 +271,14 @@ class KnowledgeGraph:
return False return False
def create_collection(self, collection): def create_collection(self, collection):
"""Create collection by inserting marker row in triples_collection table""" """Create collection by inserting metadata row"""
try: try:
# Insert a system marker triple to establish collection exists import datetime
# This won't interfere with data as it uses a reserved system namespace
self.session.execute( self.session.execute(
f"INSERT INTO {self.collection_table} (collection, s, p, o) VALUES (%s, %s, %s, %s)", f"INSERT INTO {self.collection_metadata_table} (collection, created_at) VALUES (%s, %s)",
(collection, "__system__", "__collection_created__", "__marker__") (collection, datetime.datetime.now())
) )
logger.info(f"Created collection marker for {collection}") logger.info(f"Created collection metadata for {collection}")
except Exception as e: except Exception as e:
logger.error(f"Error creating collection: {e}") logger.error(f"Error creating collection: {e}")
raise e raise e
@ -324,6 +334,12 @@ class KnowledgeGraph:
if count % 100 != 0: if count % 100 != 0:
self.session.execute(batch) 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}") logger.info(f"Deleted {count} triples from collection {collection}")
def close(self): def close(self):

View file

@ -39,20 +39,9 @@ class Processor(DocumentEmbeddingsQueryService):
self.qdrant = QdrantClient(url=store_uri, api_key=api_key) self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def ensure_collection_exists(self, collection, dim): def collection_exists(self, collection):
"""Ensure collection exists, create if it doesn't""" """Check if collection exists (no implicit creation)"""
if not self.qdrant.collection_exists(collection): return 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
async def query_document_embeddings(self, msg): async def query_document_embeddings(self, msg):
@ -60,15 +49,16 @@ class Processor(DocumentEmbeddingsQueryService):
chunks = [] chunks = []
for vec in msg.vectors:
dim = len(vec)
collection = ( collection = (
"d_" + msg.user + "_" + msg.collection "d_" + msg.user + "_" + msg.collection
) )
self.ensure_collection_exists(collection, dim) # 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:
search_result = self.qdrant.query_points( search_result = self.qdrant.query_points(
collection_name=collection, collection_name=collection,
query=vec, query=vec,

View file

@ -39,20 +39,9 @@ class Processor(GraphEmbeddingsQueryService):
self.qdrant = QdrantClient(url=store_uri, api_key=api_key) self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def ensure_collection_exists(self, collection, dim): def collection_exists(self, collection):
"""Ensure collection exists, create if it doesn't""" """Check if collection exists (no implicit creation)"""
if not self.qdrant.collection_exists(collection): return 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 create_value(self, ent): def create_value(self, ent):
if ent.startswith("http://") or ent.startswith("https://"): if ent.startswith("http://") or ent.startswith("https://"):
@ -67,14 +56,16 @@ class Processor(GraphEmbeddingsQueryService):
entity_set = set() entity_set = set()
entities = [] entities = []
for vec in msg.vectors:
dim = len(vec)
collection = ( collection = (
"t_" + msg.user + "_" + msg.collection "t_" + msg.user + "_" + msg.collection
) )
self.ensure_collection_exists(collection, dim) # 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:
# Heuristic hack, get (2*limit), so that we have more chance # Heuristic hack, get (2*limit), so that we have more chance
# of getting (limit) entities # of getting (limit) entities

View file

@ -140,7 +140,9 @@ class Processor(DocumentEmbeddingsStoreService):
logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}")
try: 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) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
@ -161,6 +163,39 @@ class Processor(DocumentEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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): async def handle_delete_collection(self, request):
"""Delete the collection for document embeddings""" """Delete the collection for document embeddings"""
try: try:

View file

@ -147,7 +147,9 @@ class Processor(GraphEmbeddingsStoreService):
logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}") logger.info(f"Storage management request: {request.operation} for {request.user}/{request.collection}")
try: 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) await self.handle_delete_collection(request)
else: else:
response = StorageManagementResponse( response = StorageManagementResponse(
@ -168,6 +170,39 @@ class Processor(GraphEmbeddingsStoreService):
) )
await self.storage_response_producer.send(response) 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): async def handle_delete_collection(self, request):
"""Delete the collection for graph embeddings""" """Delete the collection for graph embeddings"""
try: try:

View file

@ -434,7 +434,16 @@ class Processor(FlowProcessor):
logger.info(f"Received storage management request: {msg.operation} for {msg.user}/{msg.collection}") logger.info(f"Received storage management request: {msg.operation} for {msg.user}/{msg.collection}")
try: 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) await self.delete_collection(msg.user, msg.collection)
# Send success response # Send success response
@ -467,6 +476,24 @@ class Processor(FlowProcessor):
) )
await self.storage_response_producer.send(response) 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): async def delete_collection(self, user: str, collection: str):
"""Delete all data for a specific collection""" """Delete all data for a specific collection"""
# Connect if not already connected # Connect if not already connected