new-api 代码学习

new-api:新一代大模型网关与AI资产管理系统.

程序入口

程序入口：main.go

embed 指令

//go:embed web/dist
var buildFS embed.FS

embed.FS 是一个文件系统，用来嵌入整个目录。

这段代码表示将当前目录下的 web/dist 目录及其所有内容嵌入最终二进制文件。

//go:embed web/dist/index.html
var indexPage []byte

这里直接嵌入单个 HTML 文件到字节切片中，适合快速访问单个文件，比如作为模板加载或者直接返回给客户端。

读取配置文件

err := godotenv.Load(".env")
if err != nil {
    common.SysLog("Support for .env file is disabled")
}

godotenv.Load(".env") 用于加载 .env 文件,将配置读取到环境变量中（通过 os.Getenv("xxx") 来获取）。

命令行参数读取

common/init.go

使用到的包。 flag、log、os、path/filepath.

flag 用于解析命令行参数。
log 记录日志
os 读取环境变量和读取文件
path/filepath 跨平台文件路径处理

flag 解析命令行参数

flag.Int、flag.String、flag.Bool 返回值是一个指针类型，在使用时需要解引用（在变量前面加上*）。

flag.Int("port",3000,"the listening port")
flag.Bool("version", false, "print version and exit")
flag.String("log-dir", "./logs", "specify the log directory")

os 读取环境变量

ss := os.Getenv("SESSION_SECRET")

os 操作文件

*os.Stat(LogDir) 获取文件信息
os.IsNotExist(err) 判断文件是否存在
*os.Mkdir(LogDir, 0777) 创建文件夹，权限为 0777

if _, err := os.Stat(*LogDir); os.IsNotExist(err) {
    err = os.Mkdir(*LogDir, 0777) // 文件不存在时，创建文件夹，权限为 0777
    if err != nil {
        log.Fatal(err)
    }
}

path/filepath 获取绝对路径

// 通过命令行获取到 log-dir 参数
LogDir = flag.String("log-dir", "./logs", "specify the log directory")
// 获取到 logdir 的绝对路径
*LogDir, err = filepath.Abs(*LogDir) 

源码

package common

import (
	"flag"          // 处理命令行参数
	"fmt"           // 格式化输入输出
	"log"           // 记录日志
	"os"            // 读取环境变量、文件操作等
	"path/filepath" // 跨平台文件路径处理
)

var (
	Port         = flag.Int("port", 3000, "the listening port")
	PrintVersion = flag.Bool("version", false, "print version and exit")
	PrintHelp    = flag.Bool("help", false, "print help and exit")
	LogDir       = flag.String("log-dir", "./logs", "specify the log directory")
)

func printHelp() {
	fmt.Println("New API " + Version + " - All in one API service for OpenAI API.")
	fmt.Println("Copyright (C) 2023 JustSong. All rights reserved.")
	fmt.Println("GitHub: https://github.com/songquanpeng/one-api")
	fmt.Println("Usage: one-api [--port <port>] [--log-dir <log directory>] [--version] [--help]")
}

func LoadEnv() {
	flag.Parse() // 解析命令行参数，赋值给全局变量

	if *PrintVersion {
		fmt.Println(Version)
		os.Exit(0)
	}

	if *PrintHelp {
		printHelp()
		os.Exit(0)
	}

	if os.Getenv("SESSION_SECRET") != "" {
		ss := os.Getenv("SESSION_SECRET")
		if ss == "random_string" {
			log.Println("WARNING: SESSION_SECRET is set to the default value 'random_string', please change it to a random string.")
			log.Println("警告：SESSION_SECRET被设置为默认值'random_string'，请修改为随机字符串。")
			log.Fatal("Please set SESSION_SECRET to a random string.")
		} else {
			SessionSecret = ss
		}
	}
    // delete something.... 
    
	if *LogDir != "" {
		var err error
		*LogDir, err = filepath.Abs(*LogDir) // 获取 LogDir 的绝对路径
		if err != nil {
			log.Fatal(err) // log.Fatal() 会打印错误信息并退出程序
		}
		if _, err := os.Stat(*LogDir); os.IsNotExist(err) {
			err = os.Mkdir(*LogDir, 0777) // 文件不存在时，创建文件夹，权限为 0777
			if err != nil {
				log.Fatal(err)
			}
		}
	}
}

日志记录

核心:一个基于gin的日志系统，支持日志文件切换、并发安全、不同级别的日志记录，以及配额格式化等功能。但需要注意潜在的竞态条件和资源管理问题。

const (
	loggerINFO  = "INFO"
	loggerWarn  = "WARN"
	loggerError = "ERR"
)

// 日志文件最大行数
const maxLogCount = 1000000

// 计数器
var logCount int

