8000 fix: 修复没有找到 database 代码问题 · TTB-Network/python-openbmclapi@27ab567 · GitHub
[go: up one dir, main page]
Skip to content

Commit 27ab567

Browse files
tianxiu2b2ttianxiu2b2t
committed
fix: 修复没有找到 database 代码问题
1 parent ba5d5ef commit 27ab567

File tree

7 files changed

+270
-2
lines changed

7 files changed

+270
-2
lines changed

‎.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,6 @@ logs
55
bmclapi
66
config
77
cache
8-
database
8+
database/*.db
99

1010
**/__pycache__

‎core/config.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@ def __init__(self):
1111

1212
def load(self):
1313
try:
14-
with open(ROOT_PATH / "config" / "config.yml", "r", encoding="utf-8") as f:
14+
file = ROOT_PATH / "config" / "config.yml"
15+
if not file.exists():
16+
return
17+
with open(file, "r", encoding="utf-8") as f:
1518
self._data = yaml.safe_load(f) or {}
1619
except:
1720
print("[Config] Failed to load config.yml")

‎core/database/__init__.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import importlib
2+
from typing import Any
3+
from .abc import DataBase
4+
5+
6+
loaders = {
7+
"local": "SqliteDB",
8+
"mongo": "MongoDB",
9+
"memory": "MemoryDataBase"
10+
}
11+
12+
db = None
13+
14+
async def init(
15+
config: dict[str, Any]
16+
):
17+
global db
18+
type = config.get('type', 'memory')
19+
module = importlib.import_module(f'.{type}', __package__)
20+
instance = getattr(module, loaders[type])
21+
database_name = config.get('database_name', 'pyoba')
22+
db = instance(
23+
database_name=database_name,
24+
**config
25+
)
26+
27+
28+
def get_db() -> DataBase:
29+
assert db is not None, 'Database not initialized'
30+
return db

