Bläddra i källkod

First pass at as_data, mostly untested for now

Kirk Trombley 4 år sedan
förälder
incheckning
2213ec3104
4 ändrade filer med 68 tillägg och 1 borttagningar
  1. 18 0
      data_test.py
  2. 1 1
      lib/rollbot/command/__init__.py
  3. 47 0
      lib/rollbot/command/decorators.py
  4. 2 0
      lib/rollbot/command/injection.py

+ 18 - 0
data_test.py

@@ -0,0 +1,18 @@
+from rollbot import as_data
+
+
+@as_data
+class Test:
+    foo: int
+    bar: str
+    baz: list[float]
+
+
+t = Test(0, "x", [0.5, 0.7])
+print(t)
+print(t.__field_queries)
+print(t.__create_query)
+print(t.__select_query)
+print(t.__save_query)
+print(t.__to_blob())
+print(Test.__from_blob('{"foo": 7, "bar": "z", "baz": []}'))

+ 1 - 1
lib/rollbot/command/__init__.py

@@ -1,3 +1,3 @@
-from .decorators import as_command, on_startup, on_shutdown, get_command_config
+from .decorators import as_command, on_startup, on_shutdown, get_command_config, as_data
 from .failure import RollbotFailure
 from .injection import Args, ArgList, ArgListSplitOn, ArgParse, Database, Lazy

+ 47 - 0
lib/rollbot/command/decorators.py

@@ -3,6 +3,8 @@ from typing import Union
 from functools import wraps
 import inspect
 import asyncio
+import dataclasses
+import json
 
 from ..types import (
     Message,
@@ -103,3 +105,48 @@ def get_command_config() -> CommandConfiguration:
         startup=decorated_startup,
         shutdown=decorated_shutdown,
     )
+
+
+def as_data(cls):
+    table_name = "".join(("_" + c.lower()) if "A" <= c <= "Z" else c for c in cls.__name__).strip("_")
+
+    columns = [
+        "key TEXT NOT NULL PRIMARY KEY",
+        'body TEXT DEFAULT ""',
+    ]
+    queries = {}
+
+    for k, v in cls.__annotations__.items():
+        if v == int:
+            t = "INT"
+        elif v == float:
+            t = "REAL"
+        elif v == str:
+            t = "TEXT"
+        else:
+            continue
+        columns.append(f"__{k} {t} GENERATED ALWAYS AS (json_extract(body, '$.{k}')) VIRTUAL")
+        queries[k] = f"SELECT body FROM {table_name} WHERE __{k} = ?"
+
+    create_query = f"CREATE TABLE IF NOT EXISTS {table_name} ({', '.join(columns)})"
+
+    @on_startup
+    async def create_table(context: Context):
+        async with context.database() as db:
+            await db.execute(create_query)
+            await db.commit()
+
+    new_class = type(
+        cls.__name__,
+        (dataclasses.dataclass(cls),),
+        dict(
+            __field_queries=queries,
+            __create_query=create_query,
+            __select_query=f"SELECT body FROM {table_name} WHERE key=:key",
+            __save_query=f"INSERT INTO {table_name} VALUES (:key, :body) ON CONFLICT(key) DO UPDATE SET body=:body",
+            __to_blob=lambda self: json.dumps(dataclasses.asdict(self)),
+            __from_blob=staticmethod(lambda blob: new_class(**json.loads(blob)))
+        ),
+    )
+
+    return new_class

+ 2 - 0
lib/rollbot/command/injection.py

@@ -69,3 +69,5 @@ class Lazy(Injector[Callable[[], Coroutine[None, None, Dep]]]):
 Args = ArgsInjector()
 ArgList = ArgListSplitOn()
 Database = DatabaseInjector()
+
+# TODO data query handler