// 互斥锁，并发控制，防止多个goroutine同时操作日志文件。
// 保证日志初始化操作的原子性。
var setupLogLock sync.Mutex
var setupLogWorking bool

func SetupLogger() {
	if *LogDir != "" {
		ok := setupLogLock.TryLock()
		if !ok {
			// 获取失败表示已经有 goroutine 在处理。
			log.Println("setup log is already working")
			return
		}
		// 在函数退出时释放锁
		defer func() {
			setupLogLock.Unlock()
			setupLogWorking = false
		}()
		// 日志文件
		logPath := filepath.Join(*LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102150405")))
		// 打开文件, 权限：追加、创建、只写
		fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			log.Fatal("failed to open log file")
		}
		// 将gin的默认写和错误写设置为同时输出到标准输出和日志文件。这说明该日志系统同时支持控制台和文件输出。
		gin.DefaultWriter = io.MultiWriter(os.Stdout, fd)
		gin.DefaultErrorWriter = io.MultiWriter(os.Stderr, fd)
	}
}

func logHelper(ctx context.Context, level string, msg string) {
	writer := gin.DefaultErrorWriter
	if level == loggerINFO {
		writer = gin.DefaultWriter
	}
	// 获取请求ID, 从请求中获取  X-Oneapi-Request-Id 请求头
	id := ctx.Value(RequestIdKey)
	now := time.Now()
	_, _ = fmt.Fprintf(writer, "[%s] %v | %s | %s \n", level, now.Format("2006/01/02 - 15:04:05"), id, msg)
	logCount++ // we don't need accurate count, so no lock here
	if logCount > maxLogCount && !setupLogWorking {
		logCount = 0
		setupLogWorking = true
		// 通过gopool.Go启动一个goroutine调用SetupLogger
		gopool.Go(func() {
			SetupLogger()
		})
	}
}

Gin release mode

Gin框架有不同的模式，比如debug和release.

debug(默认模式): 有更多的日志输出(会打印路由信息、错误详情)。 开发环境使用，以便获得更详细的调试信息，比如请求的日志、错误信息等。
release: 可能更注重性能，减少日志输出。生产环境使用
test: 用于测试。测试环境使用

gorm 数据库操作

连接到数据库

https://gorm.io/zh_CN/docs/connecting_to_the_database.html

db, err := gorm.Open(mysql.New(mysql.Config{
  DSN: "gorm:gorm@tcp(127.0.0.1:3306)/gorm?charset=utf8&parseTime=True&loc=Local", // DSN data source name
  DefaultStringSize: 256, // string 类型字段的默认长度
  DisableDatetimePrecision: true, // 禁用 datetime 精度，MySQL 5.6 之前的数据库不支持
  DontSupportRenameIndex: true, // 重命名索引时采用删除并新建的方式，MySQL 5.7 之前的数据库和 MariaDB 不支持重命名索引
  DontSupportRenameColumn: true, // 用 `change` 重命名列，MySQL 8 之前的数据库和 MariaDB 不支持重命名列
  SkipInitializeWithVersion: false, // 根据当前 MySQL 版本自动配置
}), &gorm.Config{})

func chooseDB(envName string) (*gorm.DB, error) {
	defer func() {
		initCol()
	}()
	dsn := os.Getenv(envName)
	if dsn != "" {
		// use pg 数据库
		if strings.HasPrefix(dsn, "postgres://") {
			// Use PostgreSQL
			common.SysLog("using PostgreSQL as database")
			common.UsingPostgreSQL = true
			return gorm.Open(postgres.New(postgres.Config{
				DSN:                  dsn,
				PreferSimpleProtocol: true, // disables implicit prepared statement usage
			}), &gorm.Config{
				PrepareStmt: true, // precompile SQL
			})
		}
		// use SQLite
		if strings.HasPrefix(dsn, "local") {
			common.SysLog("SQL_DSN not set, using SQLite as database")
			common.UsingSQLite = true
			return useSQLite()
		}
		// Use MySQL
		common.SysLog("using MySQL as database")
		// check parseTime
		if !strings.Contains(dsn, "parseTime") {
			if strings.Contains(dsn, "?") {
				dsn += "&parseTime=true"
			} else {
				dsn += "?parseTime=true"
			}
		}
		common.UsingMySQL = true
		return gorm.Open(mysql.Open(dsn), &gorm.Config{
			PrepareStmt: true, // precompile SQL
		})
	}
	// Use SQLite
	common.SysLog("SQL_DSN not set, using SQLite as database")
	common.UsingSQLite = true
	return useSQLite()
}

// 使用 SQLite 数据库
func useSQLite() (*gorm.DB, error) {
	return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
		PrepareStmt: true, // precompile SQL
	})
}

