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.3.0-alpha10
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 01cdf2e Dawsh 2y ago def record_user_stats(params: list, node_id: int, consumption_factor: int = 1):
15 9a6b0fa Dawsh 2y ago if not params:
16 9a6b0fa Dawsh 2y ago return
17 9a6b0fa Dawsh 2y ago
18 01cdf2e Dawsh 2y ago created_at = datetime.fromisoformat(datetime.utcnow().strftime("%Y-%m-%dT%H:00:00"))
19 9a6b0fa Dawsh 2y ago
20 9a6b0fa Dawsh 2y ago with GetDB() as db:
21 9a6b0fa Dawsh 2y ago # make user usage row if doesn't exist
22 01cdf2e Dawsh 2y ago select_stmt = select(NodeUserUsage.user_id).where(
23 01cdf2e Dawsh 2y ago and_(
24 01cdf2e Dawsh 2y ago NodeUserUsage.node_id == node_id, NodeUserUsage.created_at == created_at
25 01cdf2e Dawsh 2y ago )
26 01cdf2e Dawsh 2y ago )
27 9a6b0fa Dawsh 2y ago existings = [r[0] for r in db.execute(select_stmt).fetchall()]
28 9a6b0fa Dawsh 2y ago uids_to_insert = set()
29 9a6b0fa Dawsh 2y ago
30 9a6b0fa Dawsh 2y ago for p in params:
31 01cdf2e Dawsh 2y ago uid = p["uid"]
32 9a6b0fa Dawsh 2y ago if uid in existings:
33 9a6b0fa Dawsh 2y ago continue
34 9a6b0fa Dawsh 2y ago uids_to_insert.add(uid)
35 9a6b0fa Dawsh 2y ago
36 9a6b0fa Dawsh 2y ago if uids_to_insert:
37 9a6b0fa Dawsh 2y ago stmt = insert(NodeUserUsage).values(
38 01cdf2e Dawsh 2y ago user_id=bindparam("uid"),
39 9a6b0fa Dawsh 2y ago created_at=created_at,
40 9a6b0fa Dawsh 2y ago node_id=node_id,
41 01cdf2e Dawsh 2y ago used_traffic=0,
42 9a6b0fa Dawsh 2y ago )
43 01cdf2e Dawsh 2y ago db.execute(stmt, [{"uid": uid} for uid in uids_to_insert])
44 9a6b0fa Dawsh 2y ago
45 9a6b0fa Dawsh 2y ago # record
46 01cdf2e Dawsh 2y ago stmt = (
47 01cdf2e Dawsh 2y ago update(NodeUserUsage)
48 01cdf2e Dawsh 2y ago .values(
49 01cdf2e Dawsh 2y ago used_traffic=NodeUserUsage.used_traffic
50 01cdf2e Dawsh 2y ago + bindparam("value") * consumption_factor
51 01cdf2e Dawsh 2y ago )
52 01cdf2e Dawsh 2y ago .where(
53 01cdf2e Dawsh 2y ago and_(
54 01cdf2e Dawsh 2y ago NodeUserUsage.user_id == bindparam("uid"),
55 01cdf2e Dawsh 2y ago NodeUserUsage.node_id == node_id,
56 01cdf2e Dawsh 2y ago NodeUserUsage.created_at == created_at,
57 01cdf2e Dawsh 2y ago )
58 01cdf2e Dawsh 2y ago )
59 01cdf2e Dawsh 2y ago )
60 01cdf2e Dawsh 2y ago db.connection().execute(
61 01cdf2e Dawsh 2y ago stmt, params, execution_options={"synchronize_session": None}
62 01cdf2e Dawsh 2y ago )
63 9a6b0fa Dawsh 2y ago db.commit()
64 9a6b0fa Dawsh 2y ago
65 9a6b0fa Dawsh 2y ago
66 9a6b0fa Dawsh 2y ago def record_node_stats(params: dict, node_id: int):
67 9a6b0fa Dawsh 2y ago if not params:
68 9a6b0fa Dawsh 2y ago return
69 9a6b0fa Dawsh 2y ago
70 01cdf2e Dawsh 2y ago created_at = datetime.fromisoformat(datetime.utcnow().strftime("%Y-%m-%dT%H:00:00"))
71 9a6b0fa Dawsh 2y ago
72 9a6b0fa Dawsh 2y ago with GetDB() as db:
73 9a6b0fa Dawsh 2y ago
74 9a6b0fa Dawsh 2y ago # make node usage row if doesn't exist
75 01cdf2e Dawsh 2y ago select_stmt = select(NodeUsage.node_id).where(
76 01cdf2e Dawsh 2y ago and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at)
77 01cdf2e Dawsh 2y ago )
78 9a6b0fa Dawsh 2y ago notfound = db.execute(select_stmt).first() is None
79 9a6b0fa Dawsh 2y ago if notfound:
80 01cdf2e Dawsh 2y ago stmt = insert(NodeUsage).values(
81 01cdf2e Dawsh 2y ago created_at=created_at, node_id=node_id, uplink=0, downlink=0
82 01cdf2e Dawsh 2y ago )
83 9a6b0fa Dawsh 2y ago db.execute(stmt)
84 9a6b0fa Dawsh 2y ago db.commit()
85 9a6b0fa Dawsh 2y ago
86 9a6b0fa Dawsh 2y ago # record
87 01cdf2e Dawsh 2y ago stmt = (
88 01cdf2e Dawsh 2y ago update(NodeUsage)
89 01cdf2e Dawsh 2y ago .values(
90 01cdf2e Dawsh 2y ago uplink=NodeUsage.uplink + bindparam("up"),
91 01cdf2e Dawsh 2y ago downlink=NodeUsage.downlink + bindparam("down"),
92 01cdf2e Dawsh 2y ago )
93 01cdf2e Dawsh 2y ago .where(
94 01cdf2e Dawsh 2y ago and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at)
95 01cdf2e Dawsh 2y ago )
96 01cdf2e Dawsh 2y ago )
97 9a6b0fa Dawsh 2y ago
98 9a6b0fa Dawsh 2y ago db.execute(stmt, params)
99 9a6b0fa Dawsh 2y ago
100 9a6b0fa Dawsh 2y ago
101 9a6b0fa Dawsh 2y ago async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[dict]]:
102 9a6b0fa Dawsh 2y ago try:
103 9a6b0fa Dawsh 2y ago params = list()
104 9a6b0fa Dawsh 2y ago for stat in await node.fetch_users_stats():
105 9a6b0fa Dawsh 2y ago if stat.usage:
106 9a6b0fa Dawsh 2y ago params.append({"uid": stat.uid, "value": stat.usage})
107 9a6b0fa Dawsh 2y ago return node_id, params
108 f6a32ea Dawsh 2y ago except:
109 9a6b0fa Dawsh 2y ago return node_id, []
110 9a6b0fa Dawsh 2y ago
111 9a6b0fa Dawsh 2y ago
112 9a6b0fa Dawsh 2y ago async def record_user_usages():
113 9a6b0fa Dawsh 2y ago # usage_coefficient = {None: 1} # default usage coefficient for the main api instance
114 9a6b0fa Dawsh 2y ago
115 9a6b0fa Dawsh 2y ago results = await asyncio.gather(
116 01cdf2e Dawsh 2y ago *[get_users_stats(node_id, node) for node_id, node in marznode.nodes.items()]
117 01cdf2e Dawsh 2y ago )
118 9a6b0fa Dawsh 2y ago api_params = {node_id: params for node_id, params in list(results)}
119 9a6b0fa Dawsh 2y ago
120 9a6b0fa Dawsh 2y ago users_usage = defaultdict(int)
121 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
122 01cdf2e Dawsh 2y ago coefficient = (
123 01cdf2e Dawsh 2y ago node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1
124 01cdf2e Dawsh 2y ago )
125 9a6b0fa Dawsh 2y ago for param in params:
126 01cdf2e Dawsh 2y ago users_usage[param["uid"]] += (
127 01cdf2e Dawsh 2y ago param["value"] * coefficient
128 01cdf2e Dawsh 2y ago ) # apply the usage coefficient
129 9a6b0fa Dawsh 2y ago
130 01cdf2e Dawsh 2y ago users_usage = list(
131 01cdf2e Dawsh 2y ago {"id": uid, "value": value} for uid, value in users_usage.items()
132 01cdf2e Dawsh 2y ago )
133 9a6b0fa Dawsh 2y ago if not users_usage:
134 9a6b0fa Dawsh 2y ago return
135 9a6b0fa Dawsh 2y ago
136 9a6b0fa Dawsh 2y ago # record users usage
137 9a6b0fa Dawsh 2y ago with GetDB() as db:
138 01cdf2e Dawsh 2y ago stmt = update(User).values(
139 01cdf2e Dawsh 2y ago used_traffic=User.used_traffic + bindparam("value"),
140 01cdf2e Dawsh 2y ago online_at=datetime.utcnow(),
141 9a6b0fa Dawsh 2y ago )
142 9a6b0fa Dawsh 2y ago
143 9a6b0fa Dawsh 2y ago db.execute(stmt, users_usage, execution_options={"synchronize_session": None})
144 9a6b0fa Dawsh 2y ago db.commit()
145 9a6b0fa Dawsh 2y ago
146 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
147 01cdf2e Dawsh 2y ago record_user_stats(
148 01cdf2e Dawsh 2y ago params,
149 01cdf2e Dawsh 2y ago node_id,
150 01cdf2e Dawsh 2y ago node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1,
151 01cdf2e Dawsh 2y ago )