Skip to content

Commit

Permalink
query roughly working
Browse files Browse the repository at this point in the history
  • Loading branch information
@calypsomatic
calypsomatic committed Nov 8, 2024
1 parent c5c736a commit 81f843b
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 34 deletions.
96 changes: 73 additions & 23 deletions core/src/main/scala/org/broadinstitute/dsde/rawls/spendreporting/SpendReportingService.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,65 @@ object SpendReportingService {
SpendReportingResults(aggregations.map(aggregateRows(allRows, _)).toList, summary)
}

def extractCrossBillingProjectSpendReportingResults(
allRows: List[FieldValueList],
start: DateTime,
end: DateTime,
names: Map[GoogleProjectId, WorkspaceName]
): List[SpendReportingResults] = {

val spendDetails = allRows.flatMap { row =>
val projectId = GoogleProjectId(row.get("project_id").getStringValue)
val workspaceName = names.getOrElse(
projectId,
throw RawlsExceptionWithErrorReport(
StatusCodes.InternalServerError,
s"unexpected project $projectId returned by BigQuery"
)
)
val totalCost = row.get("total_cost").getDoubleValue.toString
val computeCost = row.get("compute_cost").getDoubleValue.toString
val storageCost = row.get("storage_cost").getDoubleValue.toString
val currency = row.get("currency").getStringValue

// Each row gets a summary of compute, storage, and total
List(
SpendReportingForDateRange(
totalCost,
"0", // Ignoring credits for now; do we want to include them?
currency,
Option(start),
Option(end),
workspace = Some(workspaceName),
googleProjectId = Some(GoogleProject(projectId.value))
),
SpendReportingForDateRange(
computeCost,
"0",
currency,
Option(start),
Option(end),
workspace = Some(workspaceName),
googleProjectId = Some(GoogleProject(projectId.value)),
category = Some(TerraSpendCategories.Compute)
),
SpendReportingForDateRange(
storageCost,
"0",
currency,
Option(start),
Option(end),
workspace = Some(workspaceName),
googleProjectId = Some(GoogleProject(projectId.value)),
category = Some(TerraSpendCategories.Storage)
)
)
}

// TODO what format should the result be? Do we need a new model?
spendDetails.map(details => SpendReportingResults(List.empty, details))
}

}

