dfs.py

import redis
import toml
import os


CONFIG = toml.load(os.path.join(os.path.dirname(__file__), "config", "config.toml"))


class DFS:
    def __init__(self, location="naoc"):
        if location not in CONFIG.keys():
            print("Available DFS locations: ", list(CONFIG.keys()))
            raise ValueError(f"Unknown dfs location: {location}")
        self.location = location
        self.config = CONFIG[location]
        print("Setting DFS config:")
        for k, v in self.config["dfs"].items():
            if k not in os.environ.keys():
                os.environ.setdefault(k, str(v))
                print(f" - Set {k}: {v}")
            else:
                print(f" - {k}: {os.environ[k]}")

        self.redis = Redis(location)


class Redis(redis.Redis):
    def __init__(self, location="naoc"):
        if location not in CONFIG.keys():
            print("Available redis locations: ", list(CONFIG.keys()))
            raise ValueError(f"Unknown redis location: {location}")
        super().__init__(
            host=CONFIG[location]["redis"]["host"],
            port=CONFIG[location]["redis"]["port"],
            db=CONFIG[location]["redis"]["db"],
            password=CONFIG[location]["redis"]["password"],
        )
        self.qname = password = CONFIG[location]["redis"]["qname"]
        self.config = CONFIG[location]["redis"]

        print("Setting redis config:")
        for k, v in self.config.items():
            print(f" - {k}: {v}")

    def push(self, msg):
        self.lpush(self.qname, msg)

    def pop(self):
        return self.rpop(self.qname)

    def get_all(self):
        return self.lrange(self.qname, 0, -1)


# msgs = r.lrange(name, 0, -1)
# for chipid in range(6, 26):
#     this_msg = gen_msg(
#         dag_id="csst-msc-l1-mbi", obsid="11009101682009", chipid=f"{chipid:02d}"
#     )
#     print(this_msg)
#     r.lpush(name, this_msg)
#
# msgs = r.lrange(name, 0, -1)
# print(msgs)
#
# msgs_later = r.lrange(name, 0, -1)
# print(msgs_later)