from math import ceil import aioredis import manticoresearch import ujson from manticoresearch.rest import ApiException from numpy import array_split from redis import StrictRedis import util from schemas import mc_s configuration = manticoresearch.Configuration(host="http://monolith-db-1:9308") api_client = manticoresearch.ApiClient(configuration) api_instance = manticoresearch.IndexApi(api_client) log = util.get_logger("db") # Redis (legacy) r = StrictRedis(unix_socket_path="/var/run/redis/redis.sock", db=0) # AIORedis ar = aioredis.from_url("unix:///var/run/redis/redis.sock", db=0) TYPES_MAIN = [ "msg", "notice", "action", "part", "join", "kick", "quit", "nick", "mode", "topic", "update", ] TYPES_META = ["who"] TYPES_INT = ["conn", "highlight", "znc", "query", "self"] def store_message(msg): """ Store a message into Manticore :param msg: dict """ # Duplicated to avoid extra function call if msg["type"] in TYPES_MAIN: index = "main" schema = mc_s.schema_main elif msg["type"] in TYPES_META: index = "meta" schema = mc_s.schema_meta elif msg["type"] in TYPES_INT: index = "internal" schema = mc_s.schema_int # normalise fields for key, value in list(msg.items()): if value is None: del msg[key] if key in schema: if isinstance(value, int): if schema[key].startswith("string") or schema[key].startswith("text"): msg[key] = str(value) body = [{"insert": {"index": index, "doc": msg}}] body_post = "" for item in body: body_post += ujson.dumps(item) body_post += "\n" # print(body_post) try: # Bulk index operations api_response = api_instance.bulk(body_post) # , async_req=True # print(api_response) except ApiException as e: print("Exception when calling IndexApi->bulk: %s\n" % e) print("ATTEMPT", body_post) def store_message_bulk(data): """ Store a message into Manticore :param msg: dict """ if not data: return # 10000: maximum inserts we can submit to # Manticore as of Sept 2022 split_posts = array_split(data, ceil(len(data) / 10000)) for messages in split_posts: total = [] for msg in messages: # Duplicated to avoid extra function call (see above) if msg["type"] in TYPES_MAIN: index = "main" schema = mc_s.schema_main elif msg["type"] in TYPES_META: index = "meta" schema = mc_s.schema_meta elif msg["type"] in TYPES_INT: index = "internal" schema = mc_s.schema_int # normalise fields for key, value in list(msg.items()): if value is None: del msg[key] if key in schema: if isinstance(value, int): if schema[key].startswith("string") or schema[key].startswith( "text" ): msg[key] = str(value) body = {"insert": {"index": index, "doc": msg}} total.append(body) body_post = "" for item in total: body_post += ujson.dumps(item) body_post += "\n" # print(body_post) try: # Bulk index operations api_response = api_instance.bulk(body_post) # , async_req=True print(api_response) except ApiException as e: print("Exception when calling IndexApi->bulk: %s\n" % e) print("ATTEMPT", body_post) def update_schema(): pass def create_index(api_client): util_instance = manticoresearch.UtilsApi(api_client) schemas = { "main": mc_s.schema_main, "meta": mc_s.schema_meta, "internal": mc_s.schema_int, } for name, schema in schemas.items(): schema_types = ", ".join([f"{k} {v}" for k, v in schema.items()]) create_query = ( f"create table if not exists {name}({schema_types}) engine='columnar'" ) print("Schema types", create_query) util_instance.sql(create_query) create_index(api_client) update_schema()