Files
ai/configs/sqlite3.go
2023-04-28 14:23:19 +08:00

98 lines
1.9 KiB
Go

package configs
import (
"database/sql"
"log"
"os"
_ "github.com/mattn/go-sqlite3"
)
// 使用SQLite3初始化數據庫
func init() {
// 設置日誌顯示文件名和行號
log.SetFlags(log.Lshortfile | log.LstdFlags)
// 原生golang 創建 data 目錄不存在則創建
if _, err := os.Stat("data"); os.IsNotExist(err) {
os.Mkdir("data", os.ModePerm)
}
// 初始化數據庫
db, err := sql.Open("sqlite3", "data/sqlite3.db")
if err != nil {
log.Fatal(err)
}
// 一次性創建多個數據表(自增主鍵)
_, err = db.Exec(`
CREATE TABLE IF NOT EXISTS images(
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
width INTEGER,
height INTEGER,
prompt TEXT,
negative_prompt TEXT,
num_inference_steps INTEGER,
guidance_scale REAL,
scheduler TEXT,
seed INTEGER,
from_image TEXT,
created_at TEXT,
updated_at TEXT,
user_id INTEGER
);
CREATE TABLE IF NOT EXISTS models(
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
type TEXT,
trigger_words TEXT,
base_model TEXT,
model_path TEXT,
status TEXT,
progress INTEGER,
image TEXT,
tags TEXT,
created_at TEXT,
updated_at TEXT,
user_id INTEGER
);
CREATE TABLE IF NOT EXISTS users(
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
password TEXT,
created_at TEXT,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS tags(
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
created_at TEXT,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS tasks(
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
status TEXT,
progress INTEGER,
created_at TEXT,
updated_at TEXT,
user_id INTEGER
);
`)
defer db.Close()
if err != nil {
log.Fatal(err)
}
}
// GetDB 獲取數據庫連接
func GetDB() (*sql.DB, error) {
db, err := sql.Open("sqlite3", "data/sqlite3.db")
if err != nil {
return nil, err
}
return db, nil
}