初始化数据库

func InitDB() (err error) {
	db, err := chooseDB("SQL_DSN")
	if err == nil {
		if common.DebugEnabled {
			db = db.Debug()
		}
		DB = db
		sqlDB, err := DB.DB()
		if err != nil {
			return err
		}
		sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100))
		sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000))
		sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60)))

		// 非主节点
		if !common.IsMasterNode {
			return nil
		}
		// 数据库表字段更新
		if common.UsingMySQL {
			_, _ = sqlDB.Exec("ALTER TABLE channels MODIFY model_mapping TEXT;") // TODO: delete this line when most users have upgraded
		}
		common.SysLog("database migration started")
		// 迁移
		err = migrateDB()
		return err
	} else {
		common.FatalLog(err)
	}
	return err
}

gorm 模型定义

语法

使用反引号包裹的 gorm:"..." 控制列属性。

`gorm:"[指令1];[指令2];..."`

常用指令

主键：gorm:"primaryKey"
自定义列名：gorm:"column:user_name"
指定列数据类型：gorm:"type:varchar(255)"
创建唯一约束：gorm:"uniqueIndex"
创建普通索引：gorm:"index"
非空约束：gorm:"not null"
默认值：gorm:"default:''"
自增列：gorm:"autoIncrement"
外键：gorm:"foreignKey:UserID"
自定义约束：gorm:"constraint:OnUpdate:CASCADE"

特殊字段处理

时间戳

CreatedAt time.Time // 自动设置为创建时间
UpdatedAt time.Time // 自动设置为更新时间

// 不需要时，可主动禁用
UpdatedAt time.Time `gorm:"autoUpdateTime:false"` // 禁用自动更新

软删除

DeletedAt gorm.DeletedAt `gorm:"index"`

表名

自定义表名

func (Product) TableName() string {
    return "tbl_products" // 自定义表名
}

GORM 默认表名规则：

结构体名称的复数形式，如 Product 结构体对应的表名为 products。
复数蛇形命名，如 ProductCategory 结构体对应的表名为 product_categories。

type Category struct{}    // ➔ categories
type Mouse struct{}       // ➔ mice（特殊复数形式）
type Index struct{}       // ➔ indices
type APIEndpoint struct{} // ➔ api_endpoints
type Users struct{}       // ➔ users（不会变成 userses）
type Criteria struct{}    // ➔ criteria（不规则复数保持）
type UserInfo struct{}    // ➔ user_infos
type JSONData struct{}    // ➔ json_data
type AWSConfig struct{}   // ➔ aws_configs
type URL struct{}          // ➔ urls（默认）
type NASAProject struct{}  // ➔ nasa_projects

表名全局配置项

db, err := gorm.Open(sqlite.Open("test.db"), &gorm.Config{
    NamingStrategy: schema.NamingStrategy{
        SingularTable: true,  // 禁用复数（User ➔ user）
        TablePrefix: "sys_",  // 表前缀（User ➔ sys_user）
        NameReplacer: strings.NewReplacer("Model", ""), // 替换规则（UserModel ➔ user_models）
    },
})

测试表名

stmt := db.Model(&YourModel{}).Statement
fmt.Println(stmt.Table) // 输出实际生成的表名

gorm 约定

https://gorm.io/zh_CN/docs/conventions.html

使用 ID 作为主键，可以通过 primaryKey 来将其他字段设为主键。
复数表名。

使用案例

type Product struct {
    ID          uint      `gorm:"primaryKey;autoIncrement"`
    ProductCode string    `gorm:"uniqueIndex:idx_code;type:varchar(20);not null"`
    Price       float64   `gorm:"type:decimal(10,2);default:0.0"`
    CategoryID  uint      `gorm:"foreignKey:CategoryRef"`
    Stock       int       `gorm:"check:stock >= 0"`       // 自定义检查约束
    IsActive    bool      `gorm:"default:true"`
    CreatedAt   time.Time
    UpdatedAt   time.Time
}

