From 4c904059013ff26072734bc5e42e5fcbf4b7863b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 7 May 2025 11:39:49 +0100 Subject: [PATCH] Cassandra tables for config --- trustgraph-flow/trustgraph/tables/config.py | 261 ++++++++++++++++++++ 1 file changed, 261 insertions(+) create mode 100644 trustgraph-flow/trustgraph/tables/config.py diff --git a/trustgraph-flow/trustgraph/tables/config.py b/trustgraph-flow/trustgraph/tables/config.py new file mode 100644 index 00000000..3b555491 --- /dev/null +++ b/trustgraph-flow/trustgraph/tables/config.py @@ -0,0 +1,261 @@ + +from .. schema import KnowledgeResponse, Triple, Triples, EntityEmbeddings +from .. schema import Metadata, Value, GraphEmbeddings + +from cassandra.cluster import Cluster +from cassandra.auth import PlainTextAuthProvider +from ssl import SSLContext, PROTOCOL_TLSv1_2 + +import uuid +import time +import asyncio + +class ConfigTableStore: + + def __init__( + self, + cassandra_host, cassandra_user, cassandra_password, keyspace, + ): + + self.keyspace = keyspace + + print("Connecting to Cassandra...", flush=True) + + if cassandra_user and cassandra_password: + ssl_context = SSLContext(PROTOCOL_TLSv1_2) + auth_provider = PlainTextAuthProvider( + username=cassandra_user, password=cassandra_password + ) + self.cluster = Cluster( + cassandra_host, + auth_provider=auth_provider, + ssl_context=ssl_context + ) + else: + self.cluster = Cluster(cassandra_host) + + self.cassandra = self.cluster.connect() + + print("Connected.", flush=True) + + self.ensure_cassandra_schema() + + self.prepare_statements() + + def ensure_cassandra_schema(self): + + print("Ensure Cassandra schema...", flush=True) + + print("Keyspace...", flush=True) + + # FIXME: Replication factor should be configurable + self.cassandra.execute(f""" + create keyspace if not exists {self.keyspace} + with replication = {{ + 'class' : 'SimpleStrategy', + 'replication_factor' : 1 + }}; + """); + + self.cassandra.set_keyspace(self.keyspace) + + print("config table...", flush=True) + + self.cassandra.execute(""" + CREATE TABLE IF NOT EXISTS config ( + class text, + key text, + value text, + PRIMARY KEY (class, key) + ); + """); + + print("Cassandra schema OK.", flush=True) + + def prepare_statements(self): + + self.put_config_stmt = self.cassandra.prepare(""" + INSERT INTO config ( class, key, value ) + VALUES (?, ?, ?) + """) + + self.get_classes_stmt = self.cassandra.prepare(""" + SELECT DISTINCT class FROM config; + """) + + self.get_keys_stmt = self.cassandra.prepare(""" + SELECT key FROM config WHERE class = ?; + """) + + self.get_value_stmt = self.cassandra.prepare(""" + SELECT value FROM config WHERE class = ? AND key = ?; + """) + + self.delete_key_stmt = self.cassandra.prepare(""" + DELETE FROM config + WHERE class = ? AND key = ?; + """) + + self.get_all_stmt = self.cassandra.prepare(""" + SELECT class, key, value FROM config; + """) + + self.get_values_stmt = self.cassandra.prepare(""" + SELECT key, value FROM config WHERE class = ?; + """) + + async def put_config(self, cls, key, value): + + while True: + + try: + + resp = self.cassandra.execute( + self.put_config_stmt, + ( cls, key, value ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + async def get_value(self, cls, key): + + while True: + + try: + + resp = self.cassandra.execute( + self.get_value_stmt, + ( cls, key ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + for row in resp: + return row[0] + + return None + + async def get_values(self, cls): + + while True: + + try: + + resp = self.cassandra.execute( + self.get_values_stmt, + ( cls, ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + return [ + [row[0], row[1]] + for row in resp + ] + + async def get_classes(self): + + while True: + + try: + + resp = self.cassandra.execute( + self.get_classes_stmt, + () + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + return [ + row[0] for row in resp + ] + + async def get_all(self): + + while True: + + try: + + resp = self.cassandra.execute( + self.get_all_stmt, + () + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + return [ + (row[0], row[1], row[2]) + for row in resp + ] + + async def get_keys(self, cls): + + while True: + + try: + + resp = self.cassandra.execute( + self.get_keys_stmt, + ( cls, ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + return [ + row[0] for row in resp + ] + + async def delete_key(self, cls, key): + + while True: + + try: + + resp = self.cassandra.execute( + self.delete_key_stmt, + (cls, key) + ) + + break + + except Exception as e: + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) +