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
4144func 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.
205208func 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
245248func 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
268271func 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
301304func 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
324327func 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
347350func 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
370373func 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+ }
0 commit comments