feat(cg1): implement real BigQuery billing queries in billing_agent.py
This commit is contained in:
parent
df6676e6cb
commit
93b29079a4
|
|
@ -1,4 +1,6 @@
|
|||
import os
|
||||
import datetime
|
||||
from google.cloud import bigquery
|
||||
from google.cloud import billing
|
||||
|
||||
class BillingAgent:
|
||||
|
|
@ -6,41 +8,100 @@ class BillingAgent:
|
|||
self.project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
|
||||
if not self.project_id:
|
||||
raise ValueError("GOOGLE_CLOUD_PROJECT environment variable not set.")
|
||||
|
||||
self.billing_account_id = os.environ.get("BILLING_ACCOUNT_ID")
|
||||
if not self.billing_account_id:
|
||||
raise ValueError("BILLING_ACCOUNT_ID environment variable not set.")
|
||||
|
||||
self.bq_client = bigquery.Client(project=self.project_id)
|
||||
self.billing_client = billing.CloudBillingClient()
|
||||
self.billing_table = f"{self.project_id}.billing_export.gcp_billing_export_v1_*"
|
||||
|
||||
def get_summary(self):
|
||||
"""
|
||||
Retrieves the billing summary for the last 30 days.
|
||||
"""
|
||||
# This is a simplified example. The Cloud Billing API does not have a direct
|
||||
# "summary" endpoint. You would typically query for cost data over a time range.
|
||||
# This implementation will be a placeholder.
|
||||
return {
|
||||
"summary": "Billing summary data would be here.",
|
||||
"project_id": self.project_id
|
||||
}
|
||||
query = f"""
|
||||
SELECT
|
||||
service.description as service,
|
||||
SUM(cost) as total_cost
|
||||
FROM `{self.billing_table}`
|
||||
WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY)
|
||||
GROUP BY 1
|
||||
ORDER BY total_cost DESC
|
||||
"""
|
||||
query_job = self.bq_client.query(query)
|
||||
results = query_job.result()
|
||||
|
||||
summary = []
|
||||
for row in results:
|
||||
summary.append({"service": row.service, "total_cost": row.total_cost})
|
||||
|
||||
return {"summary": summary}
|
||||
|
||||
def get_forecast(self):
|
||||
"""
|
||||
Retrieves a billing forecast based on the last 7 days.
|
||||
"""
|
||||
# The Cloud Billing API does not provide a direct forecast API.
|
||||
# This would require custom logic to calculate.
|
||||
# This implementation will be a placeholder.
|
||||
# Get total cost for the last 7 days
|
||||
query_last_7_days = f"""
|
||||
SELECT SUM(cost) as total_cost
|
||||
FROM `{self.billing_table}`
|
||||
WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 7 DAY)
|
||||
"""
|
||||
query_job = self.bq_client.query(query_last_7_days)
|
||||
results = query_job.result()
|
||||
|
||||
total_cost_last_7_days = 0
|
||||
for row in results:
|
||||
total_cost_last_7_days = row.total_cost or 0
|
||||
|
||||
daily_average = total_cost_last_7_days / 7
|
||||
|
||||
# Calculate remaining days in the month
|
||||
today = datetime.date.today()
|
||||
last_day_of_month = datetime.date(today.year, today.month, 1) + datetime.timedelta(days=32)
|
||||
last_day_of_month = last_day_of_month.replace(day=1) - datetime.timedelta(days=1)
|
||||
remaining_days = (last_day_of_month - today).days
|
||||
|
||||
forecasted_cost = daily_average * remaining_days
|
||||
|
||||
# Get current month-to-date cost
|
||||
query_mtd = f"""
|
||||
SELECT SUM(cost) as total_cost
|
||||
FROM `{self.billing_table}`
|
||||
WHERE EXTRACT(MONTH FROM _PARTITIONTIME) = EXTRACT(MONTH FROM CURRENT_DATE())
|
||||
AND EXTRACT(YEAR FROM _PARTITIONTIME) = EXTRACT(YEAR FROM CURRENT_DATE())
|
||||
"""
|
||||
query_job_mtd = self.bq_client.query(query_mtd)
|
||||
results_mtd = query_job_mtd.result()
|
||||
mtd_cost = 0
|
||||
for row in results_mtd:
|
||||
mtd_cost = row.total_cost or 0
|
||||
|
||||
total_forecast = mtd_cost + forecasted_cost
|
||||
|
||||
return {
|
||||
"forecast": "Billing forecast data would be here.",
|
||||
"project_id": self.project_id
|
||||
"daily_average_last_7_days": daily_average,
|
||||
"month_to_date_cost": mtd_cost,
|
||||
"forecasted_remaining_cost": forecasted_cost,
|
||||
"total_monthly_forecast": total_forecast,
|
||||
}
|
||||
|
||||
def get_credits(self):
|
||||
"""
|
||||
Retrieves active billing credits.
|
||||
"""
|
||||
# This is a simplified example. You would need to query for billing account
|
||||
# information and then check for credits.
|
||||
billing_account_name = f"billingAccounts/{self.billing_account_id}"
|
||||
billing_info = self.billing_client.get_billing_account(name=billing_account_name)
|
||||
|
||||
# This is a simplified representation. The actual credit information might
|
||||
# be structured differently. This is a placeholder for the structure.
|
||||
return {
|
||||
"credits": "Billing credits information would be here.",
|
||||
"project_id": self.project_id
|
||||
"billing_account_name": billing_info.name,
|
||||
"display_name": billing_info.display_name,
|
||||
"open": billing_info.open,
|
||||
"credits": "Detailed credit information would require further API calls and parsing."
|
||||
}
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user