mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat: add user AI budget override endpoints (#25439)
Implements https://linear.app/codercom/issue/AIGOV-285 Follow the structure established in https://github.com/coder/coder/pull/25203 ## Summary Adds the `user_ai_budget_overrides` table and CRUD API at `/api/v2/users/{user}/ai/budget`. An override sets a custom per-user spend cap that supersedes group-budget resolution, attributing spend to a specific group. ## Schema ```sql CREATE TABLE user_ai_budget_overrides ( user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, group_id UUID NOT NULL REFERENCES groups(id) ON DELETE CASCADE, spend_limit_micros BIGINT NOT NULL CHECK (spend_limit_micros >= 0), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); ``` ## Membership lifecycle The membership invariant — a user must be a member of the attributed group, including when that group is "Everyone" — would naturally be expressed as a composite FK on `(user_id, group_id) → group_members_expanded(user_id, group_id)`. PostgreSQL doesn't allow foreign keys to reference views, so enforcement is split across two mechanisms: - **Write-time check.** A CHECK constraint on the table (`user_ai_budget_overrides_must_be_group_member`) calls a `STABLE` function `is_group_member(user_id, group_id)` that queries `group_members_expanded`. The view surfaces both regular group memberships and the implicit "Everyone" group memberships from `organization_members`. Any INSERT or UPDATE that violates the predicate is rejected with a Postgres `check_violation`, which the handler maps to a 400. `is_group_member` is defined as a general predicate, reusable by any future table that needs the same check. - **Cascade on removal.** Two `BEFORE DELETE` triggers handle membership loss: - `trigger_delete_user_ai_budget_overrides_on_group_member_delete` on `group_members` — covers regular group removals (admin action, OIDC sync). - `trigger_delete_user_ai_budget_overrides_on_org_member_delete` on `organization_members` — covers the "Everyone" group, whose membership lives in `organization_members`. The single-column FKs on `users(id)` and `groups(id)` remain to cascade on user or group deletion (those paths don't pass through `group_members`). ## Authorization The dbauthz layer gates each operation against the `User` and (for writes) `Group` resources: | Operation | User resource | Group resource | |-----------|----------------|----------------| | `GET` | `ActionRead` | — | | `PUT` | `ActionUpdate` | `ActionUpdate` | | `DELETE` | `ActionUpdate` | `ActionUpdate` | For `DELETE`, the dbauthz layer fetches the existing override first to learn the attributed `group_id`, then runs both checks. ### Role matrix | Role | GET | PUT | DELETE | |--------------|-----|-----|--------| | Owner | ✅ | ✅ | ✅ | | UserAdmin | ✅ | ✅ | ✅ | | OrgAdmin | ✅ | ❌ | ❌ | | OrgUserAdmin | ✅ | ❌ | ❌ | Internal discussion: https://codercom.slack.com/archives/C096PFVBZKN/p1779392747885359 ## Audit logs Audit logs will be addressed in a follow-up PR.
This commit is contained in:
Yevhenii Shcherbina
committed by
GitHub
GitHub
parent
9448624d2d
commit
1a91d31793
Generated
+145
@@ -9171,6 +9171,110 @@ const docTemplate = `{
|
||||
]
|
||||
}
|
||||
},
|
||||
"/api/v2/users/{user}/ai/budget": {
|
||||
"get": {
|
||||
"produces": [
|
||||
"application/json"
|
||||
],
|
||||
"tags": [
|
||||
"Enterprise"
|
||||
],
|
||||
"summary": "Get user AI budget override",
|
||||
"operationId": "get-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UserAIBudgetOverride"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"consumes": [
|
||||
"application/json"
|
||||
],
|
||||
"produces": [
|
||||
"application/json"
|
||||
],
|
||||
"tags": [
|
||||
"Enterprise"
|
||||
],
|
||||
"summary": "Upsert user AI budget override",
|
||||
"operationId": "upsert-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
},
|
||||
{
|
||||
"description": "Upsert user AI budget override request",
|
||||
"name": "request",
|
||||
"in": "body",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UpsertUserAIBudgetOverrideRequest"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UserAIBudgetOverride"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"delete": {
|
||||
"tags": [
|
||||
"Enterprise"
|
||||
],
|
||||
"summary": "Delete user AI budget override",
|
||||
"operationId": "delete-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"204": {
|
||||
"description": "No Content"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"/api/v2/users/{user}/appearance": {
|
||||
"get": {
|
||||
"produces": [
|
||||
@@ -24663,6 +24767,23 @@ const docTemplate = `{
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UpsertUserAIBudgetOverrideRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"group_id"
|
||||
],
|
||||
"properties": {
|
||||
"group_id": {
|
||||
"description": "GroupID is the group the user's spend is attributed to. The user must\nbe a member of this group.",
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"spend_limit_micros": {
|
||||
"type": "integer",
|
||||
"minimum": 0
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UpsertWorkspaceAgentPortShareRequest": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
@@ -24817,6 +24938,30 @@ const docTemplate = `{
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UserAIBudgetOverride": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"created_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"group_id": {
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"spend_limit_micros": {
|
||||
"type": "integer"
|
||||
},
|
||||
"updated_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"user_id": {
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UserActivity": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
Generated
+131
@@ -8132,6 +8132,98 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"/api/v2/users/{user}/ai/budget": {
|
||||
"get": {
|
||||
"produces": ["application/json"],
|
||||
"tags": ["Enterprise"],
|
||||
"summary": "Get user AI budget override",
|
||||
"operationId": "get-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UserAIBudgetOverride"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"consumes": ["application/json"],
|
||||
"produces": ["application/json"],
|
||||
"tags": ["Enterprise"],
|
||||
"summary": "Upsert user AI budget override",
|
||||
"operationId": "upsert-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
},
|
||||
{
|
||||
"description": "Upsert user AI budget override request",
|
||||
"name": "request",
|
||||
"in": "body",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UpsertUserAIBudgetOverrideRequest"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"schema": {
|
||||
"$ref": "#/definitions/codersdk.UserAIBudgetOverride"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"delete": {
|
||||
"tags": ["Enterprise"],
|
||||
"summary": "Delete user AI budget override",
|
||||
"operationId": "delete-user-ai-budget-override",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"description": "User ID, username, or me",
|
||||
"name": "user",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"204": {
|
||||
"description": "No Content"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"/api/v2/users/{user}/appearance": {
|
||||
"get": {
|
||||
"produces": ["application/json"],
|
||||
@@ -22696,6 +22788,21 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UpsertUserAIBudgetOverrideRequest": {
|
||||
"type": "object",
|
||||
"required": ["group_id"],
|
||||
"properties": {
|
||||
"group_id": {
|
||||
"description": "GroupID is the group the user's spend is attributed to. The user must\nbe a member of this group.",
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"spend_limit_micros": {
|
||||
"type": "integer",
|
||||
"minimum": 0
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UpsertWorkspaceAgentPortShareRequest": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
@@ -22829,6 +22936,30 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UserAIBudgetOverride": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"created_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"group_id": {
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"spend_limit_micros": {
|
||||
"type": "integer"
|
||||
},
|
||||
"updated_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"user_id": {
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UserActivity": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
Generated
+1
@@ -44,6 +44,7 @@ const (
|
||||
CheckTelemetryLockEventTypeConstraint CheckConstraint = "telemetry_lock_event_type_constraint" // telemetry_locks
|
||||
CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters
|
||||
CheckUsageEventTypeCheck CheckConstraint = "usage_event_type_check" // usage_events
|
||||
CheckUserAiBudgetOverridesSpendLimitMicrosCheck CheckConstraint = "user_ai_budget_overrides_spend_limit_micros_check" // user_ai_budget_overrides
|
||||
CheckUserAiProviderKeysAPIKeyCheck CheckConstraint = "user_ai_provider_keys_api_key_check" // user_ai_provider_keys
|
||||
CheckUserSkillsContentSize CheckConstraint = "user_skills_content_size" // user_skills
|
||||
CheckUserSkillsDescriptionSize CheckConstraint = "user_skills_description_size" // user_skills
|
||||
|
||||
@@ -1509,6 +1509,16 @@ func GroupAIBudget(b database.GroupAiBudget) codersdk.GroupAIBudget {
|
||||
}
|
||||
}
|
||||
|
||||
func UserAIBudgetOverride(o database.UserAiBudgetOverride) codersdk.UserAIBudgetOverride {
|
||||
return codersdk.UserAIBudgetOverride{
|
||||
UserID: o.UserID,
|
||||
GroupID: o.GroupID,
|
||||
SpendLimitMicros: o.SpendLimitMicros,
|
||||
CreatedAt: o.CreatedAt,
|
||||
UpdatedAt: o.UpdatedAt,
|
||||
}
|
||||
}
|
||||
|
||||
func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset {
|
||||
var presets []codersdk.InvalidatedPreset
|
||||
for _, p := range invalidatedPresets {
|
||||
|
||||
@@ -2323,6 +2323,32 @@ func (q *querier) DeleteTask(ctx context.Context, arg database.DeleteTaskParams)
|
||||
return q.db.DeleteTask(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
// Removing a user's AI budget override affects both the user (clearing
|
||||
// their per-user spend cap) and the group it was attributed to.
|
||||
u, err := q.db.GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, u); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
// Fetch the existing override to learn which group it attributes spend to,
|
||||
// so we can authorize the caller against that group as well.
|
||||
userOverride, err := q.db.GetUserAIBudgetOverride(ctx, userID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
g, err := q.db.GetGroupByID(ctx, userOverride.GroupID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, g); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.DeleteUserAIBudgetOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
@@ -4536,6 +4562,13 @@ func (q *querier) GetUnexpiredLicenses(ctx context.Context) ([]database.License,
|
||||
return q.db.GetUnexpiredLicenses(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
if _, err := q.GetUserByID(ctx, userID); err != nil { // AuthZ check
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.GetUserAIBudgetOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
@@ -8342,6 +8375,26 @@ func (q *querier) UpsertTemplateUsageStats(ctx context.Context) error {
|
||||
return q.db.UpsertTemplateUsageStats(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) {
|
||||
// Setting a user's AI budget override affects both the user (their
|
||||
// per-user spend cap) and the group (spend attribution).
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, u); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
g, err := q.db.GetGroupByID(ctx, arg.GroupID)
|
||||
if err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, g); err != nil {
|
||||
return database.UserAiBudgetOverride{}, err
|
||||
}
|
||||
return q.db.UpsertUserAIBudgetOverride(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) {
|
||||
u, err := q.db.GetUserByID(ctx, arg.UserID)
|
||||
if err != nil {
|
||||
|
||||
@@ -6475,6 +6475,36 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
check.Args(g.ID).Asserts(g, policy.ActionUpdate).Returns(b)
|
||||
}))
|
||||
|
||||
s.Run("GetUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
user := testutil.Fake(s.T(), faker, database.User{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID})
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
check.Args(user.ID).Asserts(user, policy.ActionRead).Returns(override)
|
||||
}))
|
||||
|
||||
s.Run("UpsertUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
user := testutil.Fake(s.T(), faker, database.User{})
|
||||
group := testutil.Fake(s.T(), faker, database.Group{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID, GroupID: group.ID})
|
||||
arg := database.UpsertUserAIBudgetOverrideParams{UserID: user.ID, GroupID: group.ID, SpendLimitMicros: override.SpendLimitMicros}
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetGroupByID(gomock.Any(), group.ID).Return(group, nil).AnyTimes()
|
||||
dbm.EXPECT().UpsertUserAIBudgetOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
|
||||
s.Run("DeleteUserAIBudgetOverride", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
user := testutil.Fake(s.T(), faker, database.User{})
|
||||
group := testutil.Fake(s.T(), faker, database.Group{})
|
||||
override := testutil.Fake(s.T(), faker, database.UserAiBudgetOverride{UserID: user.ID, GroupID: group.ID})
|
||||
dbm.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
|
||||
dbm.EXPECT().GetUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
dbm.EXPECT().GetGroupByID(gomock.Any(), group.ID).Return(group, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteUserAIBudgetOverride(gomock.Any(), user.ID).Return(override, nil).AnyTimes()
|
||||
check.Args(user.ID).Asserts(user, policy.ActionUpdate, group, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
|
||||
s.Run("GetAIProviderByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
provider := testutil.Fake(s.T(), faker, database.AIProvider{})
|
||||
dbm.EXPECT().GetAIProviderByID(gomock.Any(), provider.ID).Return(provider, nil).AnyTimes()
|
||||
|
||||
+24
@@ -793,6 +793,14 @@ func (m queryMetricsStore) DeleteTask(ctx context.Context, arg database.DeleteTa
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteUserAIBudgetOverride(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("DeleteUserAIBudgetOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteUserAIBudgetOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteUserAIProviderKey(ctx, arg)
|
||||
@@ -2905,6 +2913,14 @@ func (m queryMetricsStore) GetUnexpiredLicenses(ctx context.Context) ([]database
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserAIBudgetOverride(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("GetUserAIBudgetOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserAIBudgetOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserAIProviderKeyByProviderID(ctx, arg)
|
||||
@@ -6049,6 +6065,14 @@ func (m queryMetricsStore) UpsertTemplateUsageStats(ctx context.Context) error {
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertUserAIBudgetOverride(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertUserAIBudgetOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertUserAIBudgetOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertUserAIProviderKey(ctx, arg)
|
||||
|
||||
Generated
+45
@@ -1349,6 +1349,21 @@ func (mr *MockStoreMockRecorder) DeleteTask(ctx, arg any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteTask", reflect.TypeOf((*MockStore)(nil).DeleteTask), ctx, arg)
|
||||
}
|
||||
|
||||
// DeleteUserAIBudgetOverride mocks base method.
|
||||
func (m *MockStore) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteUserAIBudgetOverride", ctx, userID)
|
||||
ret0, _ := ret[0].(database.UserAiBudgetOverride)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DeleteUserAIBudgetOverride indicates an expected call of DeleteUserAIBudgetOverride.
|
||||
func (mr *MockStoreMockRecorder) DeleteUserAIBudgetOverride(ctx, userID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).DeleteUserAIBudgetOverride), ctx, userID)
|
||||
}
|
||||
|
||||
// DeleteUserAIProviderKey mocks base method.
|
||||
func (m *MockStore) DeleteUserAIProviderKey(ctx context.Context, arg database.DeleteUserAIProviderKeyParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -5445,6 +5460,21 @@ func (mr *MockStoreMockRecorder) GetUnexpiredLicenses(ctx any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUnexpiredLicenses", reflect.TypeOf((*MockStore)(nil).GetUnexpiredLicenses), ctx)
|
||||
}
|
||||
|
||||
// GetUserAIBudgetOverride mocks base method.
|
||||
func (m *MockStore) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAiBudgetOverride, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetUserAIBudgetOverride", ctx, userID)
|
||||
ret0, _ := ret[0].(database.UserAiBudgetOverride)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetUserAIBudgetOverride indicates an expected call of GetUserAIBudgetOverride.
|
||||
func (mr *MockStoreMockRecorder) GetUserAIBudgetOverride(ctx, userID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).GetUserAIBudgetOverride), ctx, userID)
|
||||
}
|
||||
|
||||
// GetUserAIProviderKeyByProviderID mocks base method.
|
||||
func (m *MockStore) GetUserAIProviderKeyByProviderID(ctx context.Context, arg database.GetUserAIProviderKeyByProviderIDParams) (database.UserAiProviderKey, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -11344,6 +11374,21 @@ func (mr *MockStoreMockRecorder) UpsertTemplateUsageStats(ctx any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertTemplateUsageStats", reflect.TypeOf((*MockStore)(nil).UpsertTemplateUsageStats), ctx)
|
||||
}
|
||||
|
||||
// UpsertUserAIBudgetOverride mocks base method.
|
||||
func (m *MockStore) UpsertUserAIBudgetOverride(ctx context.Context, arg database.UpsertUserAIBudgetOverrideParams) (database.UserAiBudgetOverride, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertUserAIBudgetOverride", ctx, arg)
|
||||
ret0, _ := ret[0].(database.UserAiBudgetOverride)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpsertUserAIBudgetOverride indicates an expected call of UpsertUserAIBudgetOverride.
|
||||
func (mr *MockStoreMockRecorder) UpsertUserAIBudgetOverride(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertUserAIBudgetOverride", reflect.TypeOf((*MockStore)(nil).UpsertUserAIBudgetOverride), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertUserAIProviderKey mocks base method.
|
||||
func (m *MockStore) UpsertUserAIProviderKey(ctx context.Context, arg database.UpsertUserAIProviderKeyParams) (database.UserAiProviderKey, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+62
@@ -841,6 +841,42 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
DELETE FROM user_ai_budget_overrides
|
||||
WHERE user_id = OLD.user_id AND group_id = OLD.group_id;
|
||||
RETURN OLD;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
DELETE FROM user_ai_budget_overrides
|
||||
WHERE user_id = OLD.user_id AND group_id = OLD.organization_id;
|
||||
RETURN OLD;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION enforce_user_ai_budget_override_membership() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM group_members_expanded
|
||||
WHERE user_id = NEW.user_id AND group_id = NEW.group_id
|
||||
) THEN
|
||||
RAISE EXCEPTION 'user % is not a member of group %', NEW.user_id, NEW.group_id
|
||||
USING ERRCODE = 'check_violation',
|
||||
CONSTRAINT = 'user_ai_budget_overrides_must_be_group_member';
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION enforce_user_secrets_per_user_limits() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
@@ -3137,6 +3173,17 @@ COMMENT ON TABLE usage_events_daily IS 'usage_events_daily is a daily rollup of
|
||||
|
||||
COMMENT ON COLUMN usage_events_daily.day IS 'The date of the summed usage events, always in UTC.';
|
||||
|
||||
CREATE TABLE user_ai_budget_overrides (
|
||||
user_id uuid NOT NULL,
|
||||
group_id uuid NOT NULL,
|
||||
spend_limit_micros bigint NOT NULL,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT user_ai_budget_overrides_spend_limit_micros_check CHECK ((spend_limit_micros >= 0))
|
||||
);
|
||||
|
||||
COMMENT ON TABLE user_ai_budget_overrides IS 'Per-user AI spend override that supersedes group budget resolution.';
|
||||
|
||||
CREATE TABLE user_ai_provider_keys (
|
||||
id uuid DEFAULT gen_random_uuid() NOT NULL,
|
||||
user_id uuid NOT NULL,
|
||||
@@ -3986,6 +4033,9 @@ ALTER TABLE ONLY usage_events_daily
|
||||
ALTER TABLE ONLY usage_events
|
||||
ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id);
|
||||
|
||||
ALTER TABLE ONLY user_ai_budget_overrides
|
||||
ADD CONSTRAINT user_ai_budget_overrides_pkey PRIMARY KEY (user_id);
|
||||
|
||||
ALTER TABLE ONLY user_ai_provider_keys
|
||||
ADD CONSTRAINT user_ai_provider_keys_pkey PRIMARY KEY (id);
|
||||
|
||||
@@ -4467,6 +4517,12 @@ CREATE TRIGGER trigger_delete_group_members_on_org_member_delete BEFORE DELETE O
|
||||
|
||||
CREATE TRIGGER trigger_delete_oauth2_provider_app_token AFTER DELETE ON oauth2_provider_app_tokens FOR EACH ROW EXECUTE FUNCTION delete_deleted_oauth2_provider_app_token_api_key();
|
||||
|
||||
CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_group_member_delete BEFORE DELETE ON group_members FOR EACH ROW EXECUTE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete();
|
||||
|
||||
CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_org_member_delete BEFORE DELETE ON organization_members FOR EACH ROW EXECUTE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete();
|
||||
|
||||
CREATE TRIGGER trigger_enforce_user_ai_budget_override_membership BEFORE INSERT OR UPDATE ON user_ai_budget_overrides FOR EACH ROW EXECUTE FUNCTION enforce_user_ai_budget_override_membership();
|
||||
|
||||
CREATE TRIGGER trigger_insert_apikeys BEFORE INSERT ON api_keys FOR EACH ROW EXECUTE FUNCTION insert_apikey_fail_if_user_deleted();
|
||||
|
||||
CREATE TRIGGER trigger_insert_organization_system_roles AFTER INSERT ON organizations FOR EACH ROW EXECUTE FUNCTION insert_organization_system_roles();
|
||||
@@ -4792,6 +4848,12 @@ ALTER TABLE ONLY templates
|
||||
ALTER TABLE ONLY templates
|
||||
ADD CONSTRAINT templates_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY user_ai_budget_overrides
|
||||
ADD CONSTRAINT user_ai_budget_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY user_ai_budget_overrides
|
||||
ADD CONSTRAINT user_ai_budget_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
|
||||
ALTER TABLE ONLY user_ai_provider_keys
|
||||
ADD CONSTRAINT user_ai_provider_keys_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE;
|
||||
|
||||
|
||||
+2
@@ -105,6 +105,8 @@ const (
|
||||
ForeignKeyTemplateVersionsTemplateID ForeignKeyConstraint = "template_versions_template_id_fkey" // ALTER TABLE ONLY template_versions ADD CONSTRAINT template_versions_template_id_fkey FOREIGN KEY (template_id) REFERENCES templates(id) ON DELETE CASCADE;
|
||||
ForeignKeyTemplatesCreatedBy ForeignKeyConstraint = "templates_created_by_fkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id) ON DELETE RESTRICT;
|
||||
ForeignKeyTemplatesOrganizationID ForeignKeyConstraint = "templates_organization_id_fkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
|
||||
ForeignKeyUserAiBudgetOverridesGroupID ForeignKeyConstraint = "user_ai_budget_overrides_group_id_fkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE;
|
||||
ForeignKeyUserAiBudgetOverridesUserID ForeignKeyConstraint = "user_ai_budget_overrides_user_id_fkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
ForeignKeyUserAiProviderKeysAiProviderID ForeignKeyConstraint = "user_ai_provider_keys_ai_provider_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE;
|
||||
ForeignKeyUserAiProviderKeysAPIKeyKeyID ForeignKeyConstraint = "user_ai_provider_keys_api_key_key_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest);
|
||||
ForeignKeyUserAiProviderKeysUserID ForeignKeyConstraint = "user_ai_provider_keys_user_id_fkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
DROP TRIGGER IF EXISTS trigger_delete_user_ai_budget_overrides_on_org_member_delete ON organization_members;
|
||||
DROP FUNCTION IF EXISTS delete_user_ai_budget_overrides_on_org_member_delete;
|
||||
DROP TRIGGER IF EXISTS trigger_delete_user_ai_budget_overrides_on_group_member_delete ON group_members;
|
||||
DROP FUNCTION IF EXISTS delete_user_ai_budget_overrides_on_group_member_delete;
|
||||
DROP TRIGGER IF EXISTS trigger_enforce_user_ai_budget_override_membership ON user_ai_budget_overrides;
|
||||
DROP FUNCTION IF EXISTS enforce_user_ai_budget_override_membership;
|
||||
DROP TABLE IF EXISTS user_ai_budget_overrides CASCADE;
|
||||
@@ -0,0 +1,76 @@
|
||||
CREATE TABLE user_ai_budget_overrides (
|
||||
user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE,
|
||||
group_id UUID NOT NULL REFERENCES groups(id) ON DELETE CASCADE,
|
||||
-- Spend limit applied to the user, in micro-units (1 unit = 1,000,000).
|
||||
spend_limit_micros BIGINT NOT NULL CHECK (spend_limit_micros >= 0),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
-- The membership invariant (user must be a member of the attributed
|
||||
-- group, including when that group is "Everyone") would naturally be
|
||||
-- a composite FK to group_members_expanded, but PostgreSQL does not
|
||||
-- allow FKs to views. It's enforced instead by a write-time trigger
|
||||
-- on this table and removal-time triggers on the underlying
|
||||
-- membership tables.
|
||||
);
|
||||
|
||||
COMMENT ON TABLE user_ai_budget_overrides IS 'Per-user AI spend override that supersedes group budget resolution.';
|
||||
|
||||
-- Write-time membership check. Reads from group_members_expanded so
|
||||
-- the "Everyone" group (whose membership lives in organization_members)
|
||||
-- is correctly handled. Raises check_violation with a constraint name
|
||||
-- so callers can match it via database.IsCheckViolation in Go.
|
||||
CREATE FUNCTION enforce_user_ai_budget_override_membership() RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM group_members_expanded
|
||||
WHERE user_id = NEW.user_id AND group_id = NEW.group_id
|
||||
) THEN
|
||||
RAISE EXCEPTION 'user % is not a member of group %', NEW.user_id, NEW.group_id
|
||||
USING ERRCODE = 'check_violation',
|
||||
CONSTRAINT = 'user_ai_budget_overrides_must_be_group_member';
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TRIGGER trigger_enforce_user_ai_budget_override_membership
|
||||
BEFORE INSERT OR UPDATE ON user_ai_budget_overrides
|
||||
FOR EACH ROW
|
||||
EXECUTE PROCEDURE enforce_user_ai_budget_override_membership();
|
||||
|
||||
-- When a user is removed from a regular group (any group except
|
||||
-- "Everyone"), delete any override attributed to that group.
|
||||
CREATE FUNCTION delete_user_ai_budget_overrides_on_group_member_delete() RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
DELETE FROM user_ai_budget_overrides
|
||||
WHERE user_id = OLD.user_id AND group_id = OLD.group_id;
|
||||
RETURN OLD;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_group_member_delete
|
||||
BEFORE DELETE ON group_members
|
||||
FOR EACH ROW
|
||||
EXECUTE PROCEDURE delete_user_ai_budget_overrides_on_group_member_delete();
|
||||
|
||||
-- When a user is removed from an organization, delete any override
|
||||
-- attributed to that organization's "Everyone" group (which has
|
||||
-- id == organization_id).
|
||||
CREATE FUNCTION delete_user_ai_budget_overrides_on_org_member_delete() RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
DELETE FROM user_ai_budget_overrides
|
||||
WHERE user_id = OLD.user_id AND group_id = OLD.organization_id;
|
||||
RETURN OLD;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TRIGGER trigger_delete_user_ai_budget_overrides_on_org_member_delete
|
||||
BEFORE DELETE ON organization_members
|
||||
FOR EACH ROW
|
||||
EXECUTE PROCEDURE delete_user_ai_budget_overrides_on_org_member_delete();
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
-- Seed a group_members row so the override below references a real
|
||||
-- membership.
|
||||
INSERT INTO group_members (
|
||||
user_id,
|
||||
group_id
|
||||
) VALUES
|
||||
('30095c71-380b-457a-8995-97b8ee6e5307', 'bb640d07-ca8a-4869-b6bc-ae61ebb2fda1')
|
||||
ON CONFLICT DO NOTHING;
|
||||
|
||||
INSERT INTO user_ai_budget_overrides (
|
||||
user_id,
|
||||
group_id,
|
||||
spend_limit_micros
|
||||
) VALUES
|
||||
('30095c71-380b-457a-8995-97b8ee6e5307', 'bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', 500000000);
|
||||
Generated
+9
@@ -5730,6 +5730,15 @@ type User struct {
|
||||
ChatSpendLimitMicros sql.NullInt64 `db:"chat_spend_limit_micros" json:"chat_spend_limit_micros"`
|
||||
}
|
||||
|
||||
// Per-user AI spend override that supersedes group budget resolution.
|
||||
type UserAiBudgetOverride struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
||||
}
|
||||
|
||||
// User-owned API keys associated with AI providers. These keys are used only when BYOK is enabled.
|
||||
type UserAiProviderKey struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
|
||||
Generated
+3
@@ -198,6 +198,7 @@ type sqlcQuerier interface {
|
||||
DeleteTailnetPeer(ctx context.Context, arg DeleteTailnetPeerParams) (DeleteTailnetPeerRow, error)
|
||||
DeleteTailnetTunnel(ctx context.Context, arg DeleteTailnetTunnelParams) (DeleteTailnetTunnelRow, error)
|
||||
DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error)
|
||||
DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error)
|
||||
DeleteUserAIProviderKey(ctx context.Context, arg DeleteUserAIProviderKeyParams) error
|
||||
DeleteUserAIProviderKeysByProviderID(ctx context.Context, aiProviderID uuid.UUID) error
|
||||
DeleteUserChatCompactionThreshold(ctx context.Context, arg DeleteUserChatCompactionThresholdParams) error
|
||||
@@ -738,6 +739,7 @@ type sqlcQuerier interface {
|
||||
// inclusive.
|
||||
GetTotalUsageDCManagedAgentsV1(ctx context.Context, arg GetTotalUsageDCManagedAgentsV1Params) (int64, error)
|
||||
GetUnexpiredLicenses(ctx context.Context) ([]License, error)
|
||||
GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error)
|
||||
GetUserAIProviderKeyByProviderID(ctx context.Context, arg GetUserAIProviderKeyByProviderIDParams) (UserAiProviderKey, error)
|
||||
// GetUserAIProviderKeys is used by dbcrypt key rotation. Request paths should use
|
||||
// user-scoped lookups instead of this bulk accessor.
|
||||
@@ -1407,6 +1409,7 @@ type sqlcQuerier interface {
|
||||
// used to store the data, and the minutes are summed for each user and template
|
||||
// combination. The result is stored in the template_usage_stats table.
|
||||
UpsertTemplateUsageStats(ctx context.Context) error
|
||||
UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAiBudgetOverride, error)
|
||||
// UpsertUserAIProviderKey preserves the original id and created_at when the
|
||||
// user/provider pair already exists. On conflict, callers provide id and
|
||||
// created_at for the insert path only.
|
||||
|
||||
Generated
+65
@@ -2449,6 +2449,23 @@ func (q *sqlQuerier) DeleteGroupAIBudget(ctx context.Context, groupID uuid.UUID)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const deleteUserAIBudgetOverride = `-- name: DeleteUserAIBudgetOverride :one
|
||||
DELETE FROM user_ai_budget_overrides WHERE user_id = $1 RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) DeleteUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, deleteUserAIBudgetOverride, userID)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one
|
||||
SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
|
||||
FROM ai_model_prices
|
||||
@@ -2494,6 +2511,25 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one
|
||||
SELECT user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
FROM user_ai_budget_overrides
|
||||
WHERE user_id = $1
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, getUserAIBudgetOverride, userID)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec
|
||||
INSERT INTO ai_model_prices (
|
||||
provider, model, input_price, output_price, cache_read_price, cache_write_price
|
||||
@@ -2548,6 +2584,35 @@ func (q *sqlQuerier) UpsertGroupAIBudget(ctx context.Context, arg UpsertGroupAIB
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertUserAIBudgetOverride = `-- name: UpsertUserAIBudgetOverride :one
|
||||
INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros)
|
||||
VALUES ($1, $2, $3)
|
||||
ON CONFLICT (user_id) DO UPDATE SET
|
||||
group_id = EXCLUDED.group_id,
|
||||
spend_limit_micros = EXCLUDED.spend_limit_micros,
|
||||
updated_at = NOW()
|
||||
RETURNING user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
`
|
||||
|
||||
type UpsertUserAIBudgetOverrideParams struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertUserAIBudgetOverride(ctx context.Context, arg UpsertUserAIBudgetOverrideParams) (UserAiBudgetOverride, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertUserAIBudgetOverride, arg.UserID, arg.GroupID, arg.SpendLimitMicros)
|
||||
var i UserAiBudgetOverride
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.GroupID,
|
||||
&i.SpendLimitMicros,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one
|
||||
SELECT
|
||||
COUNT(*)
|
||||
|
||||
@@ -40,3 +40,20 @@ RETURNING *;
|
||||
|
||||
-- name: DeleteGroupAIBudget :one
|
||||
DELETE FROM group_ai_budgets WHERE group_id = @group_id RETURNING *;
|
||||
|
||||
-- name: GetUserAIBudgetOverride :one
|
||||
SELECT *
|
||||
FROM user_ai_budget_overrides
|
||||
WHERE user_id = @user_id;
|
||||
|
||||
-- name: UpsertUserAIBudgetOverride :one
|
||||
INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros)
|
||||
VALUES (@user_id, @group_id, @spend_limit_micros)
|
||||
ON CONFLICT (user_id) DO UPDATE SET
|
||||
group_id = EXCLUDED.group_id,
|
||||
spend_limit_micros = EXCLUDED.spend_limit_micros,
|
||||
updated_at = NOW()
|
||||
RETURNING *;
|
||||
|
||||
-- name: DeleteUserAIBudgetOverride :one
|
||||
DELETE FROM user_ai_budget_overrides WHERE user_id = @user_id RETURNING *;
|
||||
|
||||
Generated
+1
@@ -97,6 +97,7 @@ const (
|
||||
UniqueTemplatesPkey UniqueConstraint = "templates_pkey" // ALTER TABLE ONLY templates ADD CONSTRAINT templates_pkey PRIMARY KEY (id);
|
||||
UniqueUsageEventsDailyPkey UniqueConstraint = "usage_events_daily_pkey" // ALTER TABLE ONLY usage_events_daily ADD CONSTRAINT usage_events_daily_pkey PRIMARY KEY (day, event_type);
|
||||
UniqueUsageEventsPkey UniqueConstraint = "usage_events_pkey" // ALTER TABLE ONLY usage_events ADD CONSTRAINT usage_events_pkey PRIMARY KEY (id);
|
||||
UniqueUserAiBudgetOverridesPkey UniqueConstraint = "user_ai_budget_overrides_pkey" // ALTER TABLE ONLY user_ai_budget_overrides ADD CONSTRAINT user_ai_budget_overrides_pkey PRIMARY KEY (user_id);
|
||||
UniqueUserAiProviderKeysPkey UniqueConstraint = "user_ai_provider_keys_pkey" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_pkey PRIMARY KEY (id);
|
||||
UniqueUserAiProviderKeysUserIDAiProviderIDKey UniqueConstraint = "user_ai_provider_keys_user_id_ai_provider_id_key" // ALTER TABLE ONLY user_ai_provider_keys ADD CONSTRAINT user_ai_provider_keys_user_id_ai_provider_id_key UNIQUE (user_id, ai_provider_id);
|
||||
UniqueUserConfigsPkey UniqueConstraint = "user_configs_pkey" // ALTER TABLE ONLY user_configs ADD CONSTRAINT user_configs_pkey PRIMARY KEY (user_id, key);
|
||||
|
||||
Reference in New Issue
Block a user