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
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 a7dfd3a Dawsh 2y ago def record_user_usage_logs(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 a7dfd3a Dawsh 2y ago def record_node_stats(node_id: int, usage: int):
67 a7dfd3a Dawsh 2y ago if not usage:
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
85 9a6b0fa Dawsh 2y ago # record
86 01cdf2e Dawsh 2y ago stmt = (
87 01cdf2e Dawsh 2y ago update(NodeUsage)
88 01cdf2e Dawsh 2y ago .values(
89 a7dfd3a Dawsh 2y ago downlink=NodeUsage.downlink + usage,
90 01cdf2e Dawsh 2y ago )
91 01cdf2e Dawsh 2y ago .where(
92 01cdf2e Dawsh 2y ago and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at)
93 01cdf2e Dawsh 2y ago )
94 01cdf2e Dawsh 2y ago )
95 9a6b0fa Dawsh 2y ago
96 a7dfd3a Dawsh 2y ago db.execute(stmt)
97 a7dfd3a Dawsh 2y ago db.commit()
98 9a6b0fa Dawsh 2y ago
99 9a6b0fa Dawsh 2y ago
100 9a6b0fa Dawsh 2y ago async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[dict]]:
101 9a6b0fa Dawsh 2y ago try:
102 9a6b0fa Dawsh 2y ago params = list()
103 9a6b0fa Dawsh 2y ago for stat in await node.fetch_users_stats():
104 9a6b0fa Dawsh 2y ago if stat.usage:
105 9a6b0fa Dawsh 2y ago params.append({"uid": stat.uid, "value": stat.usage})
106 9a6b0fa Dawsh 2y ago return node_id, params
107 f6a32ea Dawsh 2y ago except:
108 9a6b0fa Dawsh 2y ago return node_id, []
109 9a6b0fa Dawsh 2y ago
110 9a6b0fa Dawsh 2y ago
111 9a6b0fa Dawsh 2y ago async def record_user_usages():
112 9a6b0fa Dawsh 2y ago # usage_coefficient = {None: 1} # default usage coefficient for the main api instance
113 9a6b0fa Dawsh 2y ago
114 9a6b0fa Dawsh 2y ago results = await asyncio.gather(
115 01cdf2e Dawsh 2y ago *[get_users_stats(node_id, node) for node_id, node in marznode.nodes.items()]
116 01cdf2e Dawsh 2y ago )
117 9a6b0fa Dawsh 2y ago api_params = {node_id: params for node_id, params in list(results)}
118 9a6b0fa Dawsh 2y ago
119 9a6b0fa Dawsh 2y ago users_usage = defaultdict(int)
120 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
121 01cdf2e Dawsh 2y ago coefficient = (
122 01cdf2e Dawsh 2y ago node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1
123 01cdf2e Dawsh 2y ago )
124 a7dfd3a Dawsh 2y ago node_usage = 0
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 a7dfd3a Dawsh 2y ago node_usage += param["value"]
130 a7dfd3a Dawsh 2y ago record_node_stats(node_id, node_usage)
131 9a6b0fa Dawsh 2y ago
132 01cdf2e Dawsh 2y ago users_usage = list(
133 01cdf2e Dawsh 2y ago {"id": uid, "value": value} for uid, value in users_usage.items()
134 01cdf2e Dawsh 2y ago )
135 9a6b0fa Dawsh 2y ago if not users_usage:
136 9a6b0fa Dawsh 2y ago return
137 9a6b0fa Dawsh 2y ago
138 9a6b0fa Dawsh 2y ago # record users usage
139 9a6b0fa Dawsh 2y ago with GetDB() as db:
140 01cdf2e Dawsh 2y ago stmt = update(User).values(
141 01cdf2e Dawsh 2y ago used_traffic=User.used_traffic + bindparam("value"),
142 5b5d739 Dawsh 2y ago lifetime_used_traffic=User.used_traffic + bindparam("value"),
143 01cdf2e Dawsh 2y ago online_at=datetime.utcnow(),
144 9a6b0fa Dawsh 2y ago )
145 9a6b0fa Dawsh 2y ago
146 9a6b0fa Dawsh 2y ago db.execute(stmt, users_usage, execution_options={"synchronize_session": None})
147 9a6b0fa Dawsh 2y ago db.commit()
148 9a6b0fa Dawsh 2y ago
149 9a6b0fa Dawsh 2y ago for node_id, params in api_params.items():
150 a7dfd3a Dawsh 2y ago record_user_usage_logs(
151 01cdf2e Dawsh 2y ago params,
152 01cdf2e Dawsh 2y ago node_id,
153 01cdf2e Dawsh 2y ago node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1,
154 01cdf2e Dawsh 2y ago )