mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-29 18:19:07 +02:00
[management,proxy] Agent network: per-account LLM gateway (policy, metering, multi-provider) (#6555)
* [agent-network] Shared proto, OpenAPI schema, and generated types * [agent-network] Management: store, manager, synthesizer, policy engine, provider catalog, HTTP/gRPC API Adds the account-scoped agent-network module: provider/policy/budget CRUD and store, the reverse-proxy service synthesizer, policy selection + limit enforcement, the provider catalog (incl. Vertex AI and AWS Bedrock entries), and the management HTTP + proxy gRPC surfaces. * [management] Fix agent-network proxy-peer fan-out on affected-peer recompute The affected-peers resolver loaded only persisted reverse-proxy services, but agent-network services are synthesized on demand and never persisted. As a result the embedded proxy peer was never folded into the affected set when a client's group changed, so the proxy received no network-map update for a newly authorised client and rejected its handshake until a full resync (restart). loadProxyServices now merges the synthesized agent-network services (injected via a registration hook to avoid an import cycle), so proxy peers learn newly authorised clients immediately. * [proxy] Reverse-proxy middleware framework, chain, and request plumbing The per-target middleware chain (slots, dispatcher, mutation gate, metadata merger), body capture, access-log terminal sink, and the proxy wiring that builds + runs chains for synthesized agent-network services. * [proxy] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) Request/response parsers and SSE/event-stream metering, the embedded pricing table, and the builtin middleware set: request parser, router, policy limit-check/record, cost meter, guardrail, identity inject, response parser. Includes the path-routed providers — Google Vertex AI (keyfile:: service-account OAuth minting) and AWS Bedrock (bearer auth, invoke/converse/streaming, optional /bedrock prefix) — plus the Models allowlist and unmeterable-publisher deny. * [proxy] IPv6 in-place apply and TCP accept-loop hardening on netstack listeners * [agent-network] End-to-end test suite, module docs, and deployment preset * [agent-network] Fix codespell typos and exclude false positives - labelgen word pool: vermillion -> vermilion, racoon -> raccoon. - codespell ignore list: add flate (Go compress/flate package), recordin (a test-local identifier), and unparseable (a valid alternative spelling used consistently across identifiers + a metadata-value constant). * [management] Set LastSeen on injected proxy peer in realstack test (MySQL strict-mode) The injected embedded proxy peer had a PeerStatus with a zero LastSeen, which serializes to '0000-00-00' and is rejected by MySQL in strict mode (SQLite tolerates it). Set LastSeen to a valid time so SaveAccount succeeds on both engines. * [agent-network] Remove e2e shell-script suite from this branch The end-to-end shell scripts under scripts/e2e/ are maintained in a separate testing suite and are not part of this change set. * [agent-network] Polish module docs: remove internal review scaffolding, fix links, verify diagrams Strip PR-review framing, commit references, absolute paths, and stale internal references from the agent-network module docs; fix broken relative links; verify all diagrams against the current architecture. Remove the internal AI-reviewer prompt file. * [management] Refine session expiration handling to support 3-state encoding for SSO deadlines * [agent-network] Relocate agentnetwork package to internals/modules Move management/server/agentnetwork (and its catalog/, labelgen/, types/ subpackages) to management/internals/modules/agentnetwork, alongside the reverse-proxy module, and rewrite all importers. Pure relocation: package names, the synthesizer + affectedpeers registration hook, and store access (shared store.Store) are unchanged, so no import cycle is introduced (affectedpeers still depends only on the agentnetwork/types leaf). * [agent-network] Co-locate HTTP handlers in the module (RegisterEndpoints) Move the agent-network HTTP handlers from server/http/handlers/agentnetwork into the module at internals/modules/agentnetwork/handlers (package handlers) and rename the entrypoint AddEndpoints -> RegisterEndpoints, matching the reverse-proxy module convention. Wiring in http/handler.go updated accordingly.
This commit is contained in:
Maycon Santos
@@ -37,6 +37,7 @@ import (
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
|
||||
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
|
||||
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
|
||||
@@ -137,6 +138,10 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
|
||||
&networkTypes.Network{}, &routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{}, &types.AccountOnboarding{},
|
||||
&types.Job{}, &zones.Zone{}, &records.Record{}, &types.UserInviteRecord{}, &rpservice.Service{}, &rpservice.Target{}, &domain.Domain{},
|
||||
&accesslogs.AccessLogEntry{}, &proxy.Proxy{},
|
||||
&agentNetworkTypes.Provider{}, &agentNetworkTypes.Policy{}, &agentNetworkTypes.Guardrail{}, &agentNetworkTypes.Settings{},
|
||||
&agentNetworkTypes.Consumption{}, &agentNetworkTypes.AccountBudgetRule{},
|
||||
&agentNetworkTypes.AgentNetworkAccessLog{}, &agentNetworkTypes.AgentNetworkAccessLogGroup{},
|
||||
&agentNetworkTypes.AgentNetworkUsage{}, &agentNetworkTypes.AgentNetworkUsageGroup{},
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("auto migratePreAuto: %w", err)
|
||||
@@ -5573,6 +5578,255 @@ func (s *SqlStore) CreateAccessLog(ctx context.Context, logEntry *accesslogs.Acc
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateAgentNetworkAccessLog persists a flattened agent-network access-log
|
||||
// entry together with its authorising-group child rows in a single
|
||||
// transaction.
|
||||
func (s *SqlStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Idempotent on the log id / (log_id, group_id) so a proxy resend of the
|
||||
// same entry can't fail the request.
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(entry).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(groups) > 0 {
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).WithFields(log.Fields{
|
||||
"account_id": entry.AccountID,
|
||||
"service_id": entry.ServiceID,
|
||||
"model": entry.Model,
|
||||
}).Errorf("failed to create agent-network access log entry in store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to create agent-network access log entry in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateAgentNetworkUsage persists a stripped agent-network usage record
|
||||
// together with its authorising-group child rows in a single transaction.
|
||||
func (s *SqlStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Idempotent on the usage id / (usage_id, group_id) so a proxy resend of
|
||||
// the same entry can't fail the request.
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(usage).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(groups) > 0 {
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).WithFields(log.Fields{
|
||||
"account_id": usage.AccountID,
|
||||
"model": usage.Model,
|
||||
}).Errorf("failed to create agent-network usage record in store: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to create agent-network usage record in store")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteOldAgentNetworkAccessLogs deletes an account's access-log rows (and
|
||||
// their authorising-group child rows) older than the cutoff. Usage records are
|
||||
// untouched — they are the long-term aggregate. Returns the number of log rows
|
||||
// deleted.
|
||||
func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
|
||||
var deleted int64
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Remove group child rows for the soon-to-be-deleted logs first.
|
||||
if err := tx.Exec(
|
||||
"DELETE FROM agent_network_access_log_group WHERE account_id = ? AND log_id IN (SELECT id FROM agent_network_access_log WHERE account_id = ? AND timestamp < ?)",
|
||||
accountID, accountID, olderThan,
|
||||
).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
res := tx.Where("account_id = ? AND timestamp < ?", accountID, olderThan).
|
||||
Delete(&agentNetworkTypes.AgentNetworkAccessLog{})
|
||||
if res.Error != nil {
|
||||
return res.Error
|
||||
}
|
||||
deleted = res.RowsAffected
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete old agent-network access logs for account %s: %v", accountID, err)
|
||||
return 0, status.Errorf(status.Internal, "failed to delete old agent-network access logs")
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkUsageRows returns the stripped usage rows for an account that
|
||||
// match the filter (date / user / group / provider / model). Aggregation into
|
||||
// time buckets happens in the manager so granularities stay engine-portable.
|
||||
func (s *SqlStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
|
||||
var rows []*agentNetworkTypes.AgentNetworkUsage
|
||||
|
||||
query := s.applyAgentNetworkUsageFilters(
|
||||
s.db.Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).Order("timestamp ASC")
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
if err := query.Find(&rows).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent-network usage rows from store: %v", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent-network usage rows from store")
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
// applyAgentNetworkUsageFilters applies the shared access-log filter's
|
||||
// date/user/group/provider/model conditions to a usage-table query. Pagination,
|
||||
// sort and free-text search are ignored — the overview is an aggregate.
|
||||
func (s *SqlStore) applyAgentNetworkUsageFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
|
||||
if filter.UserID != nil {
|
||||
query = query.Where("user_id = ?", *filter.UserID)
|
||||
}
|
||||
if filter.SessionID != nil {
|
||||
query = query.Where("session_id = ?", *filter.SessionID)
|
||||
}
|
||||
if len(filter.ProviderIDs) > 0 {
|
||||
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
|
||||
}
|
||||
if len(filter.Models) > 0 {
|
||||
query = query.Where("model IN ?", filter.Models)
|
||||
}
|
||||
if len(filter.GroupIDs) > 0 {
|
||||
query = query.Where(
|
||||
"id IN (SELECT usage_id FROM agent_network_request_usage_group WHERE group_id IN ?)",
|
||||
filter.GroupIDs,
|
||||
)
|
||||
}
|
||||
if filter.StartDate != nil {
|
||||
query = query.Where("timestamp >= ?", *filter.StartDate)
|
||||
}
|
||||
if filter.EndDate != nil {
|
||||
query = query.Where("timestamp <= ?", *filter.EndDate)
|
||||
}
|
||||
return query
|
||||
}
|
||||
|
||||
// GetAgentNetworkAccessLogs retrieves flattened agent-network access logs for
|
||||
// an account with server-side pagination, filtering and sorting. Authorising
|
||||
// group ids are hydrated from the group child table for the returned page.
|
||||
func (s *SqlStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
|
||||
var logs []*agentNetworkTypes.AgentNetworkAccessLog
|
||||
var totalCount int64
|
||||
|
||||
countQuery := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
)
|
||||
if err := countQuery.Count(&totalCount).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to count agent-network access logs: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access logs")
|
||||
}
|
||||
|
||||
query := s.applyAgentNetworkAccessLogFilters(
|
||||
s.db.Where(accountIDCondition, accountID),
|
||||
filter,
|
||||
).
|
||||
Order(filter.GetSortColumn() + " " + filter.GetSortOrder()).
|
||||
Limit(filter.GetLimit()).
|
||||
Offset(filter.GetOffset())
|
||||
|
||||
if lockStrength != LockingStrengthNone {
|
||||
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
if err := query.Find(&logs).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent-network access logs from store: %v", err)
|
||||
return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access logs from store")
|
||||
}
|
||||
|
||||
if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, logs); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return logs, totalCount, nil
|
||||
}
|
||||
|
||||
// applyAgentNetworkAccessLogFilters applies the filter conditions to a query.
|
||||
func (s *SqlStore) applyAgentNetworkAccessLogFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
|
||||
if filter.Search != nil {
|
||||
p := "%" + *filter.Search + "%"
|
||||
query = query.Where(
|
||||
"id LIKE ? OR host LIKE ? OR path LIKE ? OR model LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
|
||||
p, p, p, p, p, p,
|
||||
)
|
||||
}
|
||||
if filter.UserID != nil {
|
||||
query = query.Where("user_id = ?", *filter.UserID)
|
||||
}
|
||||
if filter.SessionID != nil {
|
||||
query = query.Where("session_id = ?", *filter.SessionID)
|
||||
}
|
||||
if filter.Decision != nil {
|
||||
query = query.Where("decision = ?", *filter.Decision)
|
||||
}
|
||||
if filter.PathPrefix != nil {
|
||||
query = query.Where("path LIKE ?", *filter.PathPrefix+"%")
|
||||
}
|
||||
if len(filter.ProviderIDs) > 0 {
|
||||
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
|
||||
}
|
||||
if len(filter.Models) > 0 {
|
||||
query = query.Where("model IN ?", filter.Models)
|
||||
}
|
||||
if len(filter.GroupIDs) > 0 {
|
||||
query = query.Where(
|
||||
"id IN (SELECT log_id FROM agent_network_access_log_group WHERE group_id IN ?)",
|
||||
filter.GroupIDs,
|
||||
)
|
||||
}
|
||||
if filter.StartDate != nil {
|
||||
query = query.Where("timestamp >= ?", *filter.StartDate)
|
||||
}
|
||||
if filter.EndDate != nil {
|
||||
query = query.Where("timestamp <= ?", *filter.EndDate)
|
||||
}
|
||||
return query
|
||||
}
|
||||
|
||||
// hydrateAgentNetworkAccessLogGroups loads the authorising group ids for the
|
||||
// given page of entries and assigns them onto each entry's GroupIDs field.
|
||||
func (s *SqlStore) hydrateAgentNetworkAccessLogGroups(ctx context.Context, accountID string, logs []*agentNetworkTypes.AgentNetworkAccessLog) error {
|
||||
if len(logs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ids := make([]string, 0, len(logs))
|
||||
for _, l := range logs {
|
||||
ids = append(ids, l.ID)
|
||||
}
|
||||
|
||||
var rows []agentNetworkTypes.AgentNetworkAccessLogGroup
|
||||
if err := s.db.
|
||||
Where(accountIDCondition, accountID).
|
||||
Where("log_id IN ?", ids).
|
||||
Find(&rows).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to hydrate agent-network access log groups: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to hydrate agent-network access log groups")
|
||||
}
|
||||
|
||||
byLog := make(map[string][]string, len(logs))
|
||||
for _, r := range rows {
|
||||
byLog[r.LogID] = append(byLog[r.LogID], r.GroupID)
|
||||
}
|
||||
for _, l := range logs {
|
||||
l.GroupIDs = byLog[l.ID]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAccountAccessLogs retrieves access logs for a given account with pagination and filtering
|
||||
func (s *SqlStore) GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) {
|
||||
var logs []*accesslogs.AccessLogEntry
|
||||
|
||||
@@ -0,0 +1,623 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// GetAllAgentNetworkProviders returns Agent Network providers across
|
||||
// every account. Used by the synthesizer to build the global service map.
|
||||
func (s *SqlStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var providers []*agentNetworkTypes.Provider
|
||||
if result := tx.Find(&providers); result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get all agent network providers from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get all agent network providers from store")
|
||||
}
|
||||
|
||||
for _, provider := range providers {
|
||||
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
|
||||
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
|
||||
}
|
||||
}
|
||||
|
||||
return providers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var providers []*agentNetworkTypes.Provider
|
||||
result := tx.Find(&providers, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent network providers from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network providers from store")
|
||||
}
|
||||
|
||||
for _, provider := range providers {
|
||||
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
|
||||
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
|
||||
}
|
||||
}
|
||||
|
||||
return providers, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var provider *agentNetworkTypes.Provider
|
||||
result := tx.Take(&provider, accountAndIDQueryCondition, accountID, providerID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewAgentNetworkProviderNotFoundError(providerID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get agent network provider from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network provider from store")
|
||||
}
|
||||
|
||||
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
|
||||
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
|
||||
}
|
||||
|
||||
return provider, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error {
|
||||
providerCopy := provider.Copy()
|
||||
if err := providerCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to encrypt agent network provider %s: %v", provider.ID, err)
|
||||
return status.Errorf(status.Internal, "failed to encrypt agent network provider")
|
||||
}
|
||||
|
||||
result := s.db.Save(providerCopy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save agent network provider to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save agent network provider to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error {
|
||||
result := s.db.Delete(&agentNetworkTypes.Provider{}, accountAndIDQueryCondition, accountID, providerID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete agent network provider from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete agent network provider from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewAgentNetworkProviderNotFoundError(providerID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policies []*agentNetworkTypes.Policy
|
||||
result := tx.Find(&policies, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent network policies from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network policies from store")
|
||||
}
|
||||
|
||||
return policies, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var policy *agentNetworkTypes.Policy
|
||||
result := tx.Take(&policy, accountAndIDQueryCondition, accountID, policyID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewAgentNetworkPolicyNotFoundError(policyID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get agent network policy from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network policy from store")
|
||||
}
|
||||
|
||||
return policy, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error {
|
||||
result := s.db.Save(policy)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save agent network policy to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save agent network policy to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error {
|
||||
result := s.db.Delete(&agentNetworkTypes.Policy{}, accountAndIDQueryCondition, accountID, policyID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete agent network policy from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete agent network policy from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewAgentNetworkPolicyNotFoundError(policyID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var guardrails []*agentNetworkTypes.Guardrail
|
||||
result := tx.Find(&guardrails, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent network guardrails from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network guardrails from store")
|
||||
}
|
||||
|
||||
return guardrails, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var guardrail *agentNetworkTypes.Guardrail
|
||||
result := tx.Take(&guardrail, accountAndIDQueryCondition, accountID, guardrailID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get agent network guardrail from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network guardrail from store")
|
||||
}
|
||||
|
||||
return guardrail, nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error {
|
||||
result := s.db.Save(guardrail)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save agent network guardrail to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save agent network guardrail to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SqlStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error {
|
||||
result := s.db.Delete(&agentNetworkTypes.Guardrail{}, accountAndIDQueryCondition, accountID, guardrailID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete agent network guardrail from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete agent network guardrail from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettings returns the per-account Agent Network
|
||||
// settings row. Returns status.NotFound when no row exists.
|
||||
func (s *SqlStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var settings agentNetworkTypes.Settings
|
||||
result := tx.Take(&settings, "account_id = ?", accountID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.Errorf(status.NotFound, "agent network settings for account %s not found", accountID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get agent network settings from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network settings from store")
|
||||
}
|
||||
|
||||
return &settings, nil
|
||||
}
|
||||
|
||||
// GetAllAgentNetworkSettings returns every account's settings row. Used by the
|
||||
// access-log retention sweep to learn each account's retention window.
|
||||
func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var settings []*agentNetworkTypes.Settings
|
||||
if err := tx.Find(&settings).Error; err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to list agent network settings: %v", err)
|
||||
return nil, status.Errorf(status.Internal, "failed to list agent network settings")
|
||||
}
|
||||
return settings, nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettingsByCluster returns every Settings row pinned to
|
||||
// the given proxy cluster. Used by the bootstrap label generator to
|
||||
// build the set of subdomains already taken on a cluster.
|
||||
func (s *SqlStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var settings []*agentNetworkTypes.Settings
|
||||
result := tx.Find(&settings, "cluster = ?", cluster)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent network settings by cluster from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network settings by cluster from store")
|
||||
}
|
||||
|
||||
return settings, nil
|
||||
}
|
||||
|
||||
// SaveAgentNetworkSettings upserts the per-account Agent Network
|
||||
// settings row.
|
||||
func (s *SqlStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
|
||||
result := s.db.Save(settings)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save agent network settings to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save agent network settings to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumption atomically upserts the consumption
|
||||
// row keyed on (account, dim_kind, dim_id, window_seconds, window_start)
|
||||
// and adds the supplied deltas. Concurrent calls from multiple proxy
|
||||
// nodes converge — the database performs the increment server-side via
|
||||
// ON CONFLICT DO UPDATE so no read-modify-write race exists.
|
||||
func (s *SqlStore) IncrementAgentNetworkConsumption(
|
||||
ctx context.Context,
|
||||
accountID string,
|
||||
kind agentNetworkTypes.ConsumptionDimension,
|
||||
dimID string,
|
||||
windowSeconds int64,
|
||||
windowStart time.Time,
|
||||
tokensIn, tokensOut int64,
|
||||
costUSD float64,
|
||||
) error {
|
||||
if accountID == "" || dimID == "" || windowSeconds <= 0 {
|
||||
return status.Errorf(status.InvalidArgument, "account_id, dim_id and window_seconds must be set")
|
||||
}
|
||||
// Deltas are added server-side via ON CONFLICT; a negative or non-finite
|
||||
// value would silently decrement / poison the persisted totals.
|
||||
if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
|
||||
return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
|
||||
}
|
||||
row := agentNetworkTypes.Consumption{
|
||||
AccountID: accountID,
|
||||
DimensionKind: kind,
|
||||
DimensionID: dimID,
|
||||
WindowSeconds: windowSeconds,
|
||||
WindowStartUTC: windowStart.UTC(),
|
||||
TokensInput: tokensIn,
|
||||
TokensOutput: tokensOut,
|
||||
CostUSD: costUSD,
|
||||
UpdatedAt: time.Now().UTC(),
|
||||
}
|
||||
const tbl = "agent_network_consumption"
|
||||
err := s.db.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "account_id"},
|
||||
{Name: "dim_kind"},
|
||||
{Name: "dim_id"},
|
||||
{Name: "window_seconds"},
|
||||
{Name: "window_start_utc"},
|
||||
},
|
||||
DoUpdates: clause.Assignments(map[string]any{
|
||||
"tokens_input": gorm.Expr(tbl+".tokens_input + ?", tokensIn),
|
||||
"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
|
||||
"cost_usd": gorm.Expr(tbl+".cost_usd + ?", costUSD),
|
||||
"updated_at": time.Now().UTC(),
|
||||
}),
|
||||
}).Create(&row).Error
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to increment agent network consumption: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to increment agent network consumption")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumption returns the consumption row for the exact
|
||||
// window key. Returns a zero-valued row (not found mapped to zero) so
|
||||
// callers can use the result as the headroom basis without nil checks.
|
||||
func (s *SqlStore) GetAgentNetworkConsumption(
|
||||
ctx context.Context,
|
||||
lockStrength LockingStrength,
|
||||
accountID string,
|
||||
kind agentNetworkTypes.ConsumptionDimension,
|
||||
dimID string,
|
||||
windowSeconds int64,
|
||||
windowStart time.Time,
|
||||
) (*agentNetworkTypes.Consumption, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
var row agentNetworkTypes.Consumption
|
||||
result := tx.Take(&row,
|
||||
"account_id = ? AND dim_kind = ? AND dim_id = ? AND window_seconds = ? AND window_start_utc = ?",
|
||||
accountID, kind, dimID, windowSeconds, windowStart.UTC())
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return &agentNetworkTypes.Consumption{
|
||||
AccountID: accountID,
|
||||
DimensionKind: kind,
|
||||
DimensionID: dimID,
|
||||
WindowSeconds: windowSeconds,
|
||||
WindowStartUTC: windowStart.UTC(),
|
||||
}, nil
|
||||
}
|
||||
log.WithContext(ctx).Errorf("failed to get agent network consumption: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
|
||||
}
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumptionBatch reads many consumption counters for one
|
||||
// account in a single query, returning a map keyed by the exact
|
||||
// ConsumptionKey. Missing counters are simply absent from the map (callers
|
||||
// treat absence as a zero counter). Replaces the per-cap point reads the
|
||||
// policy selector previously issued one at a time.
|
||||
func (s *SqlStore) GetAgentNetworkConsumptionBatch(
|
||||
ctx context.Context,
|
||||
lockStrength LockingStrength,
|
||||
accountID string,
|
||||
keys []agentNetworkTypes.ConsumptionKey,
|
||||
) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) {
|
||||
out := make(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, len(keys))
|
||||
if len(keys) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Collect the distinct dim ids, windows and window starts so a single
|
||||
// query scopes to exactly the current windows in play, then filter the
|
||||
// returned rows down to the exact requested keys.
|
||||
wanted := make(map[agentNetworkTypes.ConsumptionKey]struct{}, len(keys))
|
||||
dimSet := make(map[string]struct{})
|
||||
winSet := make(map[int64]struct{})
|
||||
startSet := make(map[time.Time]struct{})
|
||||
for _, k := range keys {
|
||||
k.WindowStartUTC = k.WindowStartUTC.UTC()
|
||||
wanted[k] = struct{}{}
|
||||
dimSet[k.DimID] = struct{}{}
|
||||
winSet[k.WindowSeconds] = struct{}{}
|
||||
startSet[k.WindowStartUTC] = struct{}{}
|
||||
}
|
||||
dimIDs := make([]string, 0, len(dimSet))
|
||||
for d := range dimSet {
|
||||
dimIDs = append(dimIDs, d)
|
||||
}
|
||||
windows := make([]int64, 0, len(winSet))
|
||||
for w := range winSet {
|
||||
windows = append(windows, w)
|
||||
}
|
||||
starts := make([]time.Time, 0, len(startSet))
|
||||
for t := range startSet {
|
||||
starts = append(starts, t)
|
||||
}
|
||||
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
var rows []*agentNetworkTypes.Consumption
|
||||
result := tx.Find(&rows,
|
||||
"account_id = ? AND dim_id IN ? AND window_seconds IN ? AND window_start_utc IN ?",
|
||||
accountID, dimIDs, windows, starts)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to batch-get agent network consumption: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
|
||||
}
|
||||
for _, row := range rows {
|
||||
k := agentNetworkTypes.ConsumptionKey{
|
||||
Kind: row.DimensionKind,
|
||||
DimID: row.DimensionID,
|
||||
WindowSeconds: row.WindowSeconds,
|
||||
WindowStartUTC: row.WindowStartUTC.UTC(),
|
||||
}
|
||||
if _, ok := wanted[k]; ok {
|
||||
out[k] = row
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumptionBatch applies the same usage delta to every
|
||||
// supplied counter inside a single transaction, so all per-(dimension, window)
|
||||
// counters a served request books are written atomically in one round-trip
|
||||
// instead of one upsert per counter. Keys are deduplicated by the caller.
|
||||
func (s *SqlStore) IncrementAgentNetworkConsumptionBatch(
|
||||
ctx context.Context,
|
||||
accountID string,
|
||||
keys []agentNetworkTypes.ConsumptionKey,
|
||||
tokensIn, tokensOut int64,
|
||||
costUSD float64,
|
||||
) error {
|
||||
if accountID == "" {
|
||||
return status.Errorf(status.InvalidArgument, "account_id must be set")
|
||||
}
|
||||
if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
|
||||
return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
const tbl = "agent_network_consumption"
|
||||
err := s.db.Transaction(func(tx *gorm.DB) error {
|
||||
for _, k := range keys {
|
||||
if k.DimID == "" || k.WindowSeconds <= 0 {
|
||||
return status.Errorf(status.InvalidArgument, "dim_id and window_seconds must be set")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
row := agentNetworkTypes.Consumption{
|
||||
AccountID: accountID,
|
||||
DimensionKind: k.Kind,
|
||||
DimensionID: k.DimID,
|
||||
WindowSeconds: k.WindowSeconds,
|
||||
WindowStartUTC: k.WindowStartUTC.UTC(),
|
||||
TokensInput: tokensIn,
|
||||
TokensOutput: tokensOut,
|
||||
CostUSD: costUSD,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := tx.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "account_id"},
|
||||
{Name: "dim_kind"},
|
||||
{Name: "dim_id"},
|
||||
{Name: "window_seconds"},
|
||||
{Name: "window_start_utc"},
|
||||
},
|
||||
DoUpdates: clause.Assignments(map[string]any{
|
||||
"tokens_input": gorm.Expr(tbl+".tokens_input + ?", tokensIn),
|
||||
"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
|
||||
"cost_usd": gorm.Expr(tbl+".cost_usd + ?", costUSD),
|
||||
"updated_at": now,
|
||||
}),
|
||||
}).Create(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to batch-increment agent network consumption: %v", err)
|
||||
return status.Errorf(status.Internal, "failed to increment agent network consumption")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListAgentNetworkConsumption returns every consumption row recorded
|
||||
// for the account, ordered by window_start descending. Backs the
|
||||
// dashboard's basic counter view.
|
||||
func (s *SqlStore) ListAgentNetworkConsumption(
|
||||
ctx context.Context,
|
||||
lockStrength LockingStrength,
|
||||
accountID string,
|
||||
) ([]*agentNetworkTypes.Consumption, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
var rows []*agentNetworkTypes.Consumption
|
||||
result := tx.
|
||||
Order("window_start_utc DESC").
|
||||
Find(&rows, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to list agent network consumption: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to list agent network consumption")
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkBudgetRules returns every account-level budget rule for
|
||||
// the account.
|
||||
func (s *SqlStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var rules []*agentNetworkTypes.AccountBudgetRule
|
||||
result := tx.Find(&rules, accountIDCondition, accountID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get agent network budget rules from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network budget rules from store")
|
||||
}
|
||||
|
||||
return rules, nil
|
||||
}
|
||||
|
||||
// GetAgentNetworkBudgetRuleByID returns a single budget rule scoped to the
|
||||
// account, or a NotFound error.
|
||||
func (s *SqlStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) {
|
||||
tx := s.db
|
||||
if lockStrength != LockingStrengthNone {
|
||||
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
|
||||
}
|
||||
|
||||
var rule *agentNetworkTypes.AccountBudgetRule
|
||||
result := tx.Take(&rule, accountAndIDQueryCondition, accountID, ruleID)
|
||||
if result.Error != nil {
|
||||
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
|
||||
return nil, status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Errorf("failed to get agent network budget rule from store: %v", result.Error)
|
||||
return nil, status.Errorf(status.Internal, "failed to get agent network budget rule from store")
|
||||
}
|
||||
|
||||
return rule, nil
|
||||
}
|
||||
|
||||
// SaveAgentNetworkBudgetRule upserts a budget rule.
|
||||
func (s *SqlStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error {
|
||||
result := s.db.Save(rule)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to save agent network budget rule to store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to save agent network budget rule to store")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkBudgetRule removes a budget rule scoped to the account.
|
||||
func (s *SqlStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error {
|
||||
result := s.db.Delete(&agentNetworkTypes.AccountBudgetRule{}, accountAndIDQueryCondition, accountID, ruleID)
|
||||
if result.Error != nil {
|
||||
log.WithContext(ctx).Errorf("failed to delete agent network budget rule from store: %v", result.Error)
|
||||
return status.Errorf(status.Internal, "failed to delete agent network budget rule from store")
|
||||
}
|
||||
|
||||
if result.RowsAffected == 0 {
|
||||
return status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// TestAgentNetworkUsage_RealStore_RoundTrip drives CreateAgentNetworkUsage and
|
||||
// CreateAgentNetworkAccessLog through a real sqlite store to prove the schema
|
||||
// migrates and the inserts succeed for both a populated (allowed) entry and a
|
||||
// stripped (denied) entry.
|
||||
func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err, "real sqlite test store must come up")
|
||||
defer cleanup()
|
||||
|
||||
const accountID = "acc-anet-usage-1"
|
||||
now := time.Now().UTC()
|
||||
|
||||
// Populated (allowed) usage row with two authorising groups.
|
||||
usage := &agentNetworkTypes.AgentNetworkUsage{
|
||||
ID: "log-allowed-1",
|
||||
AccountID: accountID,
|
||||
Timestamp: now,
|
||||
UserID: "user-alice",
|
||||
ResolvedProviderID: "prov-openai-1",
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o",
|
||||
SessionID: "sess-round-trip-1",
|
||||
InputTokens: 1200,
|
||||
OutputTokens: 640,
|
||||
TotalTokens: 1840,
|
||||
CostUSD: 0.0231,
|
||||
}
|
||||
usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
|
||||
{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
|
||||
{UsageID: usage.ID, GroupID: "grp-oncall", AccountID: accountID},
|
||||
}
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "populated usage insert must succeed")
|
||||
|
||||
// Stripped (denied / 403) usage row: no provider/model/tokens, no groups.
|
||||
denied := &agentNetworkTypes.AgentNetworkUsage{
|
||||
ID: "log-denied-1",
|
||||
AccountID: accountID,
|
||||
Timestamp: now,
|
||||
UserID: "user-bob",
|
||||
}
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, denied, nil), "stripped usage insert must succeed")
|
||||
|
||||
// Idempotency: re-inserting the same id must not error.
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "duplicate usage insert must be idempotent")
|
||||
|
||||
// Access-log row + group children.
|
||||
entry := &agentNetworkTypes.AgentNetworkAccessLog{
|
||||
ID: "log-allowed-1",
|
||||
AccountID: accountID,
|
||||
ServiceID: "agent-net-svc-1",
|
||||
Timestamp: now,
|
||||
UserID: "user-alice",
|
||||
StatusCode: 200,
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o",
|
||||
SessionID: "sess-round-trip-1",
|
||||
InputTokens: 1200,
|
||||
OutputTokens: 640,
|
||||
TotalTokens: 1840,
|
||||
CostUSD: 0.0231,
|
||||
}
|
||||
entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
|
||||
{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
|
||||
{LogID: entry.ID, GroupID: "grp-oncall", AccountID: accountID},
|
||||
}
|
||||
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, entry, entryGroups), "access-log insert must succeed")
|
||||
|
||||
// Read back through the filtered list + verify group hydration.
|
||||
logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
|
||||
require.NoError(t, err, "list must succeed")
|
||||
assert.Equal(t, int64(1), total, "one access-log row expected")
|
||||
require.Len(t, logs, 1)
|
||||
assert.ElementsMatch(t, []string{"grp-eng", "grp-oncall"}, logs[0].GroupIDs, "group ids must hydrate")
|
||||
assert.Equal(t, "sess-round-trip-1", logs[0].SessionID, "session id must persist and read back on the access-log row")
|
||||
|
||||
// Session filter narrows the access-log listing to one conversation.
|
||||
sessionID := "sess-round-trip-1"
|
||||
sessLogs, sessTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
|
||||
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &sessionID})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), sessTotal, "session filter must match the one row with that session id")
|
||||
require.Len(t, sessLogs, 1)
|
||||
assert.Equal(t, entry.ID, sessLogs[0].ID, "session filter must return the matching log row")
|
||||
|
||||
bogus := "no-such-session"
|
||||
_, emptyTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
|
||||
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &bogus})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), emptyTotal, "unknown session id must match nothing")
|
||||
|
||||
// Session filter also narrows the always-on usage rows.
|
||||
sessUsage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID,
|
||||
agentNetworkTypes.AgentNetworkAccessLogFilter{SessionID: &sessionID})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, sessUsage, 1, "session filter must narrow usage rows to the matching session")
|
||||
assert.Equal(t, "sess-round-trip-1", sessUsage[0].SessionID, "usage row must carry the session id")
|
||||
}
|
||||
|
||||
// TestAgentNetworkUsageOverview_DailyAggregation drives GetAgentNetworkUsageRows
|
||||
// + AggregateUsageByGranularity end-to-end against a real sqlite store, with
|
||||
// two rows on the same day and one on another, plus a model filter.
|
||||
func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err, "real sqlite test store must come up")
|
||||
defer cleanup()
|
||||
|
||||
const accountID = "acc-anet-overview-1"
|
||||
day1 := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
|
||||
day1b := time.Date(2026, 5, 5, 22, 0, 0, 0, time.UTC)
|
||||
day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
|
||||
|
||||
mk := func(id string, ts time.Time, model string, in, out int64, cost float64) *agentNetworkTypes.AgentNetworkUsage {
|
||||
return &agentNetworkTypes.AgentNetworkUsage{
|
||||
ID: id, AccountID: accountID, Timestamp: ts, Model: model,
|
||||
InputTokens: in, OutputTokens: out, TotalTokens: in + out, CostUSD: cost,
|
||||
}
|
||||
}
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u2", day1b, "gpt-4o", 200, 80, 0.20), nil))
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u3", day2, "claude-3", 10, 5, 0.01), nil))
|
||||
|
||||
rows, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rows, 3, "all three usage rows expected")
|
||||
|
||||
buckets := agentNetworkTypes.AggregateUsageByGranularity(rows, agentNetworkTypes.UsageGranularityDay)
|
||||
require.Len(t, buckets, 2, "two distinct days expected")
|
||||
assert.Equal(t, "2026-05-05", buckets[0].PeriodStart, "oldest-first ordering")
|
||||
assert.Equal(t, int64(300), buckets[0].InputTokens, "same-day input tokens summed")
|
||||
assert.Equal(t, int64(130), buckets[0].OutputTokens)
|
||||
assert.InDelta(t, 0.30, buckets[0].CostUSD, 1e-9, "same-day cost summed")
|
||||
assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
|
||||
assert.Equal(t, int64(15), buckets[1].TotalTokens)
|
||||
|
||||
// Model filter narrows to a single day.
|
||||
model := "claude-3"
|
||||
filtered, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Models: []string{model}})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, filtered, 1, "model filter must narrow rows")
|
||||
assert.Equal(t, "u3", filtered[0].ID)
|
||||
}
|
||||
|
||||
// TestDeleteOldAgentNetworkAccessLogs verifies the retention sweep removes only
|
||||
// access-log rows (and their group children) older than the cutoff, leaving
|
||||
// recent rows — and never touching usage records.
|
||||
func TestDeleteOldAgentNetworkAccessLogs(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err, "real sqlite test store must come up")
|
||||
defer cleanup()
|
||||
|
||||
const accountID = "acc-anet-retention-1"
|
||||
old := time.Now().UTC().AddDate(0, 0, -40)
|
||||
recent := time.Now().UTC().AddDate(0, 0, -1)
|
||||
|
||||
mkLog := func(id string, ts time.Time) (*agentNetworkTypes.AgentNetworkAccessLog, []agentNetworkTypes.AgentNetworkAccessLogGroup) {
|
||||
return &agentNetworkTypes.AgentNetworkAccessLog{
|
||||
ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts, StatusCode: 200, Model: "gpt-4o",
|
||||
}, []agentNetworkTypes.AgentNetworkAccessLogGroup{
|
||||
{LogID: id, GroupID: "grp-eng", AccountID: accountID},
|
||||
}
|
||||
}
|
||||
oldEntry, oldGroups := mkLog("old-1", old)
|
||||
recentEntry, recentGroups := mkLog("recent-1", recent)
|
||||
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, oldEntry, oldGroups))
|
||||
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, recentEntry, recentGroups))
|
||||
// A usage row for the old request must survive the access-log sweep.
|
||||
require.NoError(t, s.CreateAgentNetworkUsage(ctx, &agentNetworkTypes.AgentNetworkUsage{
|
||||
ID: "old-1", AccountID: accountID, Timestamp: old, Model: "gpt-4o", InputTokens: 10, TotalTokens: 10,
|
||||
}, nil))
|
||||
|
||||
cutoff := time.Now().UTC().AddDate(0, 0, -30)
|
||||
deleted, err := s.DeleteOldAgentNetworkAccessLogs(ctx, accountID, cutoff)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), deleted, "only the 40-day-old log is deleted")
|
||||
|
||||
logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), total, "the recent log remains")
|
||||
require.Len(t, logs, 1)
|
||||
assert.Equal(t, "recent-1", logs[0].ID)
|
||||
|
||||
// Usage is untouched by the access-log retention sweep.
|
||||
usage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, usage, 1, "usage record for the deleted log must survive")
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// TestAgentNetworkBudgetRule_RealStore_RoundTrip is the GC-0 no-mock guard: it
|
||||
// drives the budget-rule CRUD through a real sqlite store and asserts the full
|
||||
// object — targets and the reused PolicyLimits cap shape — survives the
|
||||
// save → gorm/JSON serialize → reload round-trip, then that delete removes it
|
||||
// and a second delete reports NotFound.
|
||||
func TestAgentNetworkBudgetRule_RealStore_RoundTrip(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err, "real sqlite test store must come up")
|
||||
defer cleanup()
|
||||
|
||||
const accountID = "acc-budgetrule-1"
|
||||
rule := agentNetworkTypes.NewAccountBudgetRule(accountID)
|
||||
rule.Name = "eng-monthly"
|
||||
rule.TargetGroups = []string{"grp-eng", "grp-oncall"}
|
||||
rule.TargetUsers = []string{"user-alice"}
|
||||
rule.Limits = agentNetworkTypes.PolicyLimits{
|
||||
TokenLimit: agentNetworkTypes.PolicyTokenLimit{
|
||||
Enabled: true, GroupCap: 100_000, UserCap: 10_000, WindowSeconds: 2_592_000,
|
||||
},
|
||||
BudgetLimit: agentNetworkTypes.PolicyBudgetLimit{
|
||||
Enabled: true, GroupCapUsd: 500, UserCapUsd: 50, WindowSeconds: 2_592_000,
|
||||
},
|
||||
}
|
||||
require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, rule), "save must succeed")
|
||||
|
||||
got, err := s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
|
||||
require.NoError(t, err, "get by id must succeed after save")
|
||||
assert.Equal(t, rule.Name, got.Name, "name must round-trip")
|
||||
assert.Equal(t, []string{"grp-eng", "grp-oncall"}, got.TargetGroups, "target groups must round-trip")
|
||||
assert.Equal(t, []string{"user-alice"}, got.TargetUsers, "target users must round-trip")
|
||||
assert.Equal(t, rule.Limits, got.Limits, "the reused PolicyLimits cap shape must round-trip intact")
|
||||
assert.True(t, got.Enabled, "enabled must round-trip")
|
||||
|
||||
list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, accountID)
|
||||
require.NoError(t, err, "list must succeed")
|
||||
require.Len(t, list, 1, "exactly the one saved rule must be listed")
|
||||
assert.Equal(t, rule.ID, list[0].ID, "listed rule id must match")
|
||||
|
||||
require.NoError(t, s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID), "delete must succeed")
|
||||
|
||||
_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
|
||||
assert.Error(t, err, "get after delete must report not found")
|
||||
|
||||
err = s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID)
|
||||
assert.Error(t, err, "deleting an absent rule must report not found")
|
||||
}
|
||||
|
||||
// TestAgentNetworkBudgetRule_RealStore_ScopedByAccount pins that rules are
|
||||
// account-scoped: a rule under one account is invisible to another.
|
||||
func TestAgentNetworkBudgetRule_RealStore_ScopedByAccount(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err)
|
||||
defer cleanup()
|
||||
|
||||
ruleA := agentNetworkTypes.NewAccountBudgetRule("acc-A")
|
||||
require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, ruleA))
|
||||
|
||||
list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, "acc-B")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, list, "account B must not see account A's budget rule")
|
||||
|
||||
_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, "acc-B", ruleA.ID)
|
||||
assert.Error(t, err, "cross-account get by id must not resolve")
|
||||
}
|
||||
|
||||
// TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip pins the GC-0
|
||||
// additive settings columns: the three collection toggles default off on a
|
||||
// fresh row and survive a save/reload at their set values.
|
||||
func TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
|
||||
require.NoError(t, err)
|
||||
defer cleanup()
|
||||
|
||||
const accountID = "acc-settings-toggles"
|
||||
require.NoError(t, s.SaveAgentNetworkSettings(ctx, &agentNetworkTypes.Settings{
|
||||
AccountID: accountID,
|
||||
Cluster: "eu.proxy.netbird.io",
|
||||
Subdomain: "violet",
|
||||
}))
|
||||
|
||||
got, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, got.EnableLogCollection, "log collection must default off")
|
||||
assert.False(t, got.EnablePromptCollection, "prompt collection must default off")
|
||||
assert.False(t, got.RedactPii, "redact pii must default off")
|
||||
|
||||
got.EnableLogCollection = true
|
||||
got.EnablePromptCollection = true
|
||||
got.RedactPii = true
|
||||
require.NoError(t, s.SaveAgentNetworkSettings(ctx, got))
|
||||
|
||||
reloaded, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, reloaded.EnableLogCollection, "log collection must round-trip on")
|
||||
assert.True(t, reloaded.EnablePromptCollection, "prompt collection must round-trip on")
|
||||
assert.True(t, reloaded.RedactPii, "redact pii must round-trip on")
|
||||
}
|
||||
@@ -37,6 +37,7 @@ import (
|
||||
"github.com/netbirdio/netbird/util"
|
||||
"github.com/netbirdio/netbird/util/crypt"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/migration"
|
||||
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
|
||||
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
|
||||
@@ -300,6 +301,11 @@ type Store interface {
|
||||
CreateAccessLog(ctx context.Context, log *accesslogs.AccessLogEntry) error
|
||||
GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error)
|
||||
DeleteOldAccessLogs(ctx context.Context, olderThan time.Time) (int64, error)
|
||||
CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error
|
||||
CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error
|
||||
GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error)
|
||||
GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error)
|
||||
DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error)
|
||||
GetServiceTargetByTargetID(ctx context.Context, lockStrength LockingStrength, accountID string, targetID string) (*rpservice.Target, error)
|
||||
GetTargetsByServiceID(ctx context.Context, lockStrength LockingStrength, accountID string, serviceID string) ([]*rpservice.Target, error)
|
||||
DeleteTarget(ctx context.Context, accountID string, serviceID string, targetID uint) error
|
||||
@@ -329,6 +335,34 @@ type Store interface {
|
||||
GetProxyMetrics(ctx context.Context) (ProxyMetrics, error)
|
||||
|
||||
GetRoutingPeerNetworks(ctx context.Context, accountID, peerID string) ([]string, error)
|
||||
|
||||
// Agent Network persistence (providers, policies, guardrails, settings).
|
||||
GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error)
|
||||
GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error)
|
||||
GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error)
|
||||
SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error
|
||||
DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error
|
||||
GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error)
|
||||
GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error)
|
||||
SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error
|
||||
DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error
|
||||
GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error)
|
||||
GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error)
|
||||
SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error
|
||||
DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error
|
||||
GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error)
|
||||
GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error)
|
||||
GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error)
|
||||
SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error
|
||||
IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error
|
||||
IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error
|
||||
GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error)
|
||||
GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error)
|
||||
ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error)
|
||||
GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error)
|
||||
GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error)
|
||||
SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error
|
||||
DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error
|
||||
}
|
||||
|
||||
// ProxyMetrics aggregates self-hosted proxy + cluster usage signals
|
||||
|
||||
@@ -0,0 +1,464 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
context "context"
|
||||
reflect "reflect"
|
||||
time "time"
|
||||
|
||||
gomock "github.com/golang/mock/gomock"
|
||||
|
||||
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// GetAllAgentNetworkProviders mocks base method.
|
||||
func (m *MockStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAllAgentNetworkProviders", ctx, lockStrength)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Provider)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAllAgentNetworkProviders indicates an expected call of GetAllAgentNetworkProviders.
|
||||
func (mr *MockStoreMockRecorder) GetAllAgentNetworkProviders(ctx, lockStrength interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkProviders), ctx, lockStrength)
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkProviders mocks base method.
|
||||
func (m *MockStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAccountAgentNetworkProviders", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Provider)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkProviders indicates an expected call of GetAccountAgentNetworkProviders.
|
||||
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkProviders(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkProviders), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkProviderByID mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkProviderByID", ctx, lockStrength, accountID, providerID)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.Provider)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkProviderByID indicates an expected call of GetAgentNetworkProviderByID.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkProviderByID(ctx, lockStrength, accountID, providerID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkProviderByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkProviderByID), ctx, lockStrength, accountID, providerID)
|
||||
}
|
||||
|
||||
// SaveAgentNetworkProvider mocks base method.
|
||||
func (m *MockStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "SaveAgentNetworkProvider", ctx, provider)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// SaveAgentNetworkProvider indicates an expected call of SaveAgentNetworkProvider.
|
||||
func (mr *MockStoreMockRecorder) SaveAgentNetworkProvider(ctx, provider interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkProvider), ctx, provider)
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkProvider mocks base method.
|
||||
func (m *MockStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAgentNetworkProvider", ctx, accountID, providerID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkProvider indicates an expected call of DeleteAgentNetworkProvider.
|
||||
func (mr *MockStoreMockRecorder) DeleteAgentNetworkProvider(ctx, accountID, providerID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkProvider), ctx, accountID, providerID)
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkPolicies mocks base method.
|
||||
func (m *MockStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAccountAgentNetworkPolicies", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Policy)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkPolicies indicates an expected call of GetAccountAgentNetworkPolicies.
|
||||
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkPolicies(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkPolicies", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkPolicies), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkPolicyByID mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkPolicyByID", ctx, lockStrength, accountID, policyID)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.Policy)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkPolicyByID indicates an expected call of GetAgentNetworkPolicyByID.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkPolicyByID(ctx, lockStrength, accountID, policyID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkPolicyByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkPolicyByID), ctx, lockStrength, accountID, policyID)
|
||||
}
|
||||
|
||||
// SaveAgentNetworkPolicy mocks base method.
|
||||
func (m *MockStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "SaveAgentNetworkPolicy", ctx, policy)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// SaveAgentNetworkPolicy indicates an expected call of SaveAgentNetworkPolicy.
|
||||
func (mr *MockStoreMockRecorder) SaveAgentNetworkPolicy(ctx, policy interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkPolicy), ctx, policy)
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkPolicy mocks base method.
|
||||
func (m *MockStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAgentNetworkPolicy", ctx, accountID, policyID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkPolicy indicates an expected call of DeleteAgentNetworkPolicy.
|
||||
func (mr *MockStoreMockRecorder) DeleteAgentNetworkPolicy(ctx, accountID, policyID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkPolicy), ctx, accountID, policyID)
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkGuardrails mocks base method.
|
||||
func (m *MockStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAccountAgentNetworkGuardrails", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Guardrail)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkGuardrails indicates an expected call of GetAccountAgentNetworkGuardrails.
|
||||
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkGuardrails(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkGuardrails", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkGuardrails), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkGuardrailByID mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkGuardrailByID", ctx, lockStrength, accountID, guardrailID)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.Guardrail)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkGuardrailByID indicates an expected call of GetAgentNetworkGuardrailByID.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkGuardrailByID(ctx, lockStrength, accountID, guardrailID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkGuardrailByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkGuardrailByID), ctx, lockStrength, accountID, guardrailID)
|
||||
}
|
||||
|
||||
// SaveAgentNetworkGuardrail mocks base method.
|
||||
func (m *MockStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "SaveAgentNetworkGuardrail", ctx, guardrail)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// SaveAgentNetworkGuardrail indicates an expected call of SaveAgentNetworkGuardrail.
|
||||
func (mr *MockStoreMockRecorder) SaveAgentNetworkGuardrail(ctx, guardrail interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkGuardrail), ctx, guardrail)
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkGuardrail mocks base method.
|
||||
func (m *MockStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAgentNetworkGuardrail", ctx, accountID, guardrailID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkGuardrail indicates an expected call of DeleteAgentNetworkGuardrail.
|
||||
func (mr *MockStoreMockRecorder) DeleteAgentNetworkGuardrail(ctx, accountID, guardrailID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkGuardrail), ctx, accountID, guardrailID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettings mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkSettings", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.Settings)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettings indicates an expected call of GetAgentNetworkSettings.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkSettings(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettings), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettingsByCluster mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkSettingsByCluster", ctx, lockStrength, cluster)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Settings)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkSettingsByCluster indicates an expected call of GetAgentNetworkSettingsByCluster.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkSettingsByCluster(ctx, lockStrength, cluster interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettingsByCluster", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettingsByCluster), ctx, lockStrength, cluster)
|
||||
}
|
||||
|
||||
// SaveAgentNetworkSettings mocks base method.
|
||||
func (m *MockStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "SaveAgentNetworkSettings", ctx, settings)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// SaveAgentNetworkSettings indicates an expected call of SaveAgentNetworkSettings.
|
||||
func (mr *MockStoreMockRecorder) SaveAgentNetworkSettings(ctx, settings interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkSettings), ctx, settings)
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumption mocks base method.
|
||||
func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumption", ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumption indicates an expected call of IncrementAgentNetworkConsumption.
|
||||
func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumption(ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumption), ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD)
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumption mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkConsumption", ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.Consumption)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumption indicates an expected call of GetAgentNetworkConsumption.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkConsumption(ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumption), ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart)
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumptionBatch mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkConsumptionBatch", ctx, lockStrength, accountID, keys)
|
||||
ret0, _ := ret[0].(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkConsumptionBatch indicates an expected call of GetAgentNetworkConsumptionBatch.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkConsumptionBatch(ctx, lockStrength, accountID, keys interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumptionBatch), ctx, lockStrength, accountID, keys)
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumptionBatch mocks base method.
|
||||
func (m *MockStore) IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumptionBatch", ctx, accountID, keys, tokensIn, tokensOut, costUSD)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// IncrementAgentNetworkConsumptionBatch indicates an expected call of IncrementAgentNetworkConsumptionBatch.
|
||||
func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumptionBatch(ctx, accountID, keys, tokensIn, tokensOut, costUSD interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumptionBatch), ctx, accountID, keys, tokensIn, tokensOut, costUSD)
|
||||
}
|
||||
|
||||
// ListAgentNetworkConsumption mocks base method.
|
||||
func (m *MockStore) ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ListAgentNetworkConsumption", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Consumption)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ListAgentNetworkConsumption indicates an expected call of ListAgentNetworkConsumption.
|
||||
func (mr *MockStoreMockRecorder) ListAgentNetworkConsumption(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).ListAgentNetworkConsumption), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkBudgetRules mocks base method.
|
||||
func (m *MockStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAccountAgentNetworkBudgetRules", ctx, lockStrength, accountID)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.AccountBudgetRule)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAccountAgentNetworkBudgetRules indicates an expected call of GetAccountAgentNetworkBudgetRules.
|
||||
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkBudgetRules(ctx, lockStrength, accountID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkBudgetRules", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkBudgetRules), ctx, lockStrength, accountID)
|
||||
}
|
||||
|
||||
// GetAgentNetworkBudgetRuleByID mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkBudgetRuleByID", ctx, lockStrength, accountID, ruleID)
|
||||
ret0, _ := ret[0].(*agentNetworkTypes.AccountBudgetRule)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkBudgetRuleByID indicates an expected call of GetAgentNetworkBudgetRuleByID.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkBudgetRuleByID(ctx, lockStrength, accountID, ruleID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkBudgetRuleByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkBudgetRuleByID), ctx, lockStrength, accountID, ruleID)
|
||||
}
|
||||
|
||||
// SaveAgentNetworkBudgetRule mocks base method.
|
||||
func (m *MockStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "SaveAgentNetworkBudgetRule", ctx, rule)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// SaveAgentNetworkBudgetRule indicates an expected call of SaveAgentNetworkBudgetRule.
|
||||
func (mr *MockStoreMockRecorder) SaveAgentNetworkBudgetRule(ctx, rule interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkBudgetRule), ctx, rule)
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkBudgetRule mocks base method.
|
||||
func (m *MockStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteAgentNetworkBudgetRule", ctx, accountID, ruleID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteAgentNetworkBudgetRule indicates an expected call of DeleteAgentNetworkBudgetRule.
|
||||
func (mr *MockStoreMockRecorder) DeleteAgentNetworkBudgetRule(ctx, accountID, ruleID interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkBudgetRule), ctx, accountID, ruleID)
|
||||
}
|
||||
|
||||
// CreateAgentNetworkAccessLog mocks base method.
|
||||
func (m *MockStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "CreateAgentNetworkAccessLog", ctx, entry, groups)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// CreateAgentNetworkAccessLog indicates an expected call of CreateAgentNetworkAccessLog.
|
||||
func (mr *MockStoreMockRecorder) CreateAgentNetworkAccessLog(ctx, entry, groups interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkAccessLog", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkAccessLog), ctx, entry, groups)
|
||||
}
|
||||
|
||||
// CreateAgentNetworkUsage mocks base method.
|
||||
func (m *MockStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "CreateAgentNetworkUsage", ctx, usage, groups)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// CreateAgentNetworkUsage indicates an expected call of CreateAgentNetworkUsage.
|
||||
func (mr *MockStoreMockRecorder) CreateAgentNetworkUsage(ctx, usage, groups interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkUsage", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkUsage), ctx, usage, groups)
|
||||
}
|
||||
|
||||
// GetAgentNetworkAccessLogs mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkAccessLogs", ctx, lockStrength, accountID, filter)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkAccessLog)
|
||||
ret1, _ := ret[1].(int64)
|
||||
ret2, _ := ret[2].(error)
|
||||
return ret0, ret1, ret2
|
||||
}
|
||||
|
||||
// GetAgentNetworkAccessLogs indicates an expected call of GetAgentNetworkAccessLogs.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkAccessLogs(ctx, lockStrength, accountID, filter interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkAccessLogs), ctx, lockStrength, accountID, filter)
|
||||
}
|
||||
|
||||
// GetAgentNetworkUsageRows mocks base method.
|
||||
func (m *MockStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAgentNetworkUsageRows", ctx, lockStrength, accountID, filter)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkUsage)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAgentNetworkUsageRows indicates an expected call of GetAgentNetworkUsageRows.
|
||||
func (mr *MockStoreMockRecorder) GetAgentNetworkUsageRows(ctx, lockStrength, accountID, filter interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkUsageRows", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkUsageRows), ctx, lockStrength, accountID, filter)
|
||||
}
|
||||
|
||||
// DeleteOldAgentNetworkAccessLogs mocks base method.
|
||||
func (m *MockStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteOldAgentNetworkAccessLogs", ctx, accountID, olderThan)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DeleteOldAgentNetworkAccessLogs indicates an expected call of DeleteOldAgentNetworkAccessLogs.
|
||||
func (mr *MockStoreMockRecorder) DeleteOldAgentNetworkAccessLogs(ctx, accountID, olderThan interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).DeleteOldAgentNetworkAccessLogs), ctx, accountID, olderThan)
|
||||
}
|
||||
|
||||
// GetAllAgentNetworkSettings mocks base method.
|
||||
func (m *MockStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAllAgentNetworkSettings", ctx, lockStrength)
|
||||
ret0, _ := ret[0].([]*agentNetworkTypes.Settings)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAllAgentNetworkSettings indicates an expected call of GetAllAgentNetworkSettings.
|
||||
func (mr *MockStoreMockRecorder) GetAllAgentNetworkSettings(ctx, lockStrength interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkSettings), ctx, lockStrength)
|
||||
}
|
||||
Reference in New Issue
Block a user