66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
|
from datetime import datetime
|
||
|
|
from app.utils.calc_api_usage import calculate_api_usage
|
||
|
|
from flask import current_app
|
||
|
|
from typing import List, Dict, Any, Optional
|
||
|
|
from bson import ObjectId
|
||
|
|
|
||
|
|
class MongoBillingService:
|
||
|
|
def __init__(self):
|
||
|
|
self.mongo_client = current_app.mongo_client
|
||
|
|
|
||
|
|
def update_usage_total_cost(self,
|
||
|
|
product: str,
|
||
|
|
start_date: str,
|
||
|
|
end_date: str,
|
||
|
|
price: float,
|
||
|
|
billing_unit: int,
|
||
|
|
company_ids: Optional[List[str]] = None, ) -> int:
|
||
|
|
query = {
|
||
|
|
"product": product,
|
||
|
|
"createdAt": {
|
||
|
|
"$gte": datetime.strptime(f"{start_date} 00:00:00", "%Y-%m-%d %H:%M:%S"),
|
||
|
|
"$lte": datetime.strptime(f"{end_date} 23:59:59", "%Y-%m-%d %H:%M:%S")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if company_ids:
|
||
|
|
query["companyId"] = {"$in": company_ids}
|
||
|
|
|
||
|
|
collection = self.mongo_client["billing-api"]["api_usages"]
|
||
|
|
|
||
|
|
cursor = collection.find(query)
|
||
|
|
|
||
|
|
count = 0
|
||
|
|
|
||
|
|
for doc in cursor:
|
||
|
|
usage = float(doc.get("usage", 0))
|
||
|
|
new_total_cost = calculate_api_usage(float(price), int(billing_unit), float(usage))
|
||
|
|
|
||
|
|
collection.update_one(
|
||
|
|
{"_id": doc["_id"]},
|
||
|
|
{
|
||
|
|
"$set": {
|
||
|
|
"total_cost": new_total_cost,
|
||
|
|
"price": f"{price}",
|
||
|
|
"billingUnit": billing_unit,
|
||
|
|
"updatedAt": datetime.utcnow()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
)
|
||
|
|
count+=1
|
||
|
|
return count
|
||
|
|
|
||
|
|
def update_model_policy_price(self, id:str, update_data: Dict[str, Any]) -> int:
|
||
|
|
|
||
|
|
collection = self.mongo_client["billing-api"]["api_pricings"]
|
||
|
|
|
||
|
|
update_data["updatedAt"] = datetime.utcnow()
|
||
|
|
|
||
|
|
result = collection.update_one(
|
||
|
|
{"_id": ObjectId(id)},
|
||
|
|
{"$set": update_data}
|
||
|
|
)
|
||
|
|
|
||
|
|
return result.modified_count
|
||
|
|
|
||
|
|
|