rlt_salary/app/database/methods.py

62 lines
1.8 KiB
Python
Raw Normal View History

2024-02-27 20:21:36 +01:00
from typing import List, Tuple
from motor.motor_asyncio import AsyncIOMotorCollection
import datetime as dt
2024-02-27 20:21:36 +01:00
async def aggregate_salaries(collection: AsyncIOMotorCollection, dt_from: str, dt_upto: str, group_type: str) -> Tuple[List[int], List[str]]: # type: ignore
2024-02-27 20:42:40 +01:00
group_type_format = "%Y-%m-%dT00:00:00" if group_type == "day" else "%Y-%m-01T00:00:00" if group_type == "month" else "%Y-%m-%dT%H:00:00"
2024-02-27 20:21:36 +01:00
pipeline = [
{
"$match": {
"dt": {
"$gte": iso(dt_from),
"$lte": iso(dt_upto)
}
}
},
{
"$densify": {
"field": "dt",
"range": {
"step": 1,
"unit": group_type,
"bounds": [
iso(dt_from), iso(dt_upto) + dt.timedelta(days=int(group_type=="day"), hours=int(group_type=="hour")) * (iso(dt_upto).time() == dt.time(hour=0, minute=0, second=0))
]
}
}
},
{
"$group": {
"_id": {
"$dateToString":
{
2024-02-27 20:42:40 +01:00
"format": group_type_format,
"date": "$dt"
}
},
"sum": {
"$sum": "$value"
},
}
2024-02-27 20:42:40 +01:00
},
{
"$sort": {
"_id": 1
}
}
2024-02-27 20:21:36 +01:00
]
dataset = []
labels = []
async for document in collection.aggregate(pipeline):
dataset.append(document["sum"])
labels.append(document["_id"] )
return {"dataset": dataset, "labels": labels}
def iso(date: str) -> dt.datetime:
return dt.datetime.fromisoformat(date)