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)