dfs.py

import redis
import toml
import os
from csst_dfs_client import plan, level0
from astropy.table import Table

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


class DFS:
    def __init__(self, location=None):
        # try each location
        print("Test all locations:")
        status_list = []
        for loc in CONFIG.keys():
            # print("Setting DFS config:")
            for k, v in CONFIG[loc]["dfs"].items():
                os.environ.setdefault(k, str(v))
            # print("Setting redis config:")
            redis = Redis(location=loc)
            this_dfs_status = plan.find(instrument="").success
            try:
                this_redis_status = redis.ping()
            except BaseException:
                this_redis_status = False
            this_status = dict(
                location=loc,
                status=this_dfs_status and this_redis_status,
                dfs=this_dfs_status,
                redis=this_redis_status,
            )
            status_list.append(this_status)
        status_table = Table(status_list)
        print(status_table)

        if status_table["status"].sum() == 0:
            raise ValueError("No DFS location is available")
        elif status_table["status"].sum() > 1:
            print("Multiple DFS locations are available, please specify one")
            raise ValueError("Multiple DFS locations are available")
        else:
            self.location = status_table["location"][status_table["status"]][0]
            self.config = CONFIG[self.location]
            print(f"Using DFS location: {location}")
            for k, v in CONFIG[loc]["dfs"].items():
                os.environ.setdefault(k, str(v))
            # print("Setting redis config:")
            self.redis = Redis(location=loc)


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)