From 60be1eb2bd16fc89827a50ad18377642aae176be Mon Sep 17 00:00:00 2001
From: Kyomotoi <0w0@imki.moe>
Date: Thu, 6 Apr 2023 16:07:46 +0800
Subject: =?UTF-8?q?=F0=9F=8E=A8=20=E4=BC=98=E5=8C=96=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=BA=93=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ATRI/database/__init__.py |  1 +
 ATRI/database/wrapper.py  | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 ATRI/database/wrapper.py

(limited to 'ATRI/database')

diff --git a/ATRI/database/__init__.py b/ATRI/database/__init__.py
index 5c12c6a..f43b83b 100644
--- a/ATRI/database/__init__.py
+++ b/ATRI/database/__init__.py
@@ -1,2 +1,3 @@
 from .db import init_database, close_database_connection
 from .models import *
+from .wrapper import DatabaseWrapper
diff --git a/ATRI/database/wrapper.py b/ATRI/database/wrapper.py
new file mode 100644
index 0000000..7838bd6
--- /dev/null
+++ b/ATRI/database/wrapper.py
@@ -0,0 +1,31 @@
+from typing import Type
+
+from tortoise.models import Model
+
+
+class DatabaseWrapper:
+    def __init__(self, model: Type[Model]):
+        self.model = model
+    
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    async def add_sub(self, *args, **kwargs):
+        await self.model.create(*args, **kwargs)
+    
+    async def update_sub(self, update_map: dict, **kwargs):
+        await self.model.filter(**kwargs).update(
+            **update_map
+        )
+
+    async def del_sub(self, query_map: dict):
+        await self.model.filter(**query_map).delete()
+    
+    async def get_sub_list(self, query_map: dict) -> list:
+        return await self.model.filter(**query_map)
+    
+    async def get_all_subs(self) -> list:
+        return await self.model.all()
-- 
cgit v1.2.3