Distributed-Union-Find/config/db/db_setup.py

from cassandra.cluster import Cluster
import os

def db_setup(config):
    print(" == DB SETUP SCRIPT == ")

    print(f"Attempting Cassandra connection @ {config['cassandra_addresses']}:{config['cassandra_port']}")
    cluster  = Cluster(config['cassandra_addresses'], port=config['cassandra_port'])
    session = cluster.connect()
    print(f"Connection OK")

    print(f"Creating KEYSPACE '{config['cassandra_keyspace']}'")
    with open(f"{config['setup_keyspace_dir']}/CREATE.sql") as keyspace_create:
        session.execute(keyspace_create.read())

    with open(f"{config['setup_keyspace_dir']}/USE.sql") as keyspace_use:
        session.execute(keyspace_use.read())

    for folder_name in os.listdir(config['setup_tables_dir']):
        print(f"Creating TABLE '{folder_name}'")
        with open(f"{config['setup_tables_dir']}/{folder_name}/CREATE.sql") as sql_create:
            session.execute(sql_create.read())

    print("Done!")
    cluster.shutdown()