class SpendReportingService(
Expand Down Expand Up @@ -270,6 +329,7 @@ class SpendReportingService(
| SELECT
| project.id AS project_id,
| project.name AS project_name,
| currency,
| CASE
| WHEN service.description IN ('Cloud Storage') THEN 'Storage'
| WHEN service.description IN ('Compute Engine', 'Google Kubernetes Engine') THEN 'Compute'
Expand All @@ -279,11 +339,12 @@ class SpendReportingService(
| FROM
| _BILLING_ACCOUNT_TABLE
| where
| project_id in _PROJECT_ID_LIST AND
| project.id in _PROJECT_ID_LIST AND
| _PARTITIONTIME BETWEEN @startDate AND @endDate
| GROUP BY
| project_id,
| project_name,
| currency,
| spend_category""".stripMargin.trim

val bpSubQuery = billingProjects
Expand All @@ -293,38 +354,27 @@ class SpendReportingService(
baseQuery
.replace("_PARTITIONTIME", timePartitionColumn)
.replace("_BILLING_ACCOUNT_TABLE", tableName)
.replace("_PROJECT_ID_LIST", "(" + bp._2.mkString(", ") + ")") + "\nUNION ALL"
.replace("_PROJECT_ID_LIST", "(" + bp._2.map(id => s""""${id.value}"""").mkString(", ") + ")")
}
.mkString("\n")

val allBPQuery = s"""WITH spend_categories AS (
|$bpSubQuery
| select
| project_id,
| project_name,
| spend_category,
| category_cost
| from
| `broad_materialized_view`
| where
| project_id in ('broad', 'list') AND
| _PARTITIONTIME BETWEEN @startDate AND @endDate
|)"""
.mkString("\nUNION ALL\n")

s"""
|$allBPQuery
s"""WITH spend_categories AS (
|$bpSubQuery
|)
|SELECT
| project_id,
| project_name,
| SUM(category_cost) AS total_cost,
| SUM(CASE WHEN spend_category = 'Storage' THEN category_cost ELSE 0 END) AS storage_cost,
| SUM(CASE WHEN spend_category = 'Compute' THEN category_cost ELSE 0 END) AS compute_cost,
| SUM(CASE WHEN spend_category = 'Other' THEN category_cost ELSE 0 END) AS other_cost
| SUM(CASE WHEN spend_category = 'Other' THEN category_cost ELSE 0 END) AS other_cost,
| currency
|FROM
| spend_categories
|GROUP BY
| project_id,
| project_name
| project_name,
| currency
|ORDER BY
| total_cost DESC
|limit 5
Expand Down Expand Up @@ -470,7 +520,7 @@ class SpendReportingService(
def getSpendForAllWorkspaces(
start: DateTime,
end: DateTime
): Future[SpendReportingResults] = {
): Future[List[SpendReportingResults]] = {
validateReportParameters(start, end)
for {
workspaces <- getOwnerWorkspaces()
Expand All @@ -492,7 +542,7 @@ class SpendReportingService(
StatusCodes.NotFound,
s"no spend data found between dates ${toISODateString(start)} and ${toISODateString(end)}"
) // TODO update this
case rows => extractSpendReportingResults(rows, start, end, projectNames, Set.empty)
case rows => extractCrossBillingProjectSpendReportingResults(rows, start, end, projectNames)
}
}

Expand Down
64 changes: 53 additions & 11 deletions ...c/test/scala/org/broadinstitute/dsde/rawls/spendreporting/SpendReportingServiceSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1383,17 +1383,6 @@ class SpendReportingServiceSpec extends AnyFlatSpecLike with Matchers with Mocki
| project_id,
| project_name,
| spend_category
| UNION ALL
| select
| project_id,
| project_name,
| spend_category,
| category_cost
| from
| `broad_materialized_view`
| where
| project_id in ('broad', 'list') AND
| _PARTITIONTIME BETWEEN @startDate AND @endDate
|)
|SELECT
| project_id,
Expand Down Expand Up @@ -1432,6 +1421,59 @@ class SpendReportingServiceSpec extends AnyFlatSpecLike with Matchers with Mocki

}

"extractSpendReportingResultsAcrossBillingProjects" should "should return correct summary data" in {
val storageCostWs1 = 100.582
val otherCostWs1 = 0.10111
val storageCostRoundedWs1: BigDecimal = BigDecimal(storageCostWs1).setScale(2, RoundingMode.HALF_EVEN)
val otherCostRoundedWs1: BigDecimal = BigDecimal(otherCostWs1).setScale(2, RoundingMode.HALF_EVEN)
val totalCostRoundedWs1: BigDecimal = BigDecimal(storageCostWs1 + otherCostWs1).setScale(2, RoundingMode.HALF_EVEN)

val storageCostWs2 = 20.145
val computeCostWs2 = 150.4033
val storageCostRoundedWs2: BigDecimal = BigDecimal(storageCostWs2).setScale(2, RoundingMode.HALF_EVEN)
val otherCostRoundedWs2: BigDecimal = BigDecimal(computeCostWs2).setScale(2, RoundingMode.HALF_EVEN)
val totalCostRoundedWs2: BigDecimal =
BigDecimal(storageCostWs2 + computeCostWs2).setScale(2, RoundingMode.HALF_EVEN)

val computeCostWs3 = 1111.222
val otherCostWs3 = 0.02
val storageCostRoundedWs3: BigDecimal = BigDecimal(computeCostWs3).setScale(2, RoundingMode.HALF_EVEN)
val otherCostRoundedWs3: BigDecimal = BigDecimal(otherCostWs3).setScale(2, RoundingMode.HALF_EVEN)
val totalCostRoundedWs3: BigDecimal = BigDecimal(computeCostWs3 + otherCostWs3).setScale(2, RoundingMode.HALF_EVEN)

val table: List[Map[String, String]] = List(
Map(
"storage_cost" -> s"$storageCostWs1",
"compute_cost" -> "0.0",
"other_cost" -> s"$otherCostWs1",
"googleProjectId" -> "workspace1ProjectId"
),
Map(
"storage_cost" -> s"$storageCostWs2",
"compute_cost" -> s"$computeCostWs2",
"other_cost" -> "0.0",
"googleProjectId" -> "workspace2ProjectId"
),
Map(
"storage_cost" -> "0.0",
"compute_cost" -> s"$computeCostWs3",
"other_cost" -> s"$otherCostWs3",
"googleProjectId" -> "workspace3ProjectId"
)
)

val tableResult: TableResult = createTableResult(table)

val reportingResults = SpendReportingService.extractCrossBillingProjectSpendReportingResults(
tableResult.getValues.asScala.toList,
DateTime.now().minusDays(1),
DateTime.now(),
Map()
)
reportingResults.spendSummary.cost shouldBe TestData.Workspace.totalCostRounded.toString
reportingResults.spendDetails shouldBe empty
}

"getSpendForAllWorkspaces" should "get the spend report from multiple billing projects" in {
val from = DateTime.now().minusMonths(2)
val to = from.plusMonths(1)
Expand Down

0 comments on commit 81f843b

Please sign in to comment.