mirror of
https://github.com/coder/coder.git
synced 2026-10-31 11:02:50 +08:00
Ethan
cbe24a3da5
Add the `chat_model_config` RBAC resource, actions, and API key scopes. This layer includes the scope migration and generated API types. Model list queries apply token scopes, SQL authorization filters, and row ACLs. Write and default-management handlers keep deployment configuration authorization until the next API layer replaces it. Depends on #27956 _This pull request description was generated by Coder Agents._
1343 lines
36 KiB
Go
1343 lines
36 KiB
Go
package database
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/lib/pq"
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/coderd/rbac"
|
|
"github.com/coder/coder/v2/coderd/rbac/regosql"
|
|
)
|
|
|
|
const (
|
|
authorizedQueryPlaceholder = "-- @authorize_filter"
|
|
)
|
|
|
|
// ExpectOne can be used to convert a ':many:' query into a ':one'
|
|
// query. To reduce the quantity of SQL queries, a :many with a filter is used.
|
|
// These filters sometimes are expected to return just 1 row.
|
|
//
|
|
// A :many query will never return a sql.ErrNoRows, but a :one does.
|
|
// This function will correct the error for the empty set.
|
|
func ExpectOne[T any](ret []T, err error) (T, error) {
|
|
var empty T
|
|
if err != nil {
|
|
return empty, err
|
|
}
|
|
|
|
if len(ret) == 0 {
|
|
return empty, sql.ErrNoRows
|
|
}
|
|
|
|
if len(ret) > 1 {
|
|
return empty, xerrors.Errorf("too many rows returned, expected 1")
|
|
}
|
|
|
|
return ret[0], nil
|
|
}
|
|
|
|
// customQuerier encompasses all non-generated queries.
|
|
// It provides a flexible way to write queries for cases
|
|
// where sqlc proves inadequate.
|
|
type customQuerier interface {
|
|
templateQuerier
|
|
workspaceQuerier
|
|
userQuerier
|
|
auditLogQuerier
|
|
connectionLogQuerier
|
|
aibridgeQuerier
|
|
chatQuerier
|
|
chatModelConfigQuerier
|
|
mcpServerConfigQuerier
|
|
}
|
|
|
|
type chatModelConfigQuerier interface {
|
|
GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error)
|
|
// GetDefaultChatModelConfigCandidates returns every non-deleted config in
|
|
// every organization without an authorization filter. The caller selects
|
|
// one organization's promotion candidate from this deployment-wide list.
|
|
GetDefaultChatModelConfigCandidates(ctx context.Context) ([]ChatModelConfig, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetDefaultChatModelConfigCandidates(ctx context.Context) ([]ChatModelConfig, error) {
|
|
return q.GetChatModelConfigs(ctx)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigChatModelConfigs())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getChatModelConfigs, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedChatModelConfigs :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []ChatModelConfig
|
|
for rows.Next() {
|
|
var i ChatModelConfig
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.Model,
|
|
&i.DisplayName,
|
|
&i.CreatedBy,
|
|
&i.UpdatedBy,
|
|
&i.Enabled,
|
|
&i.IsDefault,
|
|
&i.Deleted,
|
|
&i.DeletedAt,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.ContextLimit,
|
|
&i.CompressionThreshold,
|
|
&i.Options,
|
|
&i.AIProviderID,
|
|
&i.OrganizationID,
|
|
&i.GroupACL,
|
|
&i.UserACL,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type templateQuerier interface {
|
|
GetAuthorizedTemplates(ctx context.Context, arg GetTemplatesWithFilterParams, prepared rbac.PreparedAuthorized) ([]Template, error)
|
|
GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([]TemplateGroup, error)
|
|
GetTemplateUserRoles(ctx context.Context, id uuid.UUID) ([]TemplateUser, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedTemplates(ctx context.Context, arg GetTemplatesWithFilterParams, prepared rbac.PreparedAuthorized) ([]Template, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.TemplateConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getTemplatesWithFilter, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedTemplates :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.Deleted,
|
|
arg.OrganizationID,
|
|
arg.ExactName,
|
|
arg.ExactDisplayName,
|
|
arg.FuzzyName,
|
|
arg.FuzzyDisplayName,
|
|
pq.Array(arg.IDs),
|
|
arg.Deprecated,
|
|
arg.HasAITask,
|
|
arg.AgentsAllowed,
|
|
arg.AuthorID,
|
|
arg.AuthorUsername,
|
|
arg.HasExternalAgent,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []Template
|
|
for rows.Next() {
|
|
var i Template
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.OrganizationID,
|
|
&i.Deleted,
|
|
&i.Name,
|
|
&i.Provisioner,
|
|
&i.ActiveVersionID,
|
|
&i.Description,
|
|
&i.DefaultTTL,
|
|
&i.CreatedBy,
|
|
&i.Icon,
|
|
&i.UserACL,
|
|
&i.GroupACL,
|
|
&i.DisplayName,
|
|
&i.AllowUserCancelWorkspaceJobs,
|
|
&i.AllowUserAutostart,
|
|
&i.AllowUserAutostop,
|
|
&i.FailureTTL,
|
|
&i.TimeTilDormant,
|
|
&i.TimeTilDormantAutoDelete,
|
|
&i.AutostopRequirementDaysOfWeek,
|
|
&i.AutostopRequirementWeeks,
|
|
&i.AutostartBlockDaysOfWeek,
|
|
&i.RequireActiveVersion,
|
|
&i.Deprecated,
|
|
&i.ActivityBump,
|
|
&i.MaxPortSharingLevel,
|
|
&i.UseClassicParameterFlow,
|
|
&i.CorsBehavior,
|
|
&i.DisableModuleCache,
|
|
&i.TimeTilAutostopNotify,
|
|
&i.AgentsAllowed,
|
|
&i.CreatedByAvatarURL,
|
|
&i.CreatedByUsername,
|
|
&i.CreatedByName,
|
|
&i.OrganizationName,
|
|
&i.OrganizationDisplayName,
|
|
&i.OrganizationIcon,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type TemplateUser struct {
|
|
User
|
|
Actions Actions `db:"actions"`
|
|
}
|
|
|
|
func (q *sqlQuerier) GetTemplateUserRoles(ctx context.Context, id uuid.UUID) ([]TemplateUser, error) {
|
|
const query = `
|
|
SELECT
|
|
perms.value as actions, users.*
|
|
FROM
|
|
users
|
|
JOIN
|
|
(
|
|
SELECT
|
|
*
|
|
FROM
|
|
jsonb_each_text(
|
|
(
|
|
SELECT
|
|
templates.user_acl
|
|
FROM
|
|
templates
|
|
WHERE
|
|
id = $1
|
|
)
|
|
)
|
|
) AS perms
|
|
ON
|
|
users.id::text = perms.key
|
|
WHERE
|
|
users.deleted = false
|
|
AND
|
|
users.status != 'suspended';
|
|
`
|
|
|
|
var tus []TemplateUser
|
|
err := q.db.SelectContext(ctx, &tus, query, id.String())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("select user actions: %w", err)
|
|
}
|
|
|
|
return tus, nil
|
|
}
|
|
|
|
type TemplateGroup struct {
|
|
Group
|
|
Actions Actions `db:"actions"`
|
|
}
|
|
|
|
func (q *sqlQuerier) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([]TemplateGroup, error) {
|
|
const query = `
|
|
SELECT
|
|
perms.value as actions, groups.*
|
|
FROM
|
|
groups
|
|
JOIN
|
|
(
|
|
SELECT
|
|
*
|
|
FROM
|
|
jsonb_each_text(
|
|
(
|
|
SELECT
|
|
templates.group_acl
|
|
FROM
|
|
templates
|
|
WHERE
|
|
id = $1
|
|
)
|
|
)
|
|
) AS perms
|
|
ON
|
|
groups.id::text = perms.key;
|
|
`
|
|
|
|
var tgs []TemplateGroup
|
|
err := q.db.SelectContext(ctx, &tgs, query, id.String())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("select group roles: %w", err)
|
|
}
|
|
|
|
return tgs, nil
|
|
}
|
|
|
|
type workspaceQuerier interface {
|
|
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, prepared rbac.PreparedAuthorized) ([]GetWorkspacesRow, error)
|
|
GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsByOwnerIDRow, error)
|
|
}
|
|
|
|
// GetAuthorizedWorkspaces returns all workspaces that the user is authorized to access.
|
|
// This code is copied from `GetWorkspaces` and adds the authorized filter WHERE
|
|
// clause.
|
|
func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, prepared rbac.PreparedAuthorized) ([]GetWorkspacesRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigWorkspaces())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
// In order to properly use ORDER BY, OFFSET, and LIMIT, we need to inject the
|
|
// authorizedFilter between the end of the where clause and those statements.
|
|
filtered, err := insertAuthorizedFilter(getWorkspaces, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedWorkspaces :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
pq.Array(arg.IncludeAgentMetadata),
|
|
pq.Array(arg.ParamNames),
|
|
pq.Array(arg.ParamValues),
|
|
arg.Deleted,
|
|
arg.Status,
|
|
arg.OwnerID,
|
|
arg.OrganizationID,
|
|
pq.Array(arg.HasParam),
|
|
arg.OwnerUsername,
|
|
arg.TemplateName,
|
|
pq.Array(arg.TemplateIDs),
|
|
pq.Array(arg.WorkspaceIds),
|
|
arg.Name,
|
|
pq.Array(arg.HasAgentStatuses),
|
|
arg.AgentInactiveDisconnectTimeoutSeconds,
|
|
arg.Dormant,
|
|
arg.LastUsedBefore,
|
|
arg.LastUsedAfter,
|
|
arg.UsingActive,
|
|
arg.HasAITask,
|
|
arg.HasExternalAgent,
|
|
arg.Shared,
|
|
arg.SharedWithUserID,
|
|
arg.SharedWithGroupID,
|
|
arg.RequesterID,
|
|
arg.Offset,
|
|
arg.Limit,
|
|
arg.WithSummary,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetWorkspacesRow
|
|
for rows.Next() {
|
|
var i GetWorkspacesRow
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.OwnerID,
|
|
&i.OrganizationID,
|
|
&i.TemplateID,
|
|
&i.Deleted,
|
|
&i.Name,
|
|
&i.AutostartSchedule,
|
|
&i.Ttl,
|
|
&i.LastUsedAt,
|
|
&i.DormantAt,
|
|
&i.DeletingAt,
|
|
&i.AutomaticUpdates,
|
|
&i.Favorite,
|
|
&i.NextStartAt,
|
|
&i.GroupACL,
|
|
&i.UserACL,
|
|
&i.OwnerAvatarUrl,
|
|
&i.OwnerUsername,
|
|
&i.OwnerName,
|
|
&i.OrganizationName,
|
|
&i.OrganizationDisplayName,
|
|
&i.OrganizationIcon,
|
|
&i.OrganizationDescription,
|
|
&i.TemplateName,
|
|
&i.TemplateDisplayName,
|
|
&i.TemplateIcon,
|
|
&i.TemplateDescription,
|
|
&i.TaskID,
|
|
&i.GroupACLDisplayInfo,
|
|
&i.UserACLDisplayInfo,
|
|
&i.TemplateVersionID,
|
|
&i.TemplateVersionName,
|
|
&i.LatestBuildCompletedAt,
|
|
&i.LatestBuildCanceledAt,
|
|
&i.LatestBuildError,
|
|
&i.LatestBuildTransition,
|
|
&i.LatestBuildStatus,
|
|
&i.LatestBuildHasExternalAgent,
|
|
&i.LatestBuildProvisionerJobID,
|
|
&i.AgentMetadata,
|
|
&i.Count,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsByOwnerIDRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigWorkspaces())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
// In order to properly use ORDER BY, OFFSET, and LIMIT, we need to inject the
|
|
// authorizedFilter between the end of the where clause and those statements.
|
|
filtered, err := insertAuthorizedFilter(getWorkspacesAndAgentsByOwnerID, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedWorkspacesAndAgentsByOwnerID :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query, ownerID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetWorkspacesAndAgentsByOwnerIDRow
|
|
for rows.Next() {
|
|
var i GetWorkspacesAndAgentsByOwnerIDRow
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.Name,
|
|
&i.JobStatus,
|
|
&i.Transition,
|
|
pq.Array(&i.Agents),
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type userQuerier interface {
|
|
GetAuthorizedUsers(ctx context.Context, arg GetUsersParams, prepared rbac.PreparedAuthorized) ([]GetUsersRow, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedUsers(ctx context.Context, arg GetUsersParams, prepared rbac.PreparedAuthorized) ([]GetUsersRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.UserConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getUsers, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: GetAuthorizedUsers :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.AfterID,
|
|
arg.Search,
|
|
arg.Name,
|
|
arg.ExactUsername,
|
|
arg.ExactEmail,
|
|
pq.Array(arg.Status),
|
|
pq.Array(arg.RbacRole),
|
|
arg.LastSeenBefore,
|
|
arg.LastSeenAfter,
|
|
arg.CreatedBefore,
|
|
arg.CreatedAfter,
|
|
arg.IncludeSystem,
|
|
arg.GithubComUserID,
|
|
pq.Array(arg.LoginType),
|
|
arg.IsServiceAccount,
|
|
arg.OffsetOpt,
|
|
arg.LimitOpt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetUsersRow
|
|
for rows.Next() {
|
|
var i GetUsersRow
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.Email,
|
|
&i.Username,
|
|
&i.HashedPassword,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.Status,
|
|
&i.RBACRoles,
|
|
&i.LoginType,
|
|
&i.AvatarURL,
|
|
&i.Deleted,
|
|
&i.LastSeenAt,
|
|
&i.QuietHoursSchedule,
|
|
&i.Name,
|
|
&i.GithubComUserID,
|
|
&i.HashedOneTimePasscode,
|
|
&i.OneTimePasscodeExpiresAt,
|
|
&i.IsSystem,
|
|
&i.IsServiceAccount,
|
|
&i.ChatSpendLimitMicros,
|
|
&i.Count,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type auditLogQuerier interface {
|
|
GetAuthorizedAuditLogsOffset(ctx context.Context, arg GetAuditLogsOffsetParams, prepared rbac.PreparedAuthorized) ([]GetAuditLogsOffsetRow, error)
|
|
CountAuthorizedAuditLogs(ctx context.Context, arg CountAuditLogsParams, prepared rbac.PreparedAuthorized) (int64, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedAuditLogsOffset(ctx context.Context, arg GetAuditLogsOffsetParams, prepared rbac.PreparedAuthorized) ([]GetAuditLogsOffsetRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AuditLogConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getAuditLogsOffset, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: GetAuthorizedAuditLogsOffset :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.ResourceType,
|
|
arg.ResourceID,
|
|
arg.OrganizationID,
|
|
arg.ResourceTarget,
|
|
arg.Action,
|
|
arg.UserID,
|
|
arg.Username,
|
|
arg.Email,
|
|
arg.DateFrom,
|
|
arg.DateTo,
|
|
arg.BuildReason,
|
|
arg.RequestID,
|
|
arg.OffsetOpt,
|
|
arg.LimitOpt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetAuditLogsOffsetRow
|
|
for rows.Next() {
|
|
var i GetAuditLogsOffsetRow
|
|
if err := rows.Scan(
|
|
&i.AuditLog.ID,
|
|
&i.AuditLog.Time,
|
|
&i.AuditLog.UserID,
|
|
&i.AuditLog.OrganizationID,
|
|
&i.AuditLog.Ip,
|
|
&i.AuditLog.UserAgent,
|
|
&i.AuditLog.ResourceType,
|
|
&i.AuditLog.ResourceID,
|
|
&i.AuditLog.ResourceTarget,
|
|
&i.AuditLog.Action,
|
|
&i.AuditLog.Diff,
|
|
&i.AuditLog.StatusCode,
|
|
&i.AuditLog.AdditionalFields,
|
|
&i.AuditLog.RequestID,
|
|
&i.AuditLog.ResourceIcon,
|
|
&i.UserUsername,
|
|
&i.UserName,
|
|
&i.UserEmail,
|
|
&i.UserCreatedAt,
|
|
&i.UserUpdatedAt,
|
|
&i.UserLastSeenAt,
|
|
&i.UserStatus,
|
|
&i.UserLoginType,
|
|
&i.UserRoles,
|
|
&i.UserAvatarUrl,
|
|
&i.UserDeleted,
|
|
&i.UserQuietHoursSchedule,
|
|
&i.OrganizationName,
|
|
&i.OrganizationDisplayName,
|
|
&i.OrganizationIcon,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) CountAuthorizedAuditLogs(ctx context.Context, arg CountAuditLogsParams, prepared rbac.PreparedAuthorized) (int64, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AuditLogConverter(),
|
|
})
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(countAuditLogs, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: CountAuthorizedAuditLogs :one\n%s", filtered)
|
|
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.ResourceType,
|
|
arg.ResourceID,
|
|
arg.OrganizationID,
|
|
arg.ResourceTarget,
|
|
arg.Action,
|
|
arg.UserID,
|
|
arg.Username,
|
|
arg.Email,
|
|
arg.DateFrom,
|
|
arg.DateTo,
|
|
arg.BuildReason,
|
|
arg.RequestID,
|
|
arg.CountCap,
|
|
)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer rows.Close()
|
|
var count int64
|
|
for rows.Next() {
|
|
if err := rows.Scan(&count); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
type connectionLogQuerier interface {
|
|
GetAuthorizedConnectionLogsOffset(ctx context.Context, arg GetConnectionLogsOffsetParams, prepared rbac.PreparedAuthorized) ([]GetConnectionLogsOffsetRow, error)
|
|
CountAuthorizedConnectionLogs(ctx context.Context, arg CountConnectionLogsParams, prepared rbac.PreparedAuthorized) (int64, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedConnectionLogsOffset(ctx context.Context, arg GetConnectionLogsOffsetParams, prepared rbac.PreparedAuthorized) ([]GetConnectionLogsOffsetRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.ConnectionLogConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(getConnectionLogsOffset, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: GetAuthorizedConnectionLogsOffset :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.OrganizationID,
|
|
arg.WorkspaceOwner,
|
|
arg.WorkspaceOwnerID,
|
|
arg.WorkspaceOwnerEmail,
|
|
arg.Type,
|
|
arg.UserID,
|
|
arg.Username,
|
|
arg.UserEmail,
|
|
arg.ConnectedAfter,
|
|
arg.ConnectedBefore,
|
|
arg.WorkspaceID,
|
|
arg.ConnectionID,
|
|
arg.Status,
|
|
arg.OffsetOpt,
|
|
arg.LimitOpt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetConnectionLogsOffsetRow
|
|
for rows.Next() {
|
|
var i GetConnectionLogsOffsetRow
|
|
if err := rows.Scan(
|
|
&i.ConnectionLog.ID,
|
|
&i.ConnectionLog.ConnectTime,
|
|
&i.ConnectionLog.OrganizationID,
|
|
&i.ConnectionLog.WorkspaceOwnerID,
|
|
&i.ConnectionLog.WorkspaceID,
|
|
&i.ConnectionLog.WorkspaceName,
|
|
&i.ConnectionLog.AgentName,
|
|
&i.ConnectionLog.Type,
|
|
&i.ConnectionLog.Ip,
|
|
&i.ConnectionLog.Code,
|
|
&i.ConnectionLog.UserAgent,
|
|
&i.ConnectionLog.UserID,
|
|
&i.ConnectionLog.SlugOrPort,
|
|
&i.ConnectionLog.ConnectionID,
|
|
&i.ConnectionLog.DisconnectTime,
|
|
&i.ConnectionLog.DisconnectReason,
|
|
&i.UserUsername,
|
|
&i.UserName,
|
|
&i.UserEmail,
|
|
&i.UserCreatedAt,
|
|
&i.UserUpdatedAt,
|
|
&i.UserLastSeenAt,
|
|
&i.UserStatus,
|
|
&i.UserLoginType,
|
|
&i.UserRoles,
|
|
&i.UserAvatarUrl,
|
|
&i.UserDeleted,
|
|
&i.UserQuietHoursSchedule,
|
|
&i.WorkspaceOwnerUsername,
|
|
&i.OrganizationName,
|
|
&i.OrganizationDisplayName,
|
|
&i.OrganizationIcon,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) CountAuthorizedConnectionLogs(ctx context.Context, arg CountConnectionLogsParams, prepared rbac.PreparedAuthorized) (int64, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.ConnectionLogConverter(),
|
|
})
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(countConnectionLogs, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: CountAuthorizedConnectionLogs :one\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.OrganizationID,
|
|
arg.WorkspaceOwner,
|
|
arg.WorkspaceOwnerID,
|
|
arg.WorkspaceOwnerEmail,
|
|
arg.Type,
|
|
arg.UserID,
|
|
arg.Username,
|
|
arg.UserEmail,
|
|
arg.ConnectedAfter,
|
|
arg.ConnectedBefore,
|
|
arg.WorkspaceID,
|
|
arg.ConnectionID,
|
|
arg.Status,
|
|
arg.CountCap,
|
|
)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer rows.Close()
|
|
var count int64
|
|
for rows.Next() {
|
|
if err := rows.Scan(&count); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
type chatQuerier interface {
|
|
DeleteOldChatFiles(ctx context.Context, arg DeleteOldChatFilesParams) (int64, error)
|
|
LinkChatFiles(ctx context.Context, arg LinkChatFilesParams) (int32, error)
|
|
GetAuthorizedChats(ctx context.Context, arg GetChatsParams, prepared rbac.PreparedAuthorized) ([]GetChatsRow, error)
|
|
GetAuthorizedChatsByChatFileID(ctx context.Context, fileID uuid.UUID, prepared rbac.PreparedAuthorized) ([]Chat, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams, prepared rbac.PreparedAuthorized) ([]GetChatsRow, error) {
|
|
if (arg.OwnedOnly || arg.SharedOnly) && arg.ViewerID == uuid.Nil {
|
|
return nil, xerrors.New("viewer_id required when owned_only or shared_only is true")
|
|
}
|
|
if arg.SharedOnly && arg.SharedWithUserID == uuid.Nil && len(arg.SharedWithGroupIds) == 0 {
|
|
return nil, xerrors.New("shared_with_user_id or shared_with_group_ids required when shared_only is true")
|
|
}
|
|
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigChats())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getChats, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedChats :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.OwnedOnly,
|
|
arg.SharedOnly,
|
|
arg.ViewerID,
|
|
arg.SharedWithUserID,
|
|
pq.Array(arg.SharedWithGroupIds),
|
|
arg.Archived,
|
|
arg.AfterID,
|
|
arg.LabelFilter,
|
|
arg.DiffURL,
|
|
arg.TitleQuery,
|
|
arg.HasUnread,
|
|
pq.Array(arg.PullRequestStatuses),
|
|
arg.PrNumber,
|
|
arg.RepoQuery,
|
|
arg.PrTitleQuery,
|
|
arg.Search,
|
|
arg.OffsetOpt,
|
|
arg.LimitOpt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []GetChatsRow
|
|
for rows.Next() {
|
|
var i GetChatsRow
|
|
if err := rows.Scan(
|
|
&i.Chat.ID,
|
|
&i.Chat.OwnerID,
|
|
&i.Chat.WorkspaceID,
|
|
&i.Chat.Title,
|
|
&i.Chat.Status,
|
|
&i.Chat.WorkerID,
|
|
&i.Chat.StartedAt,
|
|
&i.Chat.HeartbeatAt,
|
|
&i.Chat.CreatedAt,
|
|
&i.Chat.UpdatedAt,
|
|
&i.Chat.ParentChatID,
|
|
&i.Chat.RootChatID,
|
|
&i.Chat.LastModelConfigID,
|
|
&i.Chat.LastReasoningEffort,
|
|
&i.Chat.Archived,
|
|
&i.Chat.LastError,
|
|
&i.Chat.Mode,
|
|
pq.Array(&i.Chat.MCPServerIDs),
|
|
&i.Chat.Labels,
|
|
&i.Chat.BuildID,
|
|
&i.Chat.AgentID,
|
|
&i.Chat.PinOrder,
|
|
&i.Chat.LastReadMessageID,
|
|
&i.Chat.DynamicTools,
|
|
&i.Chat.OrganizationID,
|
|
&i.Chat.PlanMode,
|
|
&i.Chat.ClientType,
|
|
&i.Chat.LastTurnSummary,
|
|
&i.Chat.Summary,
|
|
&i.Chat.SummaryGeneratedAt,
|
|
&i.Chat.SnapshotVersion,
|
|
&i.Chat.HistoryVersion,
|
|
&i.Chat.QueueVersion,
|
|
&i.Chat.GenerationAttempt,
|
|
&i.Chat.RetryState,
|
|
&i.Chat.RetryStateVersion,
|
|
&i.Chat.RunnerID,
|
|
&i.Chat.RequiresActionDeadlineAt,
|
|
&i.Chat.UserACL,
|
|
&i.Chat.GroupACL,
|
|
&i.Chat.OwnerUsername,
|
|
&i.Chat.OwnerName,
|
|
&i.Chat.ContextAggregateHash,
|
|
&i.Chat.ContextDirtySince,
|
|
&i.Chat.ContextDirtyResources,
|
|
&i.Chat.ContextError,
|
|
&i.Chat.CompactionRequestedAt,
|
|
&i.HasUnread); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID uuid.UUID, prepared rbac.PreparedAuthorized) ([]Chat, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigChats())
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getChatsByChatFileID, fmt.Sprintf(" AND %s\nLIMIT 1", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: GetAuthorizedChatsByChatFileID :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query, fileID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []Chat
|
|
for rows.Next() {
|
|
var i Chat
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.OwnerID,
|
|
&i.WorkspaceID,
|
|
&i.Title,
|
|
&i.Status,
|
|
&i.WorkerID,
|
|
&i.StartedAt,
|
|
&i.HeartbeatAt,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.ParentChatID,
|
|
&i.RootChatID,
|
|
&i.LastModelConfigID,
|
|
&i.LastReasoningEffort,
|
|
&i.Archived,
|
|
&i.LastError,
|
|
&i.Mode,
|
|
pq.Array(&i.MCPServerIDs),
|
|
&i.Labels,
|
|
&i.BuildID,
|
|
&i.AgentID,
|
|
&i.PinOrder,
|
|
&i.LastReadMessageID,
|
|
&i.DynamicTools,
|
|
&i.OrganizationID,
|
|
&i.PlanMode,
|
|
&i.ClientType,
|
|
&i.LastTurnSummary,
|
|
&i.Summary,
|
|
&i.SummaryGeneratedAt,
|
|
&i.SnapshotVersion,
|
|
&i.HistoryVersion,
|
|
&i.QueueVersion,
|
|
&i.GenerationAttempt,
|
|
&i.RetryState,
|
|
&i.RetryStateVersion,
|
|
&i.RunnerID,
|
|
&i.RequiresActionDeadlineAt,
|
|
&i.UserACL,
|
|
&i.GroupACL,
|
|
&i.OwnerUsername,
|
|
&i.OwnerName,
|
|
&i.ContextAggregateHash,
|
|
&i.ContextDirtySince,
|
|
&i.ContextDirtyResources,
|
|
&i.ContextError,
|
|
&i.CompactionRequestedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
type aibridgeQuerier interface {
|
|
ListAuthorizedAIBridgeModels(ctx context.Context, arg ListAIBridgeModelsParams, prepared rbac.PreparedAuthorized) ([]string, error)
|
|
ListAuthorizedAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, error)
|
|
ListAuthorizedAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionsRow, error)
|
|
CountAuthorizedAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) (int64, error)
|
|
ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionThreadsRow, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) ListAuthorizedAIBridgeModels(ctx context.Context, arg ListAIBridgeModelsParams, prepared rbac.PreparedAuthorized) ([]string, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AIBridgeInterceptionConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(listAIBridgeModels, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: ListAIBridgeModels :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query, arg.Model, arg.Offset, arg.Limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []string
|
|
for rows.Next() {
|
|
var model string
|
|
if err := rows.Scan(&model); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, model)
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) ListAuthorizedAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AIBridgeInterceptionConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(listAIBridgeClients, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: ListAIBridgeClients :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query, arg.Client, arg.Offset, arg.Limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []string
|
|
for rows.Next() {
|
|
var client string
|
|
if err := rows.Scan(&client); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, client)
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionsRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AIBridgeInterceptionConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(listAIBridgeSessions, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: ListAuthorizedAIBridgeSessions :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.AfterSessionID,
|
|
arg.StartedAfter,
|
|
arg.StartedBefore,
|
|
arg.InitiatorID,
|
|
arg.Provider,
|
|
arg.ProviderName,
|
|
arg.Model,
|
|
arg.Client,
|
|
arg.SessionID,
|
|
arg.Offset,
|
|
arg.Limit,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []ListAIBridgeSessionsRow
|
|
for rows.Next() {
|
|
var i ListAIBridgeSessionsRow
|
|
if err := rows.Scan(
|
|
&i.SessionID,
|
|
&i.UserID,
|
|
&i.UserUsername,
|
|
&i.UserName,
|
|
&i.UserAvatarUrl,
|
|
pq.Array(&i.Providers),
|
|
pq.Array(&i.Models),
|
|
&i.Client,
|
|
&i.Metadata,
|
|
&i.StartedAt,
|
|
&i.EndedAt,
|
|
&i.Threads,
|
|
&i.InputTokens,
|
|
&i.OutputTokens,
|
|
&i.CacheReadInputTokens,
|
|
&i.CacheWriteInputTokens,
|
|
&i.LastPrompt,
|
|
&i.LastActiveAt,
|
|
&i.NetworkCallsTotal,
|
|
&i.NetworkCallsBlocked,
|
|
&i.FirewallActive,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) CountAuthorizedAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) (int64, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AIBridgeInterceptionConverter(),
|
|
})
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(countAIBridgeSessions, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return 0, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: CountAuthorizedAIBridgeSessions :one\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.StartedAfter,
|
|
arg.StartedBefore,
|
|
arg.InitiatorID,
|
|
arg.Provider,
|
|
arg.ProviderName,
|
|
arg.Model,
|
|
arg.Client,
|
|
arg.SessionID,
|
|
)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer rows.Close()
|
|
var count int64
|
|
for rows.Next() {
|
|
if err := rows.Scan(&count); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func (q *sqlQuerier) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionThreadsRow, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.AIBridgeInterceptionConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
filtered, err := insertAuthorizedFilter(listAIBridgeSessionThreads, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
query := fmt.Sprintf("-- name: ListAuthorizedAIBridgeSessionThreads :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query,
|
|
arg.SessionID,
|
|
arg.AfterID,
|
|
arg.BeforeID,
|
|
arg.Limit,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []ListAIBridgeSessionThreadsRow
|
|
for rows.Next() {
|
|
var i ListAIBridgeSessionThreadsRow
|
|
if err := rows.Scan(
|
|
&i.ThreadID,
|
|
&i.AIBridgeInterception.ID,
|
|
&i.AIBridgeInterception.InitiatorID,
|
|
&i.AIBridgeInterception.Provider,
|
|
&i.AIBridgeInterception.Model,
|
|
&i.AIBridgeInterception.StartedAt,
|
|
&i.AIBridgeInterception.Metadata,
|
|
&i.AIBridgeInterception.EndedAt,
|
|
&i.AIBridgeInterception.APIKeyID,
|
|
&i.AIBridgeInterception.Client,
|
|
&i.AIBridgeInterception.ThreadParentID,
|
|
&i.AIBridgeInterception.ThreadRootID,
|
|
&i.AIBridgeInterception.ClientSessionID,
|
|
&i.AIBridgeInterception.SessionID,
|
|
&i.AIBridgeInterception.ProviderName,
|
|
&i.AIBridgeInterception.CredentialKind,
|
|
&i.AIBridgeInterception.CredentialHint,
|
|
&i.AIBridgeInterception.AgentFirewallSessionID,
|
|
&i.AIBridgeInterception.AgentFirewallSequenceNumber,
|
|
&i.AIBridgeInterception.ErrorType,
|
|
&i.AIBridgeInterception.ErrorMessage,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
func insertAuthorizedFilter(query string, replaceWith string) (string, error) {
|
|
if !strings.Contains(query, authorizedQueryPlaceholder) {
|
|
return "", xerrors.Errorf("query does not contain authorized replace string, this is not an authorized query")
|
|
}
|
|
filtered := strings.ReplaceAll(query, authorizedQueryPlaceholder, replaceWith)
|
|
return filtered, nil
|
|
}
|
|
|
|
// UpdateUserLinkRawJSON is a custom query for unit testing. Do not ever expose this
|
|
func (q *sqlQuerier) UpdateUserLinkRawJSON(ctx context.Context, userID uuid.UUID, data json.RawMessage) error {
|
|
_, err := q.sdb.ExecContext(ctx, "UPDATE user_links SET claims = $2 WHERE user_id = $1", userID, data)
|
|
return err
|
|
}
|
|
|
|
type mcpServerConfigQuerier interface {
|
|
GetAuthorizedMCPServerConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]MCPServerConfig, error)
|
|
}
|
|
|
|
func (q *sqlQuerier) GetAuthorizedMCPServerConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]MCPServerConfig, error) {
|
|
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
|
|
VariableConverter: regosql.MCPServerConfigConverter(),
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
|
}
|
|
|
|
filtered, err := insertAuthorizedFilter(getMCPServerConfigsByOrganization, fmt.Sprintf(" AND %s", authorizedFilter))
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("insert authorized filter: %w", err)
|
|
}
|
|
|
|
// The name comment is for metric tracking
|
|
query := fmt.Sprintf("-- name: GetAuthorizedMCPServerConfigs :many\n%s", filtered)
|
|
rows, err := q.db.QueryContext(ctx, query, organizationID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []MCPServerConfig
|
|
for rows.Next() {
|
|
var i MCPServerConfig
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.DisplayName,
|
|
&i.Slug,
|
|
&i.Description,
|
|
&i.IconURL,
|
|
&i.Transport,
|
|
&i.Url,
|
|
&i.AuthType,
|
|
&i.OAuth2ClientID,
|
|
&i.OAuth2ClientSecret,
|
|
&i.OAuth2ClientSecretKeyID,
|
|
&i.OAuth2AuthURL,
|
|
&i.OAuth2TokenURL,
|
|
&i.OAuth2Scopes,
|
|
&i.APIKeyHeader,
|
|
&i.APIKeyValue,
|
|
&i.APIKeyValueKeyID,
|
|
&i.CustomHeaders,
|
|
&i.CustomHeadersKeyID,
|
|
pq.Array(&i.ToolAllowList),
|
|
pq.Array(&i.ToolDenyList),
|
|
&i.Availability,
|
|
&i.Enabled,
|
|
&i.CreatedBy,
|
|
&i.UpdatedBy,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
&i.ModelIntent,
|
|
&i.AllowInPlanMode,
|
|
&i.ForwardCoderHeaders,
|
|
&i.OAuth2RevocationURL,
|
|
&i.OrganizationID,
|
|
&i.GroupACL,
|
|
&i.UserACL,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|