type Channel struct {
	Id                 int     `json:"id"`
	Type               int     `json:"type" gorm:"default:0"`
	Key                string  `json:"key" gorm:"not null"`
	OpenAIOrganization *string `json:"openai_organization"`
	TestModel          *string `json:"test_model"`
	Status             int     `json:"status" gorm:"default:1"`
	Name               string  `json:"name" gorm:"index"`
	Weight             *uint   `json:"weight" gorm:"default:0"`
	CreatedTime        int64   `json:"created_time" gorm:"bigint"`
	TestTime           int64   `json:"test_time" gorm:"bigint"`
	ResponseTime       int     `json:"response_time"` // in milliseconds
	BaseURL            *string `json:"base_url" gorm:"column:base_url;default:''"`
	Other              string  `json:"other"`
	Balance            float64 `json:"balance"` // in USD
	BalanceUpdatedTime int64   `json:"balance_updated_time" gorm:"bigint"`
	Models             string  `json:"models"`
	Group              string  `json:"group" gorm:"type:varchar(64);default:'default'"`
	UsedQuota          int64   `json:"used_quota" gorm:"bigint;default:0"`
	ModelMapping       *string `json:"model_mapping" gorm:"type:text"`
	//MaxInputTokens     *int    `json:"max_input_tokens" gorm:"default:0"`
	StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
	Priority          *int64  `json:"priority" gorm:"bigint;default:0"`
	AutoBan           *int    `json:"auto_ban" gorm:"default:1"`
	OtherInfo         string  `json:"other_info"`
	Tag               *string `json:"tag" gorm:"index"`
	Setting           *string `json:"setting" gorm:"type:text"`
}

迁移数据库(表结构修复)

gorm.AutoMigrate 方法

参数是 interface{} 类型
自动创建表：如果数据库中没有对应表
更新结构：如果模型有新增字段，自动添加列。gorm 不会自动删除已移除的字段。
维护约束：自动设置主键、索引、外键等，根据传入的模型定义
类型匹配：确保 go 结构体字段与数据库类型对应

func migrateDB() error {
	err := DB.AutoMigrate(&Channel{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Token{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&User{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Option{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Redemption{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Ability{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Log{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Midjourney{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&TopUp{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&QuotaData{})
	if err != nil {
		return err
	}
	err = DB.AutoMigrate(&Task{})
	if err != nil {
		return err
	}
	common.SysLog("database migrated")
	err = createRootAccountIfNeed()
	return err
}

包初始化过程

https://go.dev/ref/spec#Package_initialization

为所有包级变量分配初始值
按照在源代码中出现的顺序依次调用 init 函数
依照编译器接收文件的顺序进行处理

变量初始化的顺序

在包内部，包级变量的初始化是逐步进行的，每一步选择声明顺序最早且不依赖于任何未初始化变量的变量。通过反复初始化，知道没有就绪可初始化变量(ready for initialization)为止。

在变量声明时，左边有多个变量，右边是一个多值表达式时。这些变量会一起被初始化。

多文件中声明的变量顺序由文件呈现给编译器的顺序决定。

案例1

var x = a
var a, b = f() // a and b are initialized together, before x is initialized

a 、 b 会在 x 之前一起被初始化，

案例2

var (
	a = c + b  // == 9
	b = f()    // == 4
	c = f()    // == 5
	d = 3      // == 5 after initialization has finished
)

func f() int {
	d++
	return d
}

初始化顺序为： d、b、c、a

案例3

var x = I(T{}).ab()   // x has an undetected, hidden dependency on a and b
var _ = sideEffect()  // unrelated to x, a, or b
var a = b
var b = 42

type I interface      { ab() []int }
type T struct{}
func (T) ab() []int   { return []int{a, b} }

b 在 a 之前初始化
x 的初始化时机不确定，可能是 b 之前， b a 之间，a 之后
sideEffect 的调用时机也不确定，可能是在 x 初始化之前或之后。

init 方法

init 方法的特点

每个包可以包含多个 init 函数，包的每个源文件也可以有多个 init 函数
包中init 字符只能用于声明 init 函数
没有参数，返回值，不能被其他方法调用
先与main方法执行
包级变量的声明是早于 inti 函数的

init 与 main 的顺序

在 main 之前，所有被导入的 init 函数会在 main 函数之前自动执行。
按依赖顺序：例如包 A 导入了包 B，则 B 的 init 先于 A 的 init 执行。
单次执行：每个包的 init 函数只会执行一次（即使被多次导入）。

package main

import "pkgA"
import "pkgB"

func init() { fmt.Println("main init 1") }
func init() { fmt.Println("main init 2") }

func main() {
    pkgA.Use()
    pkgB.Use()
}

执行顺序：

pkgA 和 pkgB 的依赖包初始化（递归处理）。
pkgA 的变量 → pkgA 的 init 函数。
pkgB 的变量 → pkgB 的 init 函数。
main 包的 init 函数按声明顺序执行。
最后执行 main() 函数。

程序入口​

embed 指令​

读取配置文件​

命令行参数读取​

flag 解析命令行参数​

os 读取环境变量​

os 操作文件​

path/filepath 获取绝对路径​

源码​

日志记录​

Gin release mode​

gorm 数据库操作​

连接到数据库​

初始化数据库​

gorm 模型定义​

语法​

特殊字段处理​

表名​

gorm 约定​

使用案例​

迁移数据库(表结构修复)​

包初始化过程​

变量初始化的顺序​

init 方法​