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