Skip to content

Commit

Permalink
usage seperate by cost per second
Browse files Browse the repository at this point in the history
  • Loading branch information
@lyang2821
lyang2821 committed Nov 25, 2024
1 parent 7e74bd8 commit c1f767d
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 7 deletions.
2 changes: 1 addition & 1 deletion lapdev-conductor/src/server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -545,7 +545,7 @@ impl Conductor {
let usage = self
.enterprise
.usage
.get_monthly_cost(org.id, None, Utc::now().into(), None)
.get_monthly_cost(org.id, None, None, Utc::now().into(), None)
.await
.unwrap_or(0);
if usage as i64 >= org.usage_limit {
Expand Down
3 changes: 2 additions & 1 deletion lapdev-enterprise/src/enterprise.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ impl Enterprise {
if organization.usage_limit > 0 {
let usage = self
.usage
.get_monthly_cost(organization.id, None, Utc::now().into(), None)
.get_monthly_cost(organization.id, None, None, Utc::now().into(), None)
.await?;
if usage as i64 >= organization.usage_limit {
return Err(ApiError::InvalidRequest(
Expand All @@ -97,6 +97,7 @@ impl Enterprise {
.get_monthly_cost(
org.id,
Some(user_id),
None,
Utc::now().into(),
None,
)
Expand Down
8 changes: 4 additions & 4 deletions lapdev-enterprise/src/quota.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,12 +218,12 @@ impl Quota {
QuotaKind::Project => 0,
QuotaKind::DailyCost => {
self.usage
.get_daily_cost(organization, Some(user), Utc::now().into(), None)
.get_daily_cost(organization, Some(user), None, Utc::now().into(), None)
.await?
}
QuotaKind::MonthlyCost => {
self.usage
.get_monthly_cost(organization, Some(user), Utc::now().into(), None)
.get_monthly_cost(organization, Some(user), None, Utc::now().into(), None)
.await?
}
};
Expand Down Expand Up @@ -264,14 +264,14 @@ impl Quota {
}
QuotaKind::DailyCost => {
self.usage
.get_daily_cost(organization, None, Utc::now().into(), None)
.get_daily_cost(organization, None, None, Utc::now().into(), None)
.await?
/ 60
/ 60
}
QuotaKind::MonthlyCost => {
self.usage
.get_monthly_cost(organization, None, Utc::now().into(), None)
.get_monthly_cost(organization, None, None, Utc::now().into(), None)
.await?
/ 60
/ 60
Expand Down
24 changes: 23 additions & 1 deletion lapdev-enterprise/src/usage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,9 @@ impl Usage {
})
.collect();

let total_cost = self.get_cost(organization, user, start, end, now).await?;
let total_cost = self
.get_cost(organization, user, None, start, end, now)
.await?;
Ok(UsageResult {
total_items: items_and_pages.number_of_items,
num_pages: items_and_pages.number_of_pages,
Expand All @@ -184,6 +186,7 @@ impl Usage {
&self,
organization: Uuid,
user: Option<Uuid>,
cost_per_second: Option<i32>,
start: DateTime<FixedOffset>,
end: DateTime<FixedOffset>,
now: DateTime<FixedOffset>,
Expand All @@ -196,6 +199,9 @@ impl Usage {
if let Some(user) = user {
query = query.filter(entities::usage::Column::UserId.eq(user));
}
if let Some(cost_per_second) = cost_per_second {
query = query.filter(entities::usage::Column::CostPerSecond.eq(cost_per_second));
}
let fixed_cost: Vec<Option<i64>> = query
.select_only()
.column_as(entities::usage::Column::TotalCost.sum(), "cost")
Expand Down Expand Up @@ -228,6 +234,9 @@ impl Usage {
if let Some(user) = user {
query = query.filter(entities::usage::Column::UserId.eq(user));
}
if let Some(cost_per_second) = cost_per_second {
query = query.filter(entities::usage::Column::CostPerSecond.eq(cost_per_second));
}

let usages = query.all(&self.db.conn).await?;

Expand Down Expand Up @@ -256,6 +265,7 @@ impl Usage {
&self,
organization: Uuid,
user: Option<Uuid>,
cost_per_second: Option<i32>,
date: DateTime<FixedOffset>,
now: Option<DateTime<FixedOffset>>,
) -> Result<usize> {
Expand All @@ -268,6 +278,7 @@ impl Usage {
self.get_cost(
organization,
user,
cost_per_second,
day_start,
day_end,
now.unwrap_or_else(|| Utc::now().into()),
Expand All @@ -279,6 +290,7 @@ impl Usage {
&self,
organization: Uuid,
user: Option<Uuid>,
cost_per_second: Option<i32>,
date: DateTime<FixedOffset>,
now: Option<DateTime<FixedOffset>>,
) -> Result<usize> {
Expand All @@ -305,6 +317,7 @@ impl Usage {
self.get_cost(
organization,
user,
cost_per_second,
month_start,
month_end,
now.unwrap_or_else(|| Utc::now().into()),
Expand Down Expand Up @@ -414,6 +427,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
None,
)
Expand All @@ -436,6 +450,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
None,
)
Expand All @@ -458,6 +473,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
Some(Utc.with_ymd_and_hms(2024, 1, 29, 15, 1, 0).unwrap().into()),
)
Expand All @@ -470,6 +486,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 30, 1, 0, 0).unwrap().into(),
Some(Utc.with_ymd_and_hms(2024, 1, 30, 1, 0, 0).unwrap().into()),
)
Expand All @@ -492,6 +509,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
Some(Utc.with_ymd_and_hms(2024, 1, 29, 15, 1, 0).unwrap().into()),
)
Expand All @@ -514,6 +532,7 @@ pub mod tests {
.get_daily_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
Some(Utc.with_ymd_and_hms(2024, 1, 29, 15, 1, 0).unwrap().into()),
)
Expand Down Expand Up @@ -561,6 +580,7 @@ pub mod tests {
.get_monthly_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
None,
)
Expand All @@ -583,6 +603,7 @@ pub mod tests {
.get_monthly_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2024, 1, 29, 1, 0, 0).unwrap().into(),
None,
)
Expand All @@ -595,6 +616,7 @@ pub mod tests {
.get_monthly_cost(
organization,
Some(user.id),
None,
Utc.with_ymd_and_hms(2023, 12, 29, 1, 0, 0).unwrap().into(),
None,
)
Expand Down

0 comments on commit c1f767d

Please sign in to comment.