Skip to content

Commit 8e67042

Browse files
minor extrensibility updates
1 parent b895c3a commit 8e67042

File tree

28 files changed

+271
-40
lines changed

28 files changed

+271
-40
lines changed

framework/configstore/migrations.go

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import (
44
"context"
55
"fmt"
66

7-
"github.com/maximhq/bifrost/framework/configstore/internal/migration"
7+
"github.com/maximhq/bifrost/framework/configstore/migrator"
88
"gorm.io/gorm"
99
)
1010

@@ -34,12 +34,15 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
3434
if err := migrationAddEnableLiteLLMFallbacksColumn(ctx, db); err != nil {
3535
return err
3636
}
37+
if err := migrationTeamsTableUpdates(ctx, db); err != nil {
38+
return err
39+
}
3740
return nil
3841
}
3942

4043
// migrationInit is the first migration
4144
func migrationInit(ctx context.Context, db *gorm.DB) error {
42-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
45+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
4346
ID: "init",
4447
Migrate: func(tx *gorm.DB) error {
4548
tx = tx.WithContext(ctx)
@@ -203,7 +206,7 @@ func migrationInit(ctx context.Context, db *gorm.DB) error {
203206

204207
// createMany2ManyJoinTable creates a many-to-many join table for the given tables.
205208
func migrationMany2ManyJoinTable(ctx context.Context, db *gorm.DB) error {
206-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
209+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
207210
ID: "many2manyjoin",
208211
Migrate: func(tx *gorm.DB) error {
209212
tx = tx.WithContext(ctx)
@@ -243,7 +246,7 @@ func migrationMany2ManyJoinTable(ctx context.Context, db *gorm.DB) error {
243246

244247
// migrationAddCustomProviderConfigJSONColumn adds the custom_provider_config_json column to the provider table
245248
func migrationAddCustomProviderConfigJSONColumn(ctx context.Context, db *gorm.DB) error {
246-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
249+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
247250
ID: "addcustomproviderconfigjsoncolumn",
248251
Migrate: func(tx *gorm.DB) error {
249252
tx = tx.WithContext(ctx)
@@ -266,7 +269,7 @@ func migrationAddCustomProviderConfigJSONColumn(ctx context.Context, db *gorm.DB
266269

267270
// migrationAddVirtualKeyProviderConfigTable adds the virtual_key_provider_config table
268271
func migrationAddVirtualKeyProviderConfigTable(ctx context.Context, db *gorm.DB) error {
269-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
272+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
270273
ID: "addvirtualkeyproviderconfig",
271274
Migrate: func(tx *gorm.DB) error {
272275
tx = tx.WithContext(ctx)
@@ -299,7 +302,7 @@ func migrationAddVirtualKeyProviderConfigTable(ctx context.Context, db *gorm.DB)
299302

300303
// migrationAddOpenAIUseResponsesAPIColumn adds the open_ai_use_responses_api column to the key table
301304
func migrationAddOpenAIUseResponsesAPIColumn(ctx context.Context, db *gorm.DB) error {
302-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
305+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
303306
ID: "add_open_ai_use_responses_api_column",
304307
Migrate: func(tx *gorm.DB) error {
305308
tx = tx.WithContext(ctx)
@@ -322,7 +325,7 @@ func migrationAddOpenAIUseResponsesAPIColumn(ctx context.Context, db *gorm.DB) e
322325

323326
// migrationAddAllowedOriginsJSONColumn adds the allowed_origins_json column to the client config table
324327
func migrationAddAllowedOriginsJSONColumn(ctx context.Context, db *gorm.DB) error {
325-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
328+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
326329
ID: "add_allowed_origins_json_column",
327330
Migrate: func(tx *gorm.DB) error {
328331
tx = tx.WithContext(ctx)
@@ -345,7 +348,7 @@ func migrationAddAllowedOriginsJSONColumn(ctx context.Context, db *gorm.DB) erro
345348

346349
// migrationAddAllowDirectKeysColumn adds the allow_direct_keys column to the client config table
347350
func migrationAddAllowDirectKeysColumn(ctx context.Context, db *gorm.DB) error {
348-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
351+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
349352
ID: "add_allow_direct_keys_column",
350353
Migrate: func(tx *gorm.DB) error {
351354
tx = tx.WithContext(ctx)
@@ -368,7 +371,7 @@ func migrationAddAllowDirectKeysColumn(ctx context.Context, db *gorm.DB) error {
368371

369372
// migrationAddEnableLiteLLMFallbacksColumn adds the enable_litellm_fallbacks column to the client config table
370373
func migrationAddEnableLiteLLMFallbacksColumn(ctx context.Context, db *gorm.DB) error {
371-
m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
374+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
372375
ID: "add_enable_litellm_fallbacks_column",
373376
Migrate: func(tx *gorm.DB) error {
374377
tx = tx.WithContext(ctx)
@@ -396,3 +399,35 @@ func migrationAddEnableLiteLLMFallbacksColumn(ctx context.Context, db *gorm.DB)
396399
}
397400
return nil
398401
}
402+
403+
// migrationTeamsTableUpdates adds profile, config, and claims columns to the team table
404+
func migrationTeamsTableUpdates(ctx context.Context, db *gorm.DB) error {
405+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
406+
ID: "add_profile_config_claims_columns_to_team_table",
407+
Migrate: func(tx *gorm.DB) error {
408+
tx = tx.WithContext(ctx)
409+
migrator := tx.Migrator()
410+
if !migrator.HasColumn(&TableTeam{}, "profile") {
411+
if err := migrator.AddColumn(&TableTeam{}, "profile"); err != nil {
412+
return err
413+
}
414+
}
415+
if !migrator.HasColumn(&TableTeam{}, "config") {
416+
if err := migrator.AddColumn(&TableTeam{}, "config"); err != nil {
417+
return err
418+
}
419+
}
420+
if !migrator.HasColumn(&TableTeam{}, "claims") {
421+
if err := migrator.AddColumn(&TableTeam{}, "claims"); err != nil {
422+
return err
423+
}
424+
}
425+
return nil
426+
},
427+
}})
428+
err := m.Migrate()
429+
if err != nil {
430+
return fmt.Errorf("error while running db migration: %s", err.Error())
431+
}
432+
return nil
433+
}

framework/configstore/internal/migration/migrator.go renamed to framework/configstore/migrator/migrator.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
1919
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
2020

21-
package migration
21+
package migrator
2222

2323
import (
2424
"context"

framework/configstore/rdb.go

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"fmt"
88

99
"github.com/maximhq/bifrost/core/schemas"
10+
"github.com/maximhq/bifrost/framework/configstore/migrator"
1011
"github.com/maximhq/bifrost/framework/logstore"
1112
"github.com/maximhq/bifrost/framework/vectorstore"
1213
"gorm.io/gorm"
@@ -41,6 +42,11 @@ func (s *RDBConfigStore) UpdateClientConfig(ctx context.Context, config *ClientC
4142
})
4243
}
4344

45+
// DB returns the underlying database connection.
46+
func (s *RDBConfigStore) DB() *gorm.DB {
47+
return s.db
48+
}
49+
4450
// GetClientConfig retrieves the client configuration from the database.
4551
func (s *RDBConfigStore) GetClientConfig(ctx context.Context) (*ClientConfig, error) {
4652
var dbConfig TableClientConfig
@@ -183,7 +189,7 @@ func (s *RDBConfigStore) UpdateProvider(ctx context.Context, provider schemas.Mo
183189
dbProvider.ProxyConfig = configCopy.ProxyConfig
184190
dbProvider.SendBackRawResponse = configCopy.SendBackRawResponse
185191
dbProvider.CustomProviderConfig = configCopy.CustomProviderConfig
186-
192+
187193
// Save the updated provider
188194
if err := tx.WithContext(ctx).Save(&dbProvider).Error; err != nil {
189195
return err
@@ -1283,6 +1289,15 @@ func (s *RDBConfigStore) removeDuplicateKeysAndNullKeys(ctx context.Context) err
12831289
return nil
12841290
}
12851291

1292+
// RunMigration runs a migration.
1293+
func (s *RDBConfigStore) RunMigration(ctx context.Context, migration *migrator.Migration) error {
1294+
if migration == nil {
1295+
return fmt.Errorf("migration cannot be nil")
1296+
}
1297+
m := migrator.New(s.db, migrator.DefaultOptions, []*migrator.Migration{migration})
1298+
return m.Migrate()
1299+
}
1300+
12861301
// Close closes the SQLite config store.
12871302
func (s *RDBConfigStore) Close(ctx context.Context) error {
12881303
sqlDB, err := s.db.DB()

framework/configstore/store.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77

88
"github.com/maximhq/bifrost/core/schemas"
9+
"github.com/maximhq/bifrost/framework/configstore/migrator"
910
"github.com/maximhq/bifrost/framework/logstore"
1011
"github.com/maximhq/bifrost/framework/vectorstore"
1112
"gorm.io/gorm"
@@ -106,6 +107,12 @@ type ConfigStore interface {
106107
// Generic transaction manager
107108
ExecuteTransaction(ctx context.Context, fn func(tx *gorm.DB) error) error
108109

110+
// DB returns the underlying database connection.
111+
DB() *gorm.DB
112+
113+
// Migration manager
114+
RunMigration(ctx context.Context, migration *migrator.Migration) error
115+
109116
// Cleanup
110117
Close(ctx context.Context) error
111118
}
@@ -115,7 +122,6 @@ func NewConfigStore(ctx context.Context, config *Config, logger schemas.Logger)
115122
if config == nil {
116123
return nil, fmt.Errorf("config cannot be nil")
117124
}
118-
119125
if !config.Enabled {
120126
return nil, nil
121127
}

framework/configstore/tables.go

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"time"
77

8+
bifrost "github.com/maximhq/bifrost/core"
89
"github.com/maximhq/bifrost/core/schemas"
910
"gorm.io/gorm"
1011
)
@@ -630,10 +631,73 @@ type TableTeam struct {
630631
Budget *TableBudget `gorm:"foreignKey:BudgetID" json:"budget,omitempty"`
631632
VirtualKeys []TableVirtualKey `gorm:"foreignKey:TeamID" json:"virtual_keys"`
632633

634+
Profile *string `gorm:"type:text" json:"-"`
635+
ParsedProfile map[string]interface{} `gorm:"-" json:"profile"`
636+
637+
Config *string `gorm:"type:text" json:"-"`
638+
ParsedConfig map[string]interface{} `gorm:"-" json:"config"`
639+
640+
Claims *string `gorm:"type:text" json:"-"`
641+
ParsedClaims map[string]interface{} `gorm:"-" json:"claims"`
642+
633643
CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
634644
UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
635645
}
636646

647+
648+
// BeforeSave hook for TableTeam to serialize JSON fields
649+
func (t *TableTeam) BeforeSave(tx *gorm.DB) error {
650+
if t.ParsedProfile != nil {
651+
data, err := json.Marshal(t.ParsedProfile)
652+
if err != nil {
653+
return err
654+
}
655+
t.Profile = bifrost.Ptr(string(data))
656+
}else{
657+
t.Profile = nil
658+
}
659+
if t.ParsedConfig != nil {
660+
data, err := json.Marshal(t.ParsedConfig)
661+
if err != nil {
662+
return err
663+
}
664+
t.Config = bifrost.Ptr(string(data))
665+
}else{
666+
t.Config = nil
667+
}
668+
if t.ParsedClaims != nil {
669+
data, err := json.Marshal(t.ParsedClaims)
670+
if err != nil {
671+
return err
672+
}
673+
t.Claims = bifrost.Ptr(string(data))
674+
}else{
675+
t.Claims = nil
676+
}
677+
return nil
678+
}
679+
680+
// AfterFind hook for TableTeam to deserialize JSON fields
681+
func (t *TableTeam) AfterFind(tx *gorm.DB) error {
682+
if t.Profile != nil {
683+
if err := json.Unmarshal([]byte(*t.Profile), &t.ParsedProfile); err != nil {
684+
return err
685+
}
686+
}
687+
if t.Config != nil {
688+
if err := json.Unmarshal([]byte(*t.Config), &t.ParsedConfig); err != nil {
689+
return err
690+
}
691+
}
692+
if t.Claims != nil {
693+
if err := json.Unmarshal([]byte(*t.Claims), &t.ParsedClaims); err != nil {
694+
return err
695+
}
696+
}
697+
return nil
698+
}
699+
700+
637701
// TableVirtualKey represents a virtual key with budget, rate limits, and team/customer association
638702
type TableVirtualKey struct {
639703
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`

transports/bifrost-http/handlers/server.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,9 @@ func (s *BifrostHTTPServer) Start() error {
556556
logger.Info("server gracefully shutdown")
557557
}
558558
// Cancelling main context
559-
s.cancel()
559+
if s.cancel != nil {
560+
s.cancel()
561+
}
560562
// Wait for shutdown to complete or timeout
561563
done := make(chan struct{})
562564
go func() {
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
"use client";
2+
import FullPageLoader from "@/components/fullPageLoader";
3+
import { useRouter } from "next/navigation";
4+
import { useEffect } from "react";
5+
6+
export default function LoginView() {
7+
const router = useRouter();
8+
9+
useEffect(() => {
10+
router.push("/");
11+
}, [router]);
12+
13+
return <FullPageLoader />;
14+
}

0 commit comments

Comments
 (0)