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.6.3
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 9a6b0fa Dawsh 2y ago # make user usage row if 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 01cdf2e Dawsh 2y ago used_traffic=NodeUserUsage.used_traffic
55 01cdf2e Dawsh 2y ago + bindparam("value") * consumption_factor
56 01cdf2e Dawsh 2y ago )
57 01cdf2e Dawsh 2y ago .where(
58 01cdf2e Dawsh 2y ago and_(
59 01cdf2e Dawsh 2y ago NodeUserUsage.user_id == bindparam("uid"),
60 01cdf2e Dawsh 2y ago NodeUserUsage.node_id == node_id,
61 01cdf2e Dawsh 2y ago NodeUserUsage.created_at == created_at,
62 01cdf2e Dawsh 2y ago )
63 01cdf2e Dawsh 2y ago )
64 01cdf2e Dawsh 2y ago )
65 01cdf2e Dawsh 2y ago db.connection().execute(
66 01cdf2e Dawsh 2y ago stmt, params, execution_options={"synchronize_session": None}
67 01cdf2e Dawsh 2y ago )
68 9a6b0fa Dawsh 2y ago db.commit()
69 9a6b0fa Dawsh 2y ago
70 9a6b0fa Dawsh 2y ago
71 a7dfd3a Dawsh 2y ago def record_node_stats(node_id: int, usage: int):
72 a7dfd3a Dawsh 2y ago if not usage:
73 9a6b0fa Dawsh 2y ago return
74 9a6b0fa Dawsh 2y ago
75 f1ab16c Dawsh 2y ago created_at = datetime.fromisoformat(
76 f1ab16c Dawsh 2y ago datetime.utcnow().strftime("%Y-%m-%dT%H:00:00")
77 f1ab16c Dawsh 2y ago )
78 9a6b0fa Dawsh 2y ago
79 9a6b0fa Dawsh 2y ago with GetDB() as db:
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 8440155 mahdim10 1y ago await data_usage_percent_reached(db, users_usage)
160 8440155 mahdim10 1y ago
161 01cdf2e Dawsh 2y ago stmt = update(User).values(
162 01cdf2e Dawsh 2y ago used_traffic=User.used_traffic + bindparam("value"),
163 d1e3c6a Dawsh 2y ago lifetime_used_traffic=User.lifetime_used_traffic
164 d1e3c6a Dawsh 2y ago + bindparam("value"),
165 01cdf2e Dawsh 2y ago online_at=datetime.utcnow(),
166 9a6b0fa Dawsh 2y ago )
167 9a6b0fa Dawsh 2y ago
168 f1ab16c Dawsh 2y ago db.execute(
169 f1ab16c Dawsh 2y ago stmt, users_usage, execution_options={"synchronize_session": None}
170 f1ab16c Dawsh 2y ago )
171 9a6b0fa Dawsh 2y ago db.commit()
172 9a6b0fa Dawsh 2y ago
173 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
174 a7dfd3a Dawsh 2y ago record_user_usage_logs(
175 01cdf2e Dawsh 2y ago params,
176 01cdf2e Dawsh 2y ago node_id,
177 f1ab16c Dawsh 2y ago (
178 f1ab16c Dawsh 2y ago node.usage_coefficient
179 f1ab16c Dawsh 2y ago if (node := marznode.nodes.get(node_id))
180 f1ab16c Dawsh 2y ago else 1
181 f1ab16c Dawsh 2y ago ),
182 01cdf2e Dawsh 2y ago )