gitscope
/
← marzneshin / app/tasks/record_usages.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/tasks/record_usages.py · v0.1.2
Source
1 9a6b0fa Dawsh 2y ago import asyncio
2 9a6b0fa Dawsh 2y ago from collections import defaultdict
3 9a6b0fa Dawsh 2y ago from datetime import datetime
4 9a6b0fa Dawsh 2y ago
5 9a6b0fa Dawsh 2y ago from grpclib import GRPCError
6 9a6b0fa Dawsh 2y ago from sqlalchemy import and_, select, insert, update, bindparam
7 9a6b0fa Dawsh 2y ago
8 f7af2b3 Dawsh 2y ago from app import marznode
9 9a6b0fa Dawsh 2y ago from app.db import GetDB
10 9a6b0fa Dawsh 2y ago from app.db.models import NodeUsage, NodeUserUsage, User
11 9a6b0fa Dawsh 2y ago from app.marznode import MarzNodeBase
12 9a6b0fa Dawsh 2y ago
13 9a6b0fa Dawsh 2y ago
14 9a6b0fa Dawsh 2y ago def record_user_stats(params: list, node_id: int,
15 9a6b0fa Dawsh 2y ago consumption_factor: int = 1):
16 9a6b0fa Dawsh 2y ago if not params:
17 9a6b0fa Dawsh 2y ago return
18 9a6b0fa Dawsh 2y ago
19 9a6b0fa Dawsh 2y ago created_at = datetime.fromisoformat(datetime.utcnow().strftime('%Y-%m-%dT%H:00:00'))
20 9a6b0fa Dawsh 2y ago
21 9a6b0fa Dawsh 2y ago with GetDB() as db:
22 9a6b0fa Dawsh 2y ago # make user usage row if doesn't exist
23 9a6b0fa Dawsh 2y ago select_stmt = select(NodeUserUsage.user_id) \
24 9a6b0fa Dawsh 2y ago .where(and_(NodeUserUsage.node_id == node_id, NodeUserUsage.created_at == created_at))
25 9a6b0fa Dawsh 2y ago existings = [r[0] for r in db.execute(select_stmt).fetchall()]
26 9a6b0fa Dawsh 2y ago uids_to_insert = set()
27 9a6b0fa Dawsh 2y ago
28 9a6b0fa Dawsh 2y ago for p in params:
29 9a6b0fa Dawsh 2y ago uid = p['uid']
30 9a6b0fa Dawsh 2y ago if uid in existings:
31 9a6b0fa Dawsh 2y ago continue
32 9a6b0fa Dawsh 2y ago uids_to_insert.add(uid)
33 9a6b0fa Dawsh 2y ago
34 9a6b0fa Dawsh 2y ago if uids_to_insert:
35 9a6b0fa Dawsh 2y ago stmt = insert(NodeUserUsage).values(
36 9a6b0fa Dawsh 2y ago user_id=bindparam('uid'),
37 9a6b0fa Dawsh 2y ago created_at=created_at,
38 9a6b0fa Dawsh 2y ago node_id=node_id,
39 9a6b0fa Dawsh 2y ago used_traffic=0
40 9a6b0fa Dawsh 2y ago )
41 9a6b0fa Dawsh 2y ago db.execute(stmt, [{'uid': uid} for uid in uids_to_insert])
42 9a6b0fa Dawsh 2y ago
43 9a6b0fa Dawsh 2y ago # record
44 9a6b0fa Dawsh 2y ago stmt = update(NodeUserUsage) \
45 9a6b0fa Dawsh 2y ago .values(used_traffic=NodeUserUsage.used_traffic + bindparam('value') * consumption_factor) \
46 9a6b0fa Dawsh 2y ago .where(and_(NodeUserUsage.user_id == bindparam('uid'),
47 9a6b0fa Dawsh 2y ago NodeUserUsage.node_id == node_id,
48 9a6b0fa Dawsh 2y ago NodeUserUsage.created_at == created_at))
49 9a6b0fa Dawsh 2y ago db.connection().execute(stmt, params, execution_options={"synchronize_session": None})
50 9a6b0fa Dawsh 2y ago db.commit()
51 9a6b0fa Dawsh 2y ago
52 9a6b0fa Dawsh 2y ago
53 9a6b0fa Dawsh 2y ago def record_node_stats(params: dict, node_id: int):
54 9a6b0fa Dawsh 2y ago if not params:
55 9a6b0fa Dawsh 2y ago return
56 9a6b0fa Dawsh 2y ago
57 9a6b0fa Dawsh 2y ago created_at = datetime.fromisoformat(datetime.utcnow().strftime('%Y-%m-%dT%H:00:00'))
58 9a6b0fa Dawsh 2y ago
59 9a6b0fa Dawsh 2y ago with GetDB() as db:
60 9a6b0fa Dawsh 2y ago
61 9a6b0fa Dawsh 2y ago # make node usage row if doesn't exist
62 9a6b0fa Dawsh 2y ago select_stmt = select(NodeUsage.node_id). \
63 9a6b0fa Dawsh 2y ago where(and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at))
64 9a6b0fa Dawsh 2y ago notfound = db.execute(select_stmt).first() is None
65 9a6b0fa Dawsh 2y ago if notfound:
66 9a6b0fa Dawsh 2y ago stmt = insert(NodeUsage).values(created_at=created_at, node_id=node_id, uplink=0, downlink=0)
67 9a6b0fa Dawsh 2y ago db.execute(stmt)
68 9a6b0fa Dawsh 2y ago db.commit()
69 9a6b0fa Dawsh 2y ago
70 9a6b0fa Dawsh 2y ago # record
71 9a6b0fa Dawsh 2y ago stmt = update(NodeUsage). \
72 9a6b0fa Dawsh 2y ago values(uplink=NodeUsage.uplink + bindparam('up'), downlink=NodeUsage.downlink + bindparam('down')). \
73 9a6b0fa Dawsh 2y ago where(and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at))
74 9a6b0fa Dawsh 2y ago
75 9a6b0fa Dawsh 2y ago db.execute(stmt, params)
76 9a6b0fa Dawsh 2y ago
77 9a6b0fa Dawsh 2y ago
78 9a6b0fa Dawsh 2y ago async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[dict]]:
79 9a6b0fa Dawsh 2y ago try:
80 9a6b0fa Dawsh 2y ago params = list()
81 9a6b0fa Dawsh 2y ago for stat in await node.fetch_users_stats():
82 9a6b0fa Dawsh 2y ago if stat.usage:
83 9a6b0fa Dawsh 2y ago params.append({"uid": stat.uid, "value": stat.usage})
84 9a6b0fa Dawsh 2y ago return node_id, params
85 f6a32ea Dawsh 2y ago except:
86 9a6b0fa Dawsh 2y ago return node_id, []
87 9a6b0fa Dawsh 2y ago
88 9a6b0fa Dawsh 2y ago
89 9a6b0fa Dawsh 2y ago async def record_user_usages():
90 9a6b0fa Dawsh 2y ago # usage_coefficient = {None: 1} # default usage coefficient for the main api instance
91 9a6b0fa Dawsh 2y ago
92 9a6b0fa Dawsh 2y ago results = await asyncio.gather(
93 9a6b0fa Dawsh 2y ago *[get_users_stats(node_id, node) for node_id, node in marznode.nodes.items()]
94 9a6b0fa Dawsh 2y ago )
95 9a6b0fa Dawsh 2y ago api_params = {node_id: params for node_id, params in list(results)}
96 9a6b0fa Dawsh 2y ago
97 9a6b0fa Dawsh 2y ago users_usage = defaultdict(int)
98 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
99 9a6b0fa Dawsh 2y ago coefficient = node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1
100 9a6b0fa Dawsh 2y ago for param in params:
101 9a6b0fa Dawsh 2y ago users_usage[param['uid']] += param['value'] * coefficient # apply the usage coefficient
102 9a6b0fa Dawsh 2y ago
103 9a6b0fa Dawsh 2y ago users_usage = list({"id": uid, "value": value} for uid, value in users_usage.items())
104 9a6b0fa Dawsh 2y ago if not users_usage:
105 9a6b0fa Dawsh 2y ago return
106 9a6b0fa Dawsh 2y ago
107 9a6b0fa Dawsh 2y ago # record users usage
108 9a6b0fa Dawsh 2y ago with GetDB() as db:
109 9a6b0fa Dawsh 2y ago stmt = update(User)\
110 9a6b0fa Dawsh 2y ago .values(
111 9a6b0fa Dawsh 2y ago used_traffic=User.used_traffic + bindparam('value'),
112 9a6b0fa Dawsh 2y ago online_at=datetime.utcnow()
113 9a6b0fa Dawsh 2y ago )
114 9a6b0fa Dawsh 2y ago
115 9a6b0fa Dawsh 2y ago db.execute(stmt, users_usage, execution_options={"synchronize_session": None})
116 9a6b0fa Dawsh 2y ago db.commit()
117 9a6b0fa Dawsh 2y ago
118 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
119 9a6b0fa Dawsh 2y ago record_user_stats(params, node_id,
120 9a6b0fa Dawsh 2y ago node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1)
121 9a6b0fa Dawsh 2y ago