add some scriptsmain

1 files changed, 38 insertions, 3 deletions

diff --git a/internal/controlplane/snapshot.go b/internal/controlplane/snapshot.go
index c8931ab..bc04e7c 100644
--- a/internal/controlplane/snapshot.go
+++ b/internal/controlplane/snapshot.go
@@ -37,12 +37,40 @@ func (s *Store) LoadSnapshot(ctx context.Context) (Snapshot, error) {
 	if err != nil {
 		return Snapshot{}, err
 	}
+	result.Limits, err = loadLimitPolicies(ctx, tx)
+	if err != nil {
+		return Snapshot{}, err
+	}
 	if err := tx.Commit(ctx); err != nil {
 		return Snapshot{}, fmt.Errorf("commit snapshot transaction: %w", err)
 	}
 	return result, nil
 }
 
+func loadLimitPolicies(ctx context.Context, tx pgx.Tx) ([]domain.LimitPolicy, error) {
+	rows, err := tx.Query(ctx, `
+		SELECT tenant_id::text, project_id::text, requests_per_minute, tokens_per_minute,
+		       concurrent_requests, monthly_spend_micros
+		FROM project_limits`)
+	if err != nil {
+		return nil, fmt.Errorf("query project limits: %w", err)
+	}
+	defer rows.Close()
+	result := make([]domain.LimitPolicy, 0)
+	for rows.Next() {
+		var policy domain.LimitPolicy
+		if err := rows.Scan(&policy.TenantID, &policy.ProjectID, &policy.RequestsPerMinute,
+			&policy.TokensPerMinute, &policy.Concurrent, &policy.MonthlySpendMicros); err != nil {
+			return nil, fmt.Errorf("scan project limit: %w", err)
+		}
+		result = append(result, policy)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("read project limits: %w", err)
+	}
+	return result, nil
+}
+
 func (s *Store) loadProviders(ctx context.Context, tx pgx.Tx) (map[string]domain.Provider, error) {
 	rows, err := tx.Query(ctx, `
 		SELECT id::text, name, protocol, base_url, api_key_ciphertext
@@ -74,7 +102,9 @@ func (s *Store) loadProviders(ctx context.Context, tx pgx.Tx) (map[string]domain
 
 func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Provider) ([]domain.Model, error) {
 	rows, err := tx.Query(ctx, `
-		SELECT m.public_id, m.owned_by, r.provider_id::text, r.upstream_model, r.priority, r.weight
+		SELECT m.public_id, m.owned_by, m.input_price_micros_per_million, m.output_price_micros_per_million,
+		       m.cache_read_price_micros_per_million, m.cache_write_price_micros_per_million,
+		       r.provider_id::text, r.upstream_model, r.priority, r.weight
 		FROM models m
 		JOIN model_routes r ON r.model_id = m.id AND r.enabled = TRUE
 		JOIN providers p ON p.id = r.provider_id AND p.enabled = TRUE
@@ -88,8 +118,9 @@ func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Prov
 	index := make(map[string]int)
 	for rows.Next() {
 		var publicID, ownedBy, providerID, upstreamModel string
+		var inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64
 		var priority, weight int
-		if err := rows.Scan(&publicID, &ownedBy, &providerID, &upstreamModel, &priority, &weight); err != nil {
+		if err := rows.Scan(&publicID, &ownedBy, &inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice, &providerID, &upstreamModel, &priority, &weight); err != nil {
 			return nil, fmt.Errorf("scan model route: %w", err)
 		}
 		provider, ok := providers[providerID]
@@ -100,7 +131,11 @@ func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Prov
 		if !exists {
 			position = len(models)
 			index[publicID] = position
-			models = append(models, domain.Model{ID: publicID, OwnedBy: ownedBy})
+			models = append(models, domain.Model{
+				ID: publicID, OwnedBy: ownedBy,
+				InputPriceMicrosPerMillion: inputPrice, OutputPriceMicrosPerMillion: outputPrice,
+				CacheReadPriceMicrosPerMillion: cacheReadPrice, CacheWritePriceMicrosPerMillion: cacheWritePrice,
+			})
 		}
 		models[position].Routes = append(models[position].Routes, domain.Route{
 			Provider: provider, UpstreamModel: upstreamModel, Priority: priority, Weight: weight,