// // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // package dao import ( "fmt" "ragflow/internal/model" "ragflow/internal/server" "time" gormLogger "gorm.io/gorm/logger" "gorm.io/driver/mysql" "gorm.io/gorm" "ragflow/internal/logger" ) var DB *gorm.DB // InitDB initialize database connection func InitDB() error { cfg := server.GetConfig() dbCfg := cfg.Database dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", dbCfg.Username, dbCfg.Password, dbCfg.Host, dbCfg.Port, dbCfg.Database, dbCfg.Charset, ) // Set log level var gormLogLevel gormLogger.LogLevel if cfg.Server.Mode == "debug" { gormLogLevel = gormLogger.Info } else { gormLogLevel = gormLogger.Silent } // Connect to database var err error DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{ Logger: gormLogger.Default.LogMode(gormLogLevel), NowFunc: func() time.Time { return time.Now().Local() }, }) if err != nil { return fmt.Errorf("failed to connect database: %w", err) } // Get general database object sql.DB sqlDB, err := DB.DB() if err != nil { return fmt.Errorf("failed to get database instance: %w", err) } // Set connection pool sqlDB.SetMaxIdleConns(10) sqlDB.SetMaxOpenConns(100) sqlDB.SetConnMaxLifetime(time.Hour) // Auto migrate all models models := []interface{}{ &model.User{}, &model.Tenant{}, &model.UserTenant{}, &model.File{}, &model.File2Document{}, &model.TenantLLM{}, &model.Chat{}, &model.ChatSession{}, &model.Task{}, &model.APIToken{}, &model.API4Conversation{}, &model.Knowledgebase{}, &model.InvitationCode{}, &model.Document{}, &model.UserCanvas{}, &model.CanvasTemplate{}, &model.UserCanvasVersion{}, &model.LLMFactories{}, &model.LLM{}, &model.TenantLangfuse{}, &model.SystemSettings{}, &model.Connector{}, &model.Connector2Kb{}, &model.SyncLogs{}, &model.MCPServer{}, &model.Memory{}, &model.Search{}, &model.PipelineOperationLog{}, &model.EvaluationDataset{}, &model.EvaluationCase{}, &model.EvaluationRun{}, &model.EvaluationResult{}, } for _, m := range models { if err := DB.AutoMigrate(m); err != nil { return fmt.Errorf("failed to migrate model %T: %w", m, err) } } // Run manual migrations for complex schema changes if err := RunMigrations(DB); err != nil { return fmt.Errorf("failed to run manual migrations: %w", err) } logger.Info("Database connected and migrated successfully") return nil } // GetDB get database instance func GetDB() *gorm.DB { return DB }