Begin implementing RTS

db.py

@@ -1,18 +1,22 @@
+import asyncio
 from math import ceil
+from os import getenv
+from time import sleep
 
+import aiomysql
 import aioredis
 import manticoresearch
+import msgpack
 import orjson
 from manticoresearch.rest import ApiException
 from numpy import array_split
 from redis import StrictRedis
-import msgpack
-import asyncio
 
 import util
 from schemas import mc_s
-from os import getenv
-from time import sleep
+
+mysql_pool = None
+
 
 configuration = manticoresearch.Configuration(host="http://127.0.0.1:9308")
 api_client = manticoresearch.ApiClient(configuration)
@@ -24,24 +28,25 @@ log = util.get_logger("db")
 # r = StrictRedis(unix_socket_path="/var/run/redis/redis.sock", db=0)
 r = StrictRedis(
     host="127.0.0.1",  # Replace with your Redis server's IP address
-    port=1289,         # Replace with your Redis server's port
-    db=0               # Database number
+    port=1289,  # Replace with your Redis server's port
+    db=0,  # Database number
 )
 # AIORedis
 # ar = aioredis.from_url("unix:///var/run/redis/redis.sock", db=0)
-ar = aioredis.from_url(
-    "redis://127.0.0.1:1289",
-    db=0
-)
+ar = aioredis.from_url("redis://127.0.0.1:1289", db=0)
 
 # /var/run/neptune-redis.sock
 # db = 10
 pr = aioredis.from_url("unix://var/run/neptune-redis.sock", db=10)
-#pr = aioredis.from_url("redis://redis_neptune:6379", db=10, password=getenv("REDIS_PASSWORD"))
+# fr = aioredis.from_url("unix://var/run/fisk-redis.sock", db=10)
+fr = aioredis.from_url("unix://var/run/redis.sock", db=10)
+# pr = aioredis.from_url("redis://redis_neptune:6379", db=10, password=getenv("REDIS_PASSWORD"))
 
 KEYNAME = "queue"
 MESSAGE_KEY = "messages"
 
+OHLC_MESSAGE_KEY = "ohlc"
+
 
 TYPES_MAIN = [
     "msg",
@@ -60,50 +65,68 @@ 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)
+async def init_mysql_pool():
+    """
+    Initialize the MySQL connection pool.
+    """
+    global mysql_pool
+    mysql_pool = await aiomysql.create_pool(
+        host="127.0.0.1", port=9306, db="Manticore", minsize=1, maxsize=10
+    )
 
-#     body = [{"insert": {"index": index, "doc": msg}}]
-#     body_post = ""
-#     for item in body:
-#         body_post += orjson.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)
+async def rts_store_message(index, data):
+    """
+    Store a RTS message into MySQL using an existing connection pool.
+    Prioritizes instant PubSub delivery, with minimal data storage overhead.
+    :param index: str
+    :param data: dict
+    """
+    # Publish to Redis PubSub
+    packed_index = msgpack.packb({"index": index, "data": data}, use_bin_type=True)
+
+    try:
+        await fr.publish(OHLC_MESSAGE_KEY, packed_index)
+    except aioredis.exceptions.ConnectionError as e:
+        raise e
+        await asyncio.sleep(0.1)
+
+    # Insert data into MySQL
+    try:
+        async with mysql_pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                # Insert data into the table
+                query = f"""
+                    INSERT INTO {index} (s, o, c, h, l, v, a, i, t, t2, ts)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                """
+                # Bind the values directly
+                await cur.execute(
+                    query,
+                    (
+                        data["s"],  # symbol
+                        data["o"],  # open
+                        data["c"],  # close
+                        data["h"],  # high
+                        data["l"],  # low
+                        data["v"],  # volume_base
+                        data["a"],  # volume_quote
+                        data["i"],  # interval
+                        data["t"],  # start_time
+                        data["t2"],  # end_time
+                        data["ts"],  # event_time
+                    ),
+                )
+                await conn.commit()
+                log.debug(f"Stored data for {data['s']} in MySQL.")
+    except aiomysql.Error as e:
+        log.error(f"MySQL error: {e}")
 
 
 async def store_batch(data):
     """
     Store a message into Manticore
-    :param msg: dict
+    :param data: list
     """
     if not data:
         return
@@ -161,12 +184,11 @@ async def store_batch(data):
 
         body_post = ""
         for item in total:
-            #print("ITEM", item)
+            # print("ITEM", item)
             body_post += orjson.dumps(item).decode("utf-8")
             body_post += "\n"
 
-        #print("BODY POST INDEX", index, body_post)
-
+        # print("BODY POST INDEX", index, body_post)
 
         try:
             # Bulk index operations
@@ -186,6 +208,7 @@ def create_index(api_client):
     util_instance = manticoresearch.UtilsApi(api_client)
     schemas = {
         "main": mc_s.schema_main,
+        "rule_storage": mc_s.schema_rule_storage,
         "meta": mc_s.schema_meta,
         "internal": mc_s.schema_int,
     }
@@ -216,14 +239,3 @@ async def queue_message_bulk(data):
         # TODO: msgpack
         message = orjson.dumps(msg)
         await ar.lpush(KEYNAME, message)
-
-
-created = False
-while not created:
-    try:
-        create_index(api_client)
-        created = True
-    except Exception as e:
-        print(f"Error creating index: {e}")
-        sleep(1) # Block the thread, just wait for the DB
-update_schema()