From 66190689ee9fa891da43ef1147abae2c5e9158dd Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Wed, 2 Oct 2024 16:41:19 +0100
Subject: [PATCH] Iron out collection problems

---
 trustgraph-cli/scripts/tg-graph-show          | 26 ++++++++++++++++---
 .../storage/doc_embeddings/qdrant/write.py    |  4 +--
 .../storage/graph_embeddings/qdrant/write.py  |  4 +--
 .../storage/triples/cassandra/write.py        | 21 ++++++++++-----
 4 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/trustgraph-cli/scripts/tg-graph-show b/trustgraph-cli/scripts/tg-graph-show
index a737c97b..ac5db93f 100755
--- a/trustgraph-cli/scripts/tg-graph-show
+++ b/trustgraph-cli/scripts/tg-graph-show
@@ -9,12 +9,17 @@ import os
 from trustgraph.clients.triples_query_client import TriplesQueryClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_user = 'trustgraph'
+default_collection = 'default'
 
-def show_graph(pulsar):
+def show_graph(pulsar, user, collection):
 
     tq = TriplesQueryClient(pulsar_host=pulsar)
 
-    rows = tq.request(None, None, None, limit=10_000_000)
+    rows = tq.request(
+        user=user, collection=collection,
+        s=None, p=None, o=None, limit=10_000_000
+    )
 
     for row in rows:
         print(row.s.value, row.p.value, row.o.value)
@@ -32,11 +37,26 @@ def main():
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
 
+    parser.add_argument(
+        '-u', '--user',
+        default=default_user,
+        help=f'User ID (default: {default_user})'
+    )
+
+    parser.add_argument(
+        '-c', '--collection',
+        default=default_collection,
+        help=f'Collection ID (default: {default_collection})'
+    )
+
     args = parser.parse_args()
 
     try:
 
-        show_graph(args.pulsar_host)
+        show_graph(
+            pulsar=args.pulsar_host, user=args.user,
+            collection=args.collection,
+        )
 
     except Exception as e:
 
diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
index ab5cdb7c..813c4f29 100644
--- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
+++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
@@ -37,7 +37,6 @@ class Processor(Consumer):
         )
 
         self.last_collection = None
-        self.last_dim = None
 
         self.client = QdrantClient(url=store_uri)
 
@@ -57,7 +56,7 @@ class Processor(Consumer):
                 str(dim)
             )
 
-            if dim != self.last_dim:
+            if collection != self.last_collection:
 
                 if not self.client.collection_exists(collection):
 
@@ -73,7 +72,6 @@ class Processor(Consumer):
                         raise e
 
                 self.last_collection = collection
-                self.last_dim = dim
 
             self.client.upsert(
                 collection_name=collection,
diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
index 90d3ac77..e27c2516 100755
--- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
+++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
@@ -37,7 +37,6 @@ class Processor(Consumer):
         )
 
         self.last_collection = None
-        self.last_dim = None
 
         self.client = QdrantClient(url=store_uri)
 
@@ -55,7 +54,7 @@ class Processor(Consumer):
                 str(dim)
             )
 
-            if dim != self.last_dim:
+            if collection != self.last_collection:
 
                 if not self.client.collection_exists(collection):
 
@@ -71,7 +70,6 @@ class Processor(Consumer):
                         raise e
 
                 self.last_collection = collection
-                self.last_dim = dim
 
             self.client.upsert(
                 collection_name=collection,
diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py
index 3d63d1e1..e8074c80 100755
--- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py
+++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py
@@ -47,13 +47,22 @@ class Processor(Consumer):
 
         table = (v.metadata.user, v.metadata.collection)
 
-        if self.table is None or self.table != self.table:
-            self.tg = TrustGraph(
-                hosts=self.graph_host,
-                keyspace=v.metadata.user, table=v.metadata.collection,
-            )
-            self.table = table
+        if self.table is None or self.table != table:
 
+            self.tg = None
+
+            try:
+                self.tg = TrustGraph(
+                    hosts=self.graph_host,
+                    keyspace=v.metadata.user, table=v.metadata.collection,
+                )
+            except Exception as e:
+                print("Exception", e, flush=True)
+                time.sleep(1)
+                raise e
+
+            self.table = table
+        
         self.tg.insert(
             v.s.value,
             v.p.value,