‎core/database/abc.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
import abc
2+
from dataclasses import dataclass
3+
import datetime
4+
from pathlib import Path
5+
import time
6+
from typing import Any, Optional
7+
8+
ROOT = Path(__file__).parent.parent.parent / 'database'
9+
ROOT.mkdir(parents=True, exist_ok=True)
10+
11+
@dataclass
12+
class ClusterCounterInfo:
13+
time: datetime.datetime
14+
hits: int
15+
bytes: int
16+
17+
class DataBase(
18+
metaclass=abc.ABCMeta
19+
):
20+
"""Abstract class for database"""
21+
22+
def __init__(
23+
self,
24+
database_name: str
25+
):
26+
self.database_name = database_name
27+
28+
@abc.abstractmethod
29+
async def insert_cluster_info(
30+
self,
31+
cluster_id: str,
32+
type: str,
33+
event: str,
34+
data: Optional[Any] = None
35+
):
36+
"""Insert cluster info into database"""
37+
pass
38+
39+
@abc.abstractmethod
40+
async def upsert_cluster_counter(
41+
self,
42+
cluster_id: str,
43+
hits: int,
44+
bytes: int
45+
):
46+
"""Upsert cluster counter into database"""
47+
pass
48+
49+
@abc.abstractmethod
50+
async def get_cluster_counter(
51+
self,
52+
cluster_id: str,
53+
before_hour: int = 0
54+
) -> list[ClusterCounterInfo]:
55+
"""Get cluster counter from database"""
56+
pass
57+
58+
59+
def get_hour(self) -> int:
60+
# Asia/Shanghai
61+
# use timestamp
62+
return int(time.time() // 3600)
63+

‎core/database/local.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import datetime
2+
from typing import Any, Optional
3+
4+
from . import abc
5+
6+
try:
7+
import sqlite3 as sqlite
8+
except:
9+
try:
10+
import pysqlite3 as sqlite # type: ignore
11+
except:
12+
print("Error: sqlite3 or pysqlite3 not found")
13+
exit(1)
14+
15+
class SqliteDB(abc.DataBase):
16+
def __init__(
17+
self,
18+
database_name: str
19+
):
20+
super().__init__(database_name)
21+
22+
self.clusters = sqlite.connect(
23+
abc.ROOT / f'{self.database_name}_clusters.db'
24+
)
25+
26+
self.clusters.execute('create table if not exists logs (timestamp text, cluster_id text, type text, event text, data text)')
27+
# hits, bytes unsigned big integer
28+
self.clusters.execute('create table if not exists counters (hour text, cluster_id text, hits unsigned big integer, bytes unsigned big integer, primary key (hour, cluster_id))')
29+
self.clusters.commit()
30+
31+
async def insert_cluster_info(self, cluster_id: str, type: str, event: str, data: Optional[Any] = None):
32+
self.clusters.execute(
33+
'insert into logs values (?, ?, ?, ?, ?)',
34+
(datetime.datetime.now().isoformat(), cluster_id, type, event, data)
35+
)
36+
self.clusters.commit()
37+
38+
async def upsert_cluster_counter(self, cluster_id: str, hits: int, bytes: int):
39+
hour = self.get_hour()
40+
self.clusters.execute(
41+
'insert into counters values (?, ?, ?, ?) on conflict (hour, cluster_id) do update set hits = hits + ?, bytes = bytes + ?',
42+
(hour, cluster_id, hits, bytes, hits, bytes)
43+
)
44+
self.clusters.commit()
45+
46+
async def get_cluster_counter(self, cluster_id: str, before_hour: int = 0) -> list[abc.ClusterCounterInfo]:
47+
return [
48+
abc.ClusterCounterInfo(
49+
time=datetime.datetime.fromtimestamp(row[0] * 3600),
50+
hits=row[1],
51+
bytes=row[2]
52+
) for row in self.clusters.execute(
53+
'select hour, hits, bytes from counters where cluster_id = ? and hour >= ? and hour <= ?',
54+
(cluster_id, before_hour, self.get_hour())
55+
)
56+
]
57+
58+
def __del__(self):
59+
self.clusters.close()

‎core/database/memory.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
from collections import defaultdict
2+
from datetime import datetime
3+
from typing import Any
4+
from .abc import DataBase, ClusterCounterInfo
5+
6+
7+
class MemoryDataBase(DataBase):
8+
def __init__(
9+
self,
10+
database_name: str
11+
):
12+
super().__init__(database_name)
13+
14+
self._clusters_logs = []
15+
self._clusters_counters: defaultdict[int, defaultdict[str, defaultdict[str, int]]] = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
16+
17+
async def insert_cluster_info(self, cluster_id: str, type: str, event: str, data: Any | None = None):
18+
self._clusters_logs.append({
19+
'cluster_id': cluster_id,
20+
'type': type,
21+
'event': event,
22+
'data': data,
23+
})
24+
25+
async def upsert_cluster_counter(self, cluster_id: str, hits: int, bytes: int):
26+
hour = self.get_hour()
27+
self._clusters_counters[hour][cluster_id]['hits'] += hits
28+
self._clusters_counters[hour][cluster_id]['bytes'] += bytes
29+
30+
async def get_cluster_counter(self, cluster_id: str, before_hour: int = 0) -> list[ClusterCounterInfo]:
31+
res = []
32+
for h in range(before_hour, self.get_hour() + 1):
33+
if h not in self._clusters_counters or cluster_id not in self._clusters_counters[h]:
34+
continue
35+
res.append(ClusterCounterInfo(
36+
time=datetime.fromtimestamp(h * 3600),
37+
hits=self._clusters_counters[h][cluster_id]['hits'],
38+
bytes=self._clusters_counters[h][cluster_id]['bytes'],
39+
))
40+
return res

‎core/database/mongo.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
from datetime import datetime
2+
from typing import Any, Optional
3+
import motor.motor_asyncio as motor
4+
import urllib.parse as urlparse
5+
6+
from core.database.abc import ClusterCounterInfo
7+
8+
from .abc import DataBase
9+
10+
class MongoDB(DataBase):
11+
def __init__(
12+
self,
13+
database_name: str,
14+
host: str,
15+
port: int,
16+
**kwargs
17+
):
18+
super().__init__(
19+
database_name=database_name
20+
)
21+
22+
user = ""
23+
if "username" in kwargs and "password":
24+
user = f":".join((
25+
urlparse.quote(x) for x in (
26+
kwargs["username"],
27+
kwargs["password"]
28+
)
29+
)) + "@"
30+
31+
32+
self.connection = motor.AsyncIOMotorClient(f"mongodb://{user}{host}:{port}")
33+
34+
self.db = self.connection.get_database(f"{database_name}")
35+
36+
async def insert_cluster_info(self, cluster_id: str, type: str, event: str, data: Optional[Any] = None):
37+
await self.db["cluster_logs"].insert_one({
38+
"cluster_id": cluster_id,
39+
"type": type,
40+
"event": event,
41+
"data": data
42+
})
43+
44+
async def upsert_cluster_counter(self, cluster_id: str, hits: int, bytes: int):
45+
# add
46+
await self.db["cluster_counters"].update_one(
47+
{
48+
"cluster_id": cluster_id,
49+
"hour": self.get_hour()
50+
},
51+
{
52+
"$inc": {
53+
"hits": hits,
54+
"bytes": bytes
55+
}
56+
},
57+
upsert=True
58+
)
59+
60+
async def get_cluster_counter(self, cluster_id: str, before_hour: int = 0) -> list[ClusterCounterInfo]:
61+
res = []
62+
async for doc in self.db.get_collection("cluster_counters").find({
63+
"cluster_id": cluster_id,
64+
"hour": {"$gte": before_hour, "$lt": self.get_hour()}
65+
}):
66+
res.append(ClusterCounterInfo(
67+
time=datetime.fromtimestamp(doc["hour"] * 3600),
68+
hits=doc["hits"],
69+
bytes=doc["bytes"]
70+
))
71+
return res
72+
73+

0 commit comments

Comments
 (0)
0