diff --git a/trustgraph-flow/trustgraph/direct/cassandra.py b/trustgraph-flow/trustgraph/direct/cassandra.py index 1754e090..2b577df1 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra.py +++ b/trustgraph-flow/trustgraph/direct/cassandra.py @@ -4,10 +4,16 @@ from cassandra.auth import PlainTextAuthProvider class TrustGraph: - def __init__(self, hosts=None): + def __init__( + self, hosts=None, + keyspace="trustgraph", table="default", + ): if hosts is None: hosts = ["localhost"] + + self.keyspace = keyspace + self.table = table self.cluster = Cluster(hosts) self.session = self.cluster.connect() @@ -16,26 +22,26 @@ class TrustGraph: def clear(self): - self.session.execute(""" - drop keyspace if exists trustgraph; + self.session.execute(f""" + drop keyspace if exists {self.keyspace}; """); self.init() def init(self): - self.session.execute(""" - create keyspace if not exists trustgraph - with replication = { + self.session.execute(f""" + create keyspace if not exists {self.keyspace} + with replication = {{ 'class' : 'SimpleStrategy', 'replication_factor' : 1 - }; + }}; """); - self.session.set_keyspace('trustgraph') + self.session.set_keyspace(self.keyspace) - self.session.execute(""" - create table if not exists triples ( + self.session.execute(f""" + create table if not exists {self.table} ( s text, p text, o text, @@ -43,66 +49,66 @@ class TrustGraph: ); """); - self.session.execute(""" - create index if not exists triples_p - ON triples (p); + self.session.execute(f""" + create index if not exists {self.table}_p + ON {self.table} (p); """); - self.session.execute(""" - create index if not exists triples_o - ON triples (o); + self.session.execute(f""" + create index if not exists {self.table}_o + ON {self.table} (o); """); def insert(self, s, p, o): self.session.execute( - "insert into triples (s, p, o) values (%s, %s, %s)", + f"insert into {self.table} (s, p, o) values (%s, %s, %s)", (s, p, o) ) def get_all(self, limit=50): return self.session.execute( - f"select s, p, o from triples limit {limit}" + f"select s, p, o from {self.table} limit {limit}" ) def get_s(self, s, limit=10): return self.session.execute( - f"select p, o from triples where s = %s limit {limit}", + f"select p, o from {self.table} where s = %s limit {limit}", (s,) ) def get_p(self, p, limit=10): return self.session.execute( - f"select s, o from triples where p = %s limit {limit}", + f"select s, o from {self.table} where p = %s limit {limit}", (p,) ) def get_o(self, o, limit=10): return self.session.execute( - f"select s, p from triples where o = %s limit {limit}", + f"select s, p from {self.table} where o = %s limit {limit}", (o,) ) def get_sp(self, s, p, limit=10): return self.session.execute( - f"select o from triples where s = %s and p = %s limit {limit}", + f"select o from {self.table} where s = %s and p = %s limit {limit}", (s, p) ) def get_po(self, p, o, limit=10): return self.session.execute( - f"select s from triples where p = %s and o = %s allow filtering limit {limit}", + f"select s from {self.table} where p = %s and o = %s allow filtering limit {limit}", (p, o) ) def get_os(self, o, s, limit=10): return self.session.execute( - f"select p from triples where o = %s and s = %s limit {limit}", + f"select p from {self.table} where o = %s and s = %s limit {limit}", (o, s) ) def get_spo(self, s, p, o, limit=10): return self.session.execute( - f"""select s as x from triples where s = %s and p = %s and o = %s limit {limit}""", + f"""select s as x from {self.table} where s = %s and p = %s and o = %s limit {limit}""", (s, p, o) ) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index cd350a32..7bb5133a 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -61,7 +61,7 @@ class Processor(ConsumerProducer): dim = len(vec) collection = ( - "d_" + v.user + "_" + v.collection + + "d_" + v.user + "_" + v.collection + "_" + str(dim) ) diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index 0f9a7b48..8991f9ea 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -67,7 +67,7 @@ class Processor(ConsumerProducer): dim = len(vec) collection = ( - "t_" + v.user + "_" + v.collection + + "t_" + v.user + "_" + v.collection + "_" + str(dim) ) diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 5e1e0e3e..4245784d 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -38,7 +38,8 @@ class Processor(ConsumerProducer): } ) - self.tg = TrustGraph([graph_host]) + self.graph_host = [graph_host] + self.table = None def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -52,6 +53,15 @@ class Processor(ConsumerProducer): v = msg.value() + table = (v.user, v.collection) + + if table != self.table: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.user, table=v.collection, + ) + self.table = table + # Sender-produced ID id = msg.properties()["id"] diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 8e2ee43f..ab5cdb7c 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -53,7 +53,7 @@ class Processor(Consumer): dim = len(vec) collection = ( - "d_" + v.metadata.user + "_" + v.metadata.collection + + "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + str(dim) ) diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 8f9964f6..90d3ac77 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -51,7 +51,7 @@ class Processor(Consumer): dim = len(vec) collection = ( - "t_" + v.metadata.user + "_" + v.metadata.collection + + "t_" + v.metadata.user + "_" + v.metadata.collection + "_" + str(dim) ) diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index 84c002ff..3d63d1e1 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -38,12 +38,22 @@ class Processor(Consumer): } ) - self.tg = TrustGraph([graph_host]) + self.graph_host = [graph_host] + self.table = None def handle(self, msg): v = msg.value() + 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 + self.tg.insert( v.s.value, v.p.value,