From 85689aed2cf497939a13bfc04acbaed38dd609b3 Mon Sep 17 00:00:00 2001 From: Jiahui <4543bxy@gmail.com> Date: Thu, 14 May 2026 12:04:16 +0800 Subject: [PATCH] feat(account): compat pgsql (#6958) --- .../pkg/database/cockroach/accountv2.go | 70 ++++++++++--------- 1 file changed, 36 insertions(+), 34 deletions(-) diff --git a/controllers/pkg/database/cockroach/accountv2.go b/controllers/pkg/database/cockroach/accountv2.go index 5b037f08b..9bac70cce 100644 --- a/controllers/pkg/database/cockroach/accountv2.go +++ b/controllers/pkg/database/cockroach/accountv2.go @@ -2094,25 +2094,43 @@ func (c *Cockroach) transferAccount( } func (c *Cockroach) InitTables() error { - enumTypes := []string{ - `CREATE TYPE IF NOT EXISTS subscription_status AS ENUM ('NORMAL', 'PAUSED', 'DEBT', 'DEBT_PRE_DELETION', 'DEBT_FINAL_DELETION', 'DELETED')`, - `CREATE TYPE IF NOT EXISTS subscription_operator AS ENUM ('created', 'upgraded', 'downgraded', 'canceled', 'renewed', 'deleted')`, - `CREATE TYPE IF NOT EXISTS subscription_pay_status AS ENUM ('pending', 'unpaid', 'paid', 'no_need', 'failed', 'expired', 'canceled')`, - `CREATE TYPE IF NOT EXISTS workspace_traffic_status AS ENUM ('active', 'exhausted', 'used_up', 'expired')`, - `CREATE TYPE IF NOT EXISTS subscription_transaction_status AS ENUM ('completed', 'pending', 'processing', 'failed', 'canceled')`, + createEnumSQL := createEnumTypeSQLForPostgres + if isCockroachDatabase(c.DB) { + createEnumSQL = createEnumTypeSQL } + enumTypes := []string{ + createEnumSQL( + "subscription_status", + []string{"NORMAL", "PAUSED", "DEBT", "DEBT_PRE_DELETION", "DEBT_FINAL_DELETION", "DELETED"}, + ), + createEnumSQL( + "subscription_operator", + []string{"created", "upgraded", "downgraded", "canceled", "renewed", "deleted"}, + ), + createEnumSQL( + "subscription_pay_status", + []string{"pending", "unpaid", "paid", "no_need", "failed", "expired", "canceled"}, + ), + createEnumSQL( + "workspace_traffic_status", + []string{"active", "exhausted", "used_up", "expired"}, + ), + createEnumSQL( + "subscription_transaction_status", + []string{"completed", "pending", "processing", "failed", "canceled"}, + )} // subscription_pay_status 如果不存在canceled 状态，则添加 enumTypes = append( enumTypes, - `ALTER TYPE subscription_pay_status ADD VALUE IF NOT EXISTS 'canceled'`, + alterEnumAddValueSQL("subscription_pay_status", "canceled"), ) enumTypes = append( enumTypes, - `ALTER TYPE subscription_pay_status ADD VALUE IF NOT EXISTS 'unpaid'`, + alterEnumAddValueSQL("subscription_pay_status", "unpaid"), ) enumTypes = append( enumTypes, - `ALTER TYPE subscription_transaction_status ADD VALUE IF NOT EXISTS 'canceled'`, + alterEnumAddValueSQL("subscription_transaction_status", "canceled"), ) for _, query := range enumTypes { err := c.DB.Exec(query).Error @@ -2170,10 +2188,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.Payment{}, `activityType`) { fmt.Println("add column activityType") tableName := types.Payment{}.TableName() - err := c.DB.Exec( - `ALTER TABLE "?" ADD COLUMN "activityType" TEXT;`, - gorm.Expr(tableName), - ).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"activityType" TEXT`)).Error if err != nil { return fmt.Errorf("failed to add column activityType: %w", err) } @@ -2181,10 +2196,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.Payment{}, "workspace_subscription_id") { fmt.Println("add column workspace_subscription_id") tableName := types.Payment{}.TableName() - err := c.DB.Exec( - `ALTER TABLE "?" ADD COLUMN "workspace_subscription_id" uuid;`, - gorm.Expr(tableName), - ).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"workspace_subscription_id" uuid`)).Error if err != nil { return fmt.Errorf("failed to add column workspace_subscription_id: %w", err) } @@ -2192,10 +2204,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.PaymentOrder{}, "workspace_subscription_id") { fmt.Println("add column workspace_subscription_id to PaymentOrder") tableName := types.PaymentOrder{}.TableName() - err := c.DB.Exec( - `ALTER TABLE "?" ADD COLUMN "workspace_subscription_id" uuid;`, - gorm.Expr(tableName), - ).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"workspace_subscription_id" uuid`)).Error if err != nil { return fmt.Errorf( "failed to add column workspace_subscription_id to PaymentOrder: %w", @@ -2206,7 +2215,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.Payment{}, "stripe") { fmt.Println("add column stripe to Payment") tableName := types.Payment{}.TableName() - err := c.DB.Exec(`ALTER TABLE "?" ADD COLUMN "stripe" JSON;`, gorm.Expr(tableName)).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"stripe" JSON`)).Error if err != nil { return fmt.Errorf("failed to add column stripe to Payment: %w", err) } @@ -2214,7 +2223,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.PaymentOrder{}, "stripe") { fmt.Println("add column stripe to Payment") tableName := types.PaymentOrder{}.TableName() - err := c.DB.Exec(`ALTER TABLE "?" ADD COLUMN "stripe" JSON;`, gorm.Expr(tableName)).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"stripe" JSON`)).Error if err != nil { return fmt.Errorf("failed to add column stripe to Payment: %w", err) } @@ -2222,10 +2231,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.WorkspaceSubscriptionPlan{}, "ai_quota") { fmt.Println("add column ai_quota to WorkspaceSubscriptionPlan") tableName := types.WorkspaceSubscriptionPlan{}.TableName() - err := c.DB.Exec( - `ALTER TABLE "?" ADD COLUMN "ai_quota" bigint NOT NULL DEFAULT 0;`, - gorm.Expr(tableName), - ).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"ai_quota" bigint NOT NULL DEFAULT 0`)).Error if err != nil { return fmt.Errorf("failed to add column ai_quota to WorkspaceSubscriptionPlan: %w", err) } @@ -2233,10 +2239,7 @@ func (c *Cockroach) InitTables() error { if !c.Localdb.Migrator().HasColumn(&types.WorkspaceSubscriptionPlan{}, "ai_quota") { fmt.Println("add column ai_quota to WorkspaceSubscriptionPlan") tableName := types.WorkspaceSubscriptionPlan{}.TableName() - err := c.Localdb.Exec( - `ALTER TABLE "?" ADD COLUMN "ai_quota" bigint NOT NULL DEFAULT 0;`, - gorm.Expr(tableName), - ).Error + err := c.Localdb.Exec(alterTableAddColumnSQL(tableName, `"ai_quota" bigint NOT NULL DEFAULT 0`)).Error if err != nil { return fmt.Errorf("failed to add column ai_quota to WorkspaceSubscriptionPlan: %w", err) } @@ -2250,8 +2253,7 @@ func (c *Cockroach) InitTables() error { fmt.Println("add table `Credits` column updated_at") tableName := types.Credits{}.TableName() err := c.DB.Exec( - `ALTER TABLE "?" ADD COLUMN "updated_at" TIMESTAMP(3) WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP;`, - gorm.Expr(tableName), + alterTableAddColumnSQL(tableName, `"updated_at" TIMESTAMP(3) WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP`), ).Error if err != nil { return fmt.Errorf("failed to add column updated_at: %w", err) @@ -2271,7 +2273,7 @@ func (c *Cockroach) InitTables() error { if !c.DB.Migrator().HasColumn(&types.Corporate{}, `type`) { fmt.Println("add table `Corporate` column type") tableName := types.Corporate{}.TableName() - err := c.DB.Exec(`ALTER TABLE "?" ADD COLUMN "type" TEXT;`, gorm.Expr(tableName)).Error + err := c.DB.Exec(alterTableAddColumnSQL(tableName, `"type" TEXT`)).Error if err != nil { return fmt.Errorf("failed to add column type: %w", err) }