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