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