feat alpha: support sqlite and connection stack restruct

admin/analysis.go

@@ -30,7 +30,7 @@ func getFormat(t time.Time) string {
 
 func GetSubscriptionUsers(db *sql.DB) int64 {
 	var count int64
-	err := db.QueryRow(`
+	err := globals.QueryRowDb(db, `
    		SELECT COUNT(*) FROM subscription WHERE expired_at > NOW()
    	`).Scan(&count)
 	if err != nil {
@@ -117,7 +117,7 @@ func GetUserTypeData(db *sql.DB) (UserTypeForm, error) {
 	var form UserTypeForm
 
 	// get total users
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT COUNT(*) FROM auth
 	`).Scan(&form.Total); err != nil {
 		return form, err
@@ -125,7 +125,7 @@ func GetUserTypeData(db *sql.DB) (UserTypeForm, error) {
 
 	// get subscription users count (current subscription)
 	// level 1: basic plan, level 2: standard plan, level 3: pro plan
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT
 			(SELECT COUNT(*) FROM subscription WHERE level = 1 AND expired_at > NOW()),
 			(SELECT COUNT(*) FROM subscription WHERE level = 2 AND expired_at > NOW()),
@@ -136,7 +136,7 @@ func GetUserTypeData(db *sql.DB) (UserTypeForm, error) {
 
 	// get normal users count (no subscription in `subscription` table and `quota` + `used` < initial quota in `quota` table)
 	initialQuota := channel.SystemInstance.GetInitialQuota()
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT COUNT(*) FROM auth 
 		WHERE id NOT IN (SELECT user_id FROM subscription WHERE total_month > 0)
 		AND id IN (SELECT user_id FROM quota WHERE quota + used <= ?)

admin/invitation.go

@@ -1,6 +1,7 @@
 package admin
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"errors"
@@ -12,7 +13,7 @@ import (
 func GetInvitationPagination(db *sql.DB, page int64) PaginationForm {
 	var invitations []interface{}
 	var total int64
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT COUNT(*) FROM invitation
 	`).Scan(&total); err != nil {
 		return PaginationForm{
@@ -21,7 +22,7 @@ func GetInvitationPagination(db *sql.DB, page int64) PaginationForm {
 		}
 	}
 
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT code, quota, type, used, updated_at FROM invitation
 		ORDER BY id DESC LIMIT ? OFFSET ?
 	`, pagination, page*pagination)
@@ -53,7 +54,7 @@ func GetInvitationPagination(db *sql.DB, page int64) PaginationForm {
 }
 
 func NewInvitationCode(db *sql.DB, code string, quota float32, t string) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO invitation (code, quota, type)
 		VALUES (?, ?, ?)
 	`, code, quota, t)

admin/redeem.go

@@ -1,6 +1,7 @@
 package admin
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"fmt"
@@ -10,7 +11,7 @@ import (
 func GetRedeemData(db *sql.DB) []RedeemData {
 	var data []RedeemData
 
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT quota, COUNT(*) AS total, SUM(IF(used = 0, 0, 1)) AS used
 		FROM redeem
 		GROUP BY quota
@@ -54,7 +55,7 @@ func GenerateRedeemCodes(db *sql.DB, num int, quota float32) RedeemGenerateRespo
 
 func CreateRedeemCode(db *sql.DB, quota float32) (string, error) {
 	code := fmt.Sprintf("nio-%s", utils.GenerateChar(32))
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO redeem (code, quota) VALUES (?, ?)
 	`, code, quota)
 

admin/user.go

@@ -2,6 +2,7 @@ package admin
 
 import (
 	"chat/channel"
+	"chat/globals"
 	"chat/utils"
 	"context"
 	"database/sql"
@@ -31,7 +32,7 @@ func getUsersForm(db *sql.DB, page int64, search string) PaginationForm {
 	var users []interface{}
 	var total int64
 
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT COUNT(*) FROM auth
 		WHERE username LIKE ?
 	`, "%"+search+"%").Scan(&total); err != nil {
@@ -41,7 +42,7 @@ func getUsersForm(db *sql.DB, page int64, search string) PaginationForm {
 		}
 	}
 
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT 
 		    auth.id, auth.username, auth.email, auth.is_admin,
 		    quota.quota, quota.used,
@@ -116,7 +117,7 @@ func passwordMigration(db *sql.DB, cache *redis.Client, id int64, password strin
 		return fmt.Errorf("password length must be between 6 and 36")
 	}
 
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE auth SET password = ? WHERE id = ?
 	`, utils.Sha2Encrypt(password), id)
 
@@ -126,7 +127,7 @@ func passwordMigration(db *sql.DB, cache *redis.Client, id int64, password strin
 }
 
 func emailMigration(db *sql.DB, id int64, email string) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE auth SET email = ? WHERE id = ?
 	`, email, id)
 
@@ -134,7 +135,7 @@ func emailMigration(db *sql.DB, id int64, email string) error {
 }
 
 func setAdmin(db *sql.DB, id int64, isAdmin bool) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE auth SET is_admin = ? WHERE id = ?
 	`, isAdmin, id)
 
@@ -142,7 +143,7 @@ func setAdmin(db *sql.DB, id int64, isAdmin bool) error {
 }
 
 func banUser(db *sql.DB, id int64, isBanned bool) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE auth SET is_banned = ? WHERE id = ?
 	`, isBanned, id)
 
@@ -154,7 +155,7 @@ func quotaMigration(db *sql.DB, id int64, quota float32, override bool) error {
 	// if quota is positive, then increase quota
 
 	if override {
-		_, err := db.Exec(`
+		_, err := globals.ExecDb(db, `
 			INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?)
 			ON DUPLICATE KEY UPDATE quota = ?
 		`, id, quota, 0., quota)
@@ -162,7 +163,7 @@ func quotaMigration(db *sql.DB, id int64, quota float32, override bool) error {
 		return err
 	}
 
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) 
 		ON DUPLICATE KEY UPDATE quota = quota + ?
 	`, id, quota, 0., quota)
@@ -176,7 +177,7 @@ func subscriptionMigration(db *sql.DB, id int64, month int64) error {
 
 	expireAt := time.Now().AddDate(0, int(month), 0)
 
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO subscription (user_id, total_month, expired_at) VALUES (?, ?, ?)
 		ON DUPLICATE KEY UPDATE total_month = total_month + ?, expired_at = DATE_ADD(expired_at, INTERVAL ? MONTH)
 	`, id, month, expireAt, month, month)
@@ -189,7 +190,7 @@ func subscriptionLevelMigration(db *sql.DB, id int64, level int64) error {
 		return fmt.Errorf("invalid subscription level")
 	}
 
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO subscription (user_id, level) VALUES (?, ?)
 		ON DUPLICATE KEY UPDATE level = ?
 	`, id, level, level)
@@ -199,7 +200,7 @@ func subscriptionLevelMigration(db *sql.DB, id int64, level int64) error {
 
 func releaseUsage(db *sql.DB, cache *redis.Client, id int64) error {
 	var level sql.NullInt64
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT level FROM subscription WHERE user_id = ?
 	`, id).Scan(&level); err != nil {
 		return err
@@ -225,7 +226,7 @@ func UpdateRootPassword(db *sql.DB, cache *redis.Client, password string) error
 		return fmt.Errorf("password length must be between 6 and 36")
 	}
 
-	if _, err := db.Exec(`
+	if _, err := globals.ExecDb(db, `
 		UPDATE auth SET password = ? WHERE username = 'root'
 	`, utils.Sha2Encrypt(password)); err != nil {
 		return err

app/src/api/connection.ts

@@ -1,6 +1,7 @@
 import { tokenField, websocketEndpoint } from "@/conf/bootstrap.ts";
 import { getMemory } from "@/utils/memory.ts";
 import { getErrorMessage } from "@/utils/base.ts";
+import { Mask } from "@/masks/types.ts";
 
 export const endpoint = `${websocketEndpoint}/chat`;
 export const maxRetry = 60; // 30s max websocket retry
@@ -31,11 +32,12 @@ export type ChatProps = {
   repetition_penalty?: number;
 };
 
-type StreamCallback = (message: StreamMessage) => void;
+type StreamCallback = (id: number, message: StreamMessage) => void;
 
 export class Connection {
   protected connection?: WebSocket;
   protected callback?: StreamCallback;
+  protected stack?: string;
   public id: number;
   public state: boolean;
 
@@ -66,6 +68,10 @@ export class Connection {
       const message = JSON.parse(event.data);
       this.triggerCallback(message as StreamMessage);
     };
+
+    this.connection.onclose = (event) => {
+      this.stack = `websocket connection failed (code: ${event.code}, reason: ${event.reason}, endpoint: ${endpoint})`;
+    };
   }
 
   public reconnect(): void {
@@ -99,23 +105,56 @@ export class Connection {
       );
     }
 
-    const trace = {
+    const trace =
-      message: data.message,
+      this.stack ||
-      endpoint: endpoint,
+      JSON.stringify(
-    };
+        {
+          message: data.message,
+          endpoint: endpoint,
+        },
+        null,
+        2,
+      );
 
     t &&
       this.triggerCallback({
-        message: `
+        message: `${t("request-failed")}\n\`\`\`json\n${trace}\n\`\`\`\n`,
-${t("request-failed")}
-\`\`\`json
-${JSON.stringify(trace, null, 2)}
-\`\`\`
-          `,
         end: true,
       });
   }
 
+  public sendEvent(t: any, event: string, data?: string) {
+    this.sendWithRetry(t, {
+      type: event,
+      message: data || "",
+      model: "event",
+    });
+  }
+
+  public sendStopEvent(t: any) {
+    this.sendEvent(t, "stop");
+  }
+
+  public sendRestartEvent(t: any) {
+    this.sendEvent(t, "restart");
+  }
+
+  public sendMaskEvent(t: any, mask: Mask) {
+    this.sendEvent(t, "mask", JSON.stringify(mask.context));
+  }
+
+  public sendEditEvent(t: any, id: number, message: string) {
+    this.sendEvent(t, "edit", `${id}:${message}`);
+  }
+
+  public sendRemoveEvent(t: any, id: number) {
+    this.sendEvent(t, "remove", id.toString());
+  }
+
+  public sendShareEvent(t: any, refer: string) {
+    this.sendEvent(t, "share", refer);
+  }
+
   public close(): void {
     if (!this.connection) return;
     this.connection.close();
@@ -126,13 +165,91 @@ ${JSON.stringify(trace, null, 2)}
   }
 
   protected triggerCallback(message: StreamMessage): void {
-    if (this.id === -1 && message.conversation) {
+    this.callback && this.callback(this.id, message);
-      this.setId(message.conversation);
-    }
-    this.callback && this.callback(message);
   }
 
   public setId(id: number): void {
     this.id = id;
   }
 }
+
+export class ConnectionStack {
+  protected connections: Connection[];
+  protected callback?: StreamCallback;
+
+  public constructor(callback?: StreamCallback) {
+    this.connections = [];
+    this.callback = callback;
+  }
+
+  public getConnection(id: number): Connection | undefined {
+    return this.connections.find((conn) => conn.id === id);
+  }
+
+  public addConnection(id: number): Connection {
+    const conn = new Connection(id, this.triggerCallback.bind(this));
+    this.connections.push(conn);
+    return conn;
+  }
+
+  public setCallback(callback?: StreamCallback): void {
+    this.callback = callback;
+  }
+
+  public sendEvent(id: number, t: any, event: string, data?: string) {
+    const conn = this.getConnection(id);
+    conn && conn.sendEvent(t, event, data);
+  }
+
+  public sendStopEvent(id: number, t: any) {
+    const conn = this.getConnection(id);
+    conn && conn.sendStopEvent(t);
+  }
+
+  public sendRestartEvent(id: number, t: any) {
+    const conn = this.getConnection(id);
+    conn && conn.sendRestartEvent(t);
+  }
+
+  public sendMaskEvent(id: number, t: any, mask: Mask) {
+    const conn = this.getConnection(id);
+    conn && conn.sendMaskEvent(t, mask);
+  }
+
+  public sendEditEvent(id: number, t: any, message: string) {
+    const conn = this.getConnection(id);
+    conn && conn.sendEditEvent(t, id, message);
+  }
+
+  public sendRemoveEvent(id: number, t: any, messageId: number) {
+    const conn = this.getConnection(id);
+    conn && conn.sendRemoveEvent(t, messageId);
+  }
+
+  public sendShareEvent(id: number, t: any, refer: string) {
+    const conn = this.getConnection(id);
+    conn && conn.sendShareEvent(t, refer);
+  }
+
+  public close(id: number): void {
+    const conn = this.getConnection(id);
+    conn && conn.close();
+  }
+
+  public closeAll(): void {
+    this.connections.forEach((conn) => conn.close());
+  }
+
+  public reconnect(id: number): void {
+    const conn = this.getConnection(id);
+    conn && conn.reconnect();
+  }
+
+  public reconnectAll(): void {
+    this.connections.forEach((conn) => conn.reconnect());
+  }
+
+  public triggerCallback(id: number, message: StreamMessage): void {
+    this.callback && this.callback(id, message);
+  }
+}

app/src/api/conversation.ts

@@ -9,6 +9,13 @@ import { Mask } from "@/masks/types.ts";
 
 type ConversationCallback = (idx: number, message: Message[]) => boolean;
 
+export type ConversationSerialized = {
+  model: string;
+  end: boolean;
+  mask: Mask | null;
+  messages: Message[];
+};
+
 export class Conversation {
   protected connection?: Connection;
   protected callback?: ConversationCallback;
@@ -100,14 +107,14 @@ export class Conversation {
     });
   }
 
-  public sendStopEvent() {
-    this.sendEvent("stop");
-  }
-
   public isValidIndex(idx: number): boolean {
     return idx >= 0 && idx < this.data.length;
   }
 
+  public sendStopEvent() {
+    this.sendEvent("stop");
+  }
+
   public sendRestartEvent() {
     this.sendEvent("restart");
   }
@@ -253,10 +260,6 @@ export class Conversation {
     };
   }
 
-  public getSegmentData(length: number): Message[] {
-    return this.data.slice(this.data.length - length - 1, this.data.length - 1);
-  }
-
   public send(t: any, props: ChatProps) {
     if (!this.connection) {
       this.connection = new Connection(this.id);

app/src/assets/main.less

@@ -91,7 +91,7 @@ strong {
 
 .flex-dialog {
   border-radius: var(--radius) !important;
-  max-height: calc(100vh - 2rem) !important;
+  max-height: calc(95vh - 2rem) !important;
   overflow-x: hidden;
   overflow-y: auto;
   scrollbar-width: none;
@@ -128,7 +128,7 @@ strong {
 
 .fixed-dialog {
   border-radius: var(--radius) !important;
-  max-height: calc(100vh - 2rem) !important;
+  max-height: calc(95vh - 2rem) !important;
   min-height: 60vh;
   overflow-x: hidden;
   overflow-y: auto;
@@ -197,4 +197,4 @@ strong {
 .chat-logo {
   border-radius: var(--radius);
   user-select: none;
-}
+}

app/src/assets/pages/home.less

@@ -497,7 +497,7 @@
       font-size: 14px;
       user-select: none;
 
-      &:before {
+      &:not(.loading):before {
         content: "#";
         font-size: 12px;
         margin-right: 1px;

app/src/components/admin/assemblies/ModelUsageChart.tsx

@@ -65,7 +65,7 @@ function ModelUsageChart({ labels, datasets }: ModelChartProps) {
       <div className="chart-tooltip min-w-56 w-max z-10 rounded-tremor-default border border-tremor-border bg-tremor-background p-2 text-tremor-default shadow-tremor-dropdown">
         <div className="flex flex-1 space-x-2.5">
           <div
-            className={`flex w-1.5 flex-col bg-${categoryPayload?.color}-500 rounded`}
+            className={`flex w-1.5 flex-col bg-${categoryPayload?.color} rounded`}
           />
           <div className="w-full">
             <div className="flex items-center justify-between space-x-8">

app/src/components/home/ConversationSegment.tsx

@@ -2,7 +2,13 @@ import { toggleConversation } from "@/api/history.ts";
 import { mobile } from "@/utils/device.ts";
 import { filterMessage } from "@/utils/processor.ts";
 import { setMenu } from "@/store/menu.ts";
-import { MessageSquare, MoreHorizontal, Share2, Trash2 } from "lucide-react";
+import {
+  Loader2,
+  MessageSquare,
+  MoreHorizontal,
+  Share2,
+  Trash2,
+} from "lucide-react";
 import {
   DropdownMenu,
   DropdownMenuContent,
@@ -34,6 +40,8 @@ function ConversationSegment({
   const [open, setOpen] = useState(false);
   const [offset, setOffset] = useState(0);
 
+  const loading = conversation.id <= 0;
+
   return (
     <div
       className={cn("conversation", current === conversation.id && "active")}
@@ -51,7 +59,13 @@ function ConversationSegment({
     >
       <MessageSquare className={`h-4 w-4 mr-1`} />
       <div className={`title`}>{filterMessage(conversation.name)}</div>
-      <div className={`id`}>{conversation.id}</div>
+      <div className={cn("id", loading && "loading")}>
+        {loading ? (
+          <Loader2 className={`mr-0.5 h-4 w-4 animate-spin`} />
+        ) : (
+          conversation.id
+        )}
+      </div>
       <DropdownMenu
         open={open}
         onOpenChange={(state: boolean) => {

app/src/store/chat.ts

@@ -17,10 +17,15 @@ import {
 } from "@/conf/storage.ts";
 import { CustomMask } from "@/masks/types.ts";
 import { listMasks } from "@/api/mask.ts";
+import { ConversationSerialized } from "@/api/conversation.ts";
+import { useSelector } from "react-redux";
+import { useMemo } from "react";
+import { ConnectionStack } from "@/api/connection.ts";
 
 type initialStateType = {
   history: ConversationInstance[];
   messages: Message[];
+  conversations: Record<number, ConversationSerialized>;
   model: string;
   web: boolean;
   current: number;
@@ -60,12 +65,14 @@ export function getModelList(
   return target;
 }
 
+export const stack = new ConnectionStack();
 const offline = loadPreferenceModels(getOfflineModels());
 const chatSlice = createSlice({
   name: "chat",
   initialState: {
     history: [],
     messages: [],
+    conversations: {},
     web: getBooleanMemory("web", false),
     current: -1,
     model: getModel(offline, getMemory("model")),
@@ -214,6 +221,9 @@ export const selectHistory = (state: RootState): ConversationInstance[] =>
   state.chat.history;
 export const selectMessages = (state: RootState): Message[] =>
   state.chat.messages;
+export const selectConversations = (
+  state: RootState,
+): Record<number, ConversationSerialized> => state.chat.conversations;
 export const selectModel = (state: RootState): string => state.chat.model;
 export const selectWeb = (state: RootState): boolean => state.chat.web;
 export const selectCurrent = (state: RootState): number => state.chat.current;
@@ -226,6 +236,23 @@ export const selectCustomMasks = (state: RootState): CustomMask[] =>
 export const selectSupportModels = (state: RootState): Model[] =>
   state.chat.support_models;
 
+export function useConversation(): ConversationSerialized | undefined {
+  const conversations = useSelector(selectConversations);
+  const current = useSelector(selectCurrent);
+
+  return useMemo(() => conversations[current], [conversations, current]);
+}
+
+export function useMessages(): Message[] {
+  const conversations = useSelector(selectConversations);
+  const current = useSelector(selectCurrent);
+
+  return useMemo(
+    () => conversations[current]?.messages || [],
+    [conversations, current],
+  );
+}
+
 export const updateMasks = async (dispatch: AppDispatch) => {
   const resp = await listMasks();
   resp.data.length > 0 && dispatch(setCustomMasks(resp.data));

auth/apikey.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"errors"
@@ -10,7 +11,7 @@ import (
 func (u *User) CreateApiKey(db *sql.DB) string {
 	salt := utils.Sha2Encrypt(fmt.Sprintf("%s-%s", u.Username, utils.GenerateChar(utils.GetRandomInt(720, 1024))))
 	key := fmt.Sprintf("sk-%s", salt[:64]) // 64 bytes
-	if _, err := db.Exec("INSERT INTO apikey (user_id, api_key) VALUES (?, ?)", u.GetID(db), key); err != nil {
+	if _, err := globals.ExecDb(db, "INSERT INTO apikey (user_id, api_key) VALUES (?, ?)", u.GetID(db), key); err != nil {
 		return ""
 	}
 	return key
@@ -18,14 +19,14 @@ func (u *User) CreateApiKey(db *sql.DB) string {
 
 func (u *User) GetApiKey(db *sql.DB) string {
 	var key string
-	if err := db.QueryRow("SELECT api_key FROM apikey WHERE user_id = ?", u.GetID(db)).Scan(&key); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT api_key FROM apikey WHERE user_id = ?", u.GetID(db)).Scan(&key); err != nil {
 		return u.CreateApiKey(db)
 	}
 	return key
 }
 
 func (u *User) ResetApiKey(db *sql.DB) (string, error) {
-	if _, err := db.Exec("DELETE FROM apikey WHERE user_id = ?", u.GetID(db)); err != nil && !errors.Is(err, sql.ErrNoRows) {
+	if _, err := globals.ExecDb(db, "DELETE FROM apikey WHERE user_id = ?", u.GetID(db)); err != nil && !errors.Is(err, sql.ErrNoRows) {
 		return "", err
 	}
 	return u.CreateApiKey(db), nil

auth/auth.go

@@ -47,7 +47,7 @@ func ParseApiKey(c *gin.Context, key string) *User {
 	}
 
 	var user User
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 			SELECT auth.id, auth.username, auth.password FROM auth 
 			INNER JOIN apikey ON auth.id = apikey.user_id 
 			WHERE apikey.api_key = ?
@@ -143,7 +143,7 @@ func SignUp(c *gin.Context, form RegisterForm) (string, error) {
 		Token:    utils.Sha2Encrypt(email + username),
 	}
 
-	if _, err := db.Exec(`
+	if _, err := globals.ExecDb(db, `
 			INSERT INTO auth (username, password, email, bind_id, token)
 			VALUES (?, ?, ?, ?, ?)
 			`, user.Username, user.Password, user.Email, user.BindID, user.Token); err != nil {
@@ -170,7 +170,7 @@ func Login(c *gin.Context, form LoginForm) (string, error) {
 
 	// get user from db by username (or email) and password
 	var user User
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 			SELECT auth.id, auth.username, auth.password FROM auth 
 			WHERE (auth.username = ? OR auth.email = ?) AND auth.password = ?
 			`, username, username, hash).Scan(&user.ID, &user.Username, &user.Password); err != nil {
@@ -202,7 +202,7 @@ func DeepLogin(c *gin.Context, token string) (string, error) {
 
 		// register
 		password := utils.GenerateChar(64)
-		_ = db.QueryRow("INSERT INTO auth (bind_id, username, token, password) VALUES (?, ?, ?, ?)",
+		_ = globals.QueryRowDb(db, "INSERT INTO auth (bind_id, username, token, password) VALUES (?, ?, ?, ?)",
 			user.ID, user.Username, token, password)
 		u := &User{
 			Username: user.Username,
@@ -214,9 +214,9 @@ func DeepLogin(c *gin.Context, token string) (string, error) {
 	}
 
 	// login
-	_ = db.QueryRow("UPDATE auth SET token = ? WHERE username = ?", token, user.Username)
+	_ = globals.QueryRowDb(db, "UPDATE auth SET token = ? WHERE username = ?", token, user.Username)
 	var password string
-	err := db.QueryRow("SELECT password FROM auth WHERE username = ?", user.Username).Scan(&password)
+	err := globals.QueryRowDb(db, "SELECT password FROM auth WHERE username = ?", user.Username).Scan(&password)
 	if err != nil {
 		return "", err
 	}
@@ -273,7 +273,7 @@ func Reset(c *gin.Context, form ResetForm) error {
 func (u *User) UpdatePassword(db *sql.DB, cache *redis.Client, password string) error {
 	hash := utils.Sha2Encrypt(password)
 
-	if _, err := db.Exec(`
+	if _, err := globals.ExecDb(db, `
 			UPDATE auth SET password = ? WHERE id = ?
 			`, hash, u.ID); err != nil {
 		return err
@@ -296,7 +296,7 @@ func (u *User) Validate(c *gin.Context) bool {
 
 	db := utils.GetDBFromContext(c)
 	var count int
-	if err := db.QueryRow("SELECT COUNT(*) FROM auth WHERE username = ? AND password = ?", u.Username, u.Password).Scan(&count); err != nil || count == 0 {
+	if err := globals.QueryRowDb(db, "SELECT COUNT(*) FROM auth WHERE username = ? AND password = ?", u.Username, u.Password).Scan(&count); err != nil || count == 0 {
 		if err != nil {
 			globals.Warn(fmt.Sprintf("validate user error: %s", err.Error()))
 		}
@@ -328,13 +328,13 @@ func (u *User) GenerateToken() (string, error) {
 
 func (u *User) GenerateTokenSafe(db *sql.DB) (string, error) {
 	if len(u.Username) == 0 {
-		if err := db.QueryRow("SELECT username FROM auth WHERE id = ?", u.ID).Scan(&u.Username); err != nil {
+		if err := globals.QueryRowDb(db, "SELECT username FROM auth WHERE id = ?", u.ID).Scan(&u.Username); err != nil {
 			return "", err
 		}
 	}
 
 	if len(u.Password) == 0 {
-		if err := db.QueryRow("SELECT password FROM auth WHERE id = ?", u.ID).Scan(&u.Password); err != nil {
+		if err := globals.QueryRowDb(db, "SELECT password FROM auth WHERE id = ?", u.ID).Scan(&u.Password); err != nil {
 			return "", err
 		}
 	}

auth/invitation.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"errors"
@@ -36,7 +37,7 @@ func GenerateInvitations(db *sql.DB, num int, quota float32, t string) ([]string
 }
 
 func CreateInvitationCode(db *sql.DB, code string, quota float32, t string) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO invitation (code, quota, type)
 		VALUES (?, ?, ?)
 	`, code, quota, t)
@@ -44,7 +45,7 @@ func CreateInvitationCode(db *sql.DB, code string, quota float32, t string) erro
 }
 
 func GetInvitation(db *sql.DB, code string) (*Invitation, error) {
-	row := db.QueryRow(`
+	row := globals.QueryRowDb(db, `
 		SELECT id, code, quota, type, used, used_id
 		FROM invitation
 		WHERE code = ?
@@ -69,7 +70,7 @@ func (i *Invitation) IsUsed() bool {
 }
 
 func (i *Invitation) Use(db *sql.DB, userId int64) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE invitation SET used = TRUE, used_id = ? WHERE id = ?
 	`, userId, i.Id)
 	return err

auth/package.go

@@ -1,6 +1,9 @@
 package auth
 
-import "database/sql"
+import (
+	"chat/globals"
+	"database/sql"
+)
 
 type GiftResponse struct {
 	Cert     bool `json:"cert"`
@@ -9,7 +12,7 @@ type GiftResponse struct {
 
 func (u *User) HasPackage(db *sql.DB, _t string) bool {
 	var count int
-	if err := db.QueryRow(`SELECT COUNT(*) FROM package where user_id = ? AND type = ?`, u.ID, _t).Scan(&count); err != nil {
+	if err := globals.QueryRowDb(db, `SELECT COUNT(*) FROM package where user_id = ? AND type = ?`, u.ID, _t).Scan(&count); err != nil {
 		return false
 	}
 
@@ -28,7 +31,7 @@ func NewPackage(db *sql.DB, user *User, _t string) bool {
 	id := user.GetID(db)
 
 	var count int
-	if err := db.QueryRow(`SELECT COUNT(*) FROM package where user_id = ? AND type = ?`, id, _t).Scan(&count); err != nil {
+	if err := globals.QueryRowDb(db, `SELECT COUNT(*) FROM package where user_id = ? AND type = ?`, id, _t).Scan(&count); err != nil {
 		return false
 	}
 
@@ -36,7 +39,7 @@ func NewPackage(db *sql.DB, user *User, _t string) bool {
 		return false
 	}
 
-	_ = db.QueryRow(`INSERT INTO package (user_id, type) VALUES (?, ?)`, id, _t)
+	_ = globals.QueryRowDb(db, `INSERT INTO package (user_id, type) VALUES (?, ?)`, id, _t)
 	return true
 }
 

auth/quota.go

@@ -2,11 +2,12 @@ package auth
 
 import (
 	"chat/channel"
+	"chat/globals"
 	"database/sql"
 )
 
 func (u *User) CreateInitialQuota(db *sql.DB) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?)
 	`, u.GetID(db), channel.SystemInstance.GetInitialQuota(), 0.)
 	return err == nil
@@ -14,7 +15,7 @@ func (u *User) CreateInitialQuota(db *sql.DB) bool {
 
 func (u *User) GetQuota(db *sql.DB) float32 {
 	var quota float32
-	if err := db.QueryRow("SELECT quota FROM quota WHERE user_id = ?", u.GetID(db)).Scan(&quota); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT quota FROM quota WHERE user_id = ?", u.GetID(db)).Scan(&quota); err != nil {
 		return 0.
 	}
 	return quota
@@ -22,44 +23,50 @@ func (u *User) GetQuota(db *sql.DB) float32 {
 
 func (u *User) GetUsedQuota(db *sql.DB) float32 {
 	var quota float32
-	if err := db.QueryRow("SELECT used FROM quota WHERE user_id = ?", u.GetID(db)).Scan(&quota); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT used FROM quota WHERE user_id = ?", u.GetID(db)).Scan(&quota); err != nil {
 		return 0.
 	}
 	return quota
 }
 
 func (u *User) SetQuota(db *sql.DB, quota float32) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE quota = ?
 	`, u.GetID(db), quota, 0., quota)
 	return err == nil
 }
 
 func (u *User) SetUsedQuota(db *sql.DB, used float32) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE used = ?
 	`, u.GetID(db), 0., used, used)
 	return err == nil
 }
 
 func (u *User) IncreaseQuota(db *sql.DB, quota float32) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, globals.MultiSql(`
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE quota = quota + ?
-	`, u.GetID(db), quota, 0., quota)
+	`, `
+		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE SET quota = quota + ?
+	`), u.GetID(db), quota, 0., quota)
 	return err == nil
 }
 
 func (u *User) IncreaseUsedQuota(db *sql.DB, used float32) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, globals.MultiSql(`
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE used = used + ?
-	`, u.GetID(db), 0., used, used)
+	`, `
+		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE SET used = used + ?
+	`), u.GetID(db), 0., used, used)
 	return err == nil
 }
 
 func (u *User) DecreaseQuota(db *sql.DB, quota float32) bool {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, globals.MultiSql(`
 		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE quota = quota - ?
-	`, u.GetID(db), quota, 0., quota)
+	`, `
+		INSERT INTO quota (user_id, quota, used) VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE SET quota = quota - ?
+    `), u.GetID(db), quota, 0., quota)
 	return err == nil
 }
 

auth/redeem.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"errors"
@@ -34,14 +35,14 @@ func GenerateRedeemCodes(db *sql.DB, num int, quota float32) ([]string, error) {
 }
 
 func CreateRedeemCode(db *sql.DB, code string, quota float32) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO redeem (code, quota) VALUES (?, ?)
 	`, code, quota)
 	return err
 }
 
 func GetRedeemCode(db *sql.DB, code string) (*Redeem, error) {
-	row := db.QueryRow(`
+	row := globals.QueryRowDb(db, `
 		SELECT id, code, quota, used
 		FROM redeem
 		WHERE code = ?
@@ -62,7 +63,7 @@ func (r *Redeem) IsUsed() bool {
 }
 
 func (r *Redeem) Use(db *sql.DB) error {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		UPDATE redeem SET used = TRUE WHERE id = ? AND used = FALSE
 	`, r.Id)
 	return err

auth/struct.go

@@ -22,7 +22,7 @@ type User struct {
 
 func GetUserById(db *sql.DB, id int64) *User {
 	var user User
-	if err := db.QueryRow("SELECT id, username FROM auth WHERE id = ?", id).Scan(&user.ID, &user.Username); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT id, username FROM auth WHERE id = ?", id).Scan(&user.ID, &user.Username); err != nil {
 		return nil
 	}
 	return &user
@@ -30,7 +30,7 @@ func GetUserById(db *sql.DB, id int64) *User {
 
 func GetUserByName(db *sql.DB, username string) *User {
 	var user User
-	if err := db.QueryRow("SELECT id, username FROM auth WHERE username = ?", username).Scan(&user.ID, &user.Username); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT id, username FROM auth WHERE username = ?", username).Scan(&user.ID, &user.Username); err != nil {
 		return nil
 	}
 	return &user
@@ -38,7 +38,7 @@ func GetUserByName(db *sql.DB, username string) *User {
 
 func GetUserByEmail(db *sql.DB, email string) *User {
 	var user User
-	if err := db.QueryRow("SELECT id, username FROM auth WHERE email = ?", email).Scan(&user.ID, &user.Username); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT id, username FROM auth WHERE email = ?", email).Scan(&user.ID, &user.Username); err != nil {
 		return nil
 	}
 	return &user
@@ -57,7 +57,7 @@ func (u *User) IsBanned(db *sql.DB) bool {
 	}
 
 	var banned sql.NullBool
-	if err := db.QueryRow("SELECT is_banned FROM auth WHERE username = ?", u.Username).Scan(&banned); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT is_banned FROM auth WHERE username = ?", u.Username).Scan(&banned); err != nil {
 		return false
 	}
 	u.Banned = banned.Valid && banned.Bool
@@ -71,7 +71,7 @@ func (u *User) IsAdmin(db *sql.DB) bool {
 	}
 
 	var admin sql.NullBool
-	if err := db.QueryRow("SELECT is_admin FROM auth WHERE username = ?", u.Username).Scan(&admin); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT is_admin FROM auth WHERE username = ?", u.Username).Scan(&admin); err != nil {
 		return false
 	}
 
@@ -83,7 +83,7 @@ func (u *User) GetID(db *sql.DB) int64 {
 	if u.ID > 0 {
 		return u.ID
 	}
-	if err := db.QueryRow("SELECT id FROM auth WHERE username = ?", u.Username).Scan(&u.ID); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT id FROM auth WHERE username = ?", u.Username).Scan(&u.ID); err != nil {
 		return 0
 	}
 	return u.ID
@@ -99,7 +99,7 @@ func (u *User) GetEmail(db *sql.DB) string {
 	}
 
 	var email sql.NullString
-	if err := db.QueryRow("SELECT email FROM auth WHERE username = ?", u.Username).Scan(&email); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT email FROM auth WHERE username = ?", u.Username).Scan(&email); err != nil {
 		return ""
 	}
 
@@ -109,7 +109,7 @@ func (u *User) GetEmail(db *sql.DB) string {
 
 func IsUserExist(db *sql.DB, username string) bool {
 	var count int
-	if err := db.QueryRow("SELECT COUNT(*) FROM auth WHERE username = ?", username).Scan(&count); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT COUNT(*) FROM auth WHERE username = ?", username).Scan(&count); err != nil {
 		return false
 	}
 	return count > 0
@@ -117,7 +117,7 @@ func IsUserExist(db *sql.DB, username string) bool {
 
 func IsEmailExist(db *sql.DB, email string) bool {
 	var count int
-	if err := db.QueryRow("SELECT COUNT(*) FROM auth WHERE email = ?", email).Scan(&count); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT COUNT(*) FROM auth WHERE email = ?", email).Scan(&count); err != nil {
 		return false
 	}
 	return count > 0
@@ -125,7 +125,7 @@ func IsEmailExist(db *sql.DB, email string) bool {
 
 func getMaxBindId(db *sql.DB) int64 {
 	var max int64
-	if err := db.QueryRow("SELECT MAX(bind_id) FROM auth").Scan(&max); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT MAX(bind_id) FROM auth").Scan(&max); err != nil {
 		return 0
 	}
 	return max

auth/subscription.go

@@ -2,6 +2,7 @@ package auth
 
 import (
 	"chat/channel"
+	"chat/globals"
 	"chat/utils"
 	"database/sql"
 	"errors"
@@ -21,7 +22,7 @@ func (u *User) GetSubscription(db *sql.DB) (time.Time, int) {
 	}
 
 	var expiredAt []uint8
-	if err := db.QueryRow("SELECT expired_at, level FROM subscription WHERE user_id = ?", u.GetID(db)).Scan(&expiredAt, &u.Level); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT expired_at, level FROM subscription WHERE user_id = ?", u.GetID(db)).Scan(&expiredAt, &u.Level); err != nil {
 		return time.Unix(0, 0), 0
 	}
 
@@ -62,7 +63,7 @@ func (u *User) IsEnterprise(db *sql.DB) bool {
 	}
 
 	var enterprise sql.NullBool
-	if err := db.QueryRow("SELECT enterprise FROM subscription WHERE user_id = ?", u.GetID(db)).Scan(&enterprise); err != nil {
+	if err := globals.QueryRowDb(db, "SELECT enterprise FROM subscription WHERE user_id = ?", u.GetID(db)).Scan(&enterprise); err != nil {
 		return false
 	}
 
@@ -81,7 +82,7 @@ func (u *User) AddSubscription(db *sql.DB, month int, level int) bool {
 	}
 	expiredAt := current.AddDate(0, month, 0)
 	date := utils.ConvertSqlTime(expiredAt)
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		INSERT INTO subscription (user_id, expired_at, total_month, level) VALUES (?, ?, ?, ?)
 		ON DUPLICATE KEY UPDATE expired_at = ?, total_month = total_month + ?, level = ?
 	`, u.GetID(db), date, month, level, date, month, level)
@@ -101,7 +102,7 @@ func (u *User) DowngradePlan(db *sql.DB, target int) error {
 	// ceil expired time
 	expiredAt := now.Add(time.Duration(stamp)*time.Second).AddDate(0, 0, -1)
 	date := utils.ConvertSqlTime(expiredAt)
-	_, err := db.Exec("UPDATE subscription SET level = ?, expired_at = ? WHERE user_id = ?", target, date, u.GetID(db))
+	_, err := globals.ExecDb(db, "UPDATE subscription SET level = ?, expired_at = ? WHERE user_id = ?", target, date, u.GetID(db))
 
 	return err
 }
@@ -118,7 +119,7 @@ func (u *User) CountUpgradePrice(db *sql.DB, target int) float32 {
 }
 
 func (u *User) SetSubscriptionLevel(db *sql.DB, level int) bool {
-	_, err := db.Exec("UPDATE subscription SET level = ? WHERE user_id = ?", level, u.GetID(db))
+	_, err := globals.ExecDb(db, "UPDATE subscription SET level = ? WHERE user_id = ?", level, u.GetID(db))
 	return err == nil
 }
 

cli/admin.go

@@ -7,7 +7,7 @@ import (
 )
 
 func UpdateRootCommand(args []string) {
-	db := connection.ConnectMySQL()
+	db := connection.ConnectDatabase()
 	cache := connection.ConnectRedis()
 
 	if len(args) == 0 {

cli/invite.go

@@ -8,7 +8,7 @@ import (
 )
 
 func CreateInvitationCommand(args []string) {
-	db := connection.ConnectMySQL()
+	db := connection.ConnectDatabase()
 
 	var (
 		t     = GetArgString(args, 0)

cli/token.go

@@ -8,7 +8,7 @@ import (
 )
 
 func CreateTokenCommand(args []string) {
-	db := connection.ConnectMySQL()
+	db := connection.ConnectDatabase()
 	id, _ := strconv.Atoi(args[0])
 
 	user := auth.GetUserById(db, int64(id))

connection/database.go

@@ -6,20 +6,32 @@ import (
 	"database/sql"
 	"fmt"
 	_ "github.com/go-sql-driver/mysql"
+	_ "github.com/mattn/go-sqlite3"
 	"github.com/spf13/viper"
 )
 
 var DB *sql.DB
 
 func InitMySQLSafe() *sql.DB {
-	ConnectMySQL()
+	ConnectDatabase()
 
 	// using DB as a global variable to point to the latest db connection
 	MysqlWorker(DB)
 	return DB
 }
 
-func ConnectMySQL() *sql.DB {
+func getConn() *sql.DB {
+	if viper.GetString("mysql.host") == "" {
+		globals.SqliteEngine = true
+		globals.Warn("[connection] mysql host is not set, using sqlite (chatnio.db)")
+		db, err := sql.Open("sqlite3", "chatnio.db")
+		if err != nil {
+			panic(err)
+		}
+
+		return db
+	}
+
 	// connect to MySQL
 	db, err := sql.Open("mysql", fmt.Sprintf(
 		"%s:%s@tcp(%s:%d)/%s",
@@ -29,6 +41,7 @@ func ConnectMySQL() *sql.DB {
 		viper.GetInt("mysql.port"),
 		viper.GetString("mysql.db"),
 	))
+
 	if pingErr := db.Ping(); err != nil || pingErr != nil {
 		errMsg := utils.Multi[string](err != nil, utils.GetError(err), utils.GetError(pingErr)) // err.Error() may contain nil pointer
 		globals.Warn(
@@ -40,11 +53,16 @@ func ConnectMySQL() *sql.DB {
 		utils.Sleep(5000)
 		db.Close()
 
-		return ConnectMySQL()
+		return getConn()
-	} else {
-		globals.Debug(fmt.Sprintf("[connection] connected to mysql server (host: %s)", viper.GetString("mysql.host")))
 	}
 
+	globals.Debug(fmt.Sprintf("[connection] connected to mysql server (host: %s)", viper.GetString("mysql.host")))
+	return db
+}
+
+func ConnectDatabase() *sql.DB {
+	db := getConn()
+
 	db.SetMaxOpenConns(512)
 	db.SetMaxIdleConns(64)
 
@@ -72,7 +90,7 @@ func ConnectMySQL() *sql.DB {
 func InitRootUser(db *sql.DB) {
 	// create root user if totally empty
 	var count int
-	err := db.QueryRow("SELECT COUNT(*) FROM auth").Scan(&count)
+	err := globals.QueryRowDb(db, "SELECT COUNT(*) FROM auth").Scan(&count)
 	if err != nil {
 		globals.Warn(fmt.Sprintf("[service] failed to query user count: %s", err.Error()))
 		return
@@ -80,7 +98,7 @@ func InitRootUser(db *sql.DB) {
 
 	if count == 0 {
 		globals.Debug("[service] no user found, creating root user (username: root, password: chatnio123456, email: root@example.com)")
-		_, err := db.Exec(`
+		_, err := globals.ExecDb(db, `
 			INSERT INTO auth (username, password, email, is_admin, bind_id, token)
 			VALUES (?, ?, ?, ?, ?, ?)
 		`, "root", utils.Sha2Encrypt("chatnio123456"), "root@example.com", true, 0, "root")
@@ -93,7 +111,7 @@ func InitRootUser(db *sql.DB) {
 }
 
 func CreateUserTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS auth (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  bind_id INT UNIQUE,
@@ -113,7 +131,7 @@ func CreateUserTable(db *sql.DB) {
 }
 
 func CreatePackageTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS package (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  user_id INT,
@@ -129,7 +147,7 @@ func CreatePackageTable(db *sql.DB) {
 }
 
 func CreateQuotaTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS quota (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  user_id INT UNIQUE,
@@ -146,7 +164,7 @@ func CreateQuotaTable(db *sql.DB) {
 }
 
 func CreateConversationTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS conversation (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  user_id INT,
@@ -164,7 +182,7 @@ func CreateConversationTable(db *sql.DB) {
 }
 
 func CreateMaskTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS mask (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  user_id INT,
@@ -184,7 +202,7 @@ func CreateMaskTable(db *sql.DB) {
 
 func CreateSharingTable(db *sql.DB) {
 	// refs is an array of message id, separated by comma (-1 means all messages)
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS sharing (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  hash CHAR(32) UNIQUE,
@@ -201,7 +219,7 @@ func CreateSharingTable(db *sql.DB) {
 }
 
 func CreateSubscriptionTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS subscription (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  level INT DEFAULT 1,
@@ -220,7 +238,7 @@ func CreateSubscriptionTable(db *sql.DB) {
 }
 
 func CreateApiKeyTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS apikey (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  user_id INT UNIQUE,
@@ -235,7 +253,7 @@ func CreateApiKeyTable(db *sql.DB) {
 }
 
 func CreateInvitationTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS invitation (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  code VARCHAR(255) UNIQUE,
@@ -255,7 +273,7 @@ func CreateInvitationTable(db *sql.DB) {
 }
 
 func CreateRedeemTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS redeem (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  code VARCHAR(255) UNIQUE,
@@ -271,7 +289,7 @@ func CreateRedeemTable(db *sql.DB) {
 }
 
 func CreateBroadcastTable(db *sql.DB) {
-	_, err := db.Exec(`
+	_, err := globals.ExecDb(db, `
 		CREATE TABLE IF NOT EXISTS broadcast (
 		  id INT PRIMARY KEY AUTO_INCREMENT,
 		  poster_id INT,

connection/db_migration.go

@@ -1,6 +1,7 @@
 package connection
 
 import (
+	"chat/globals"
 	"database/sql"
 	"strings"
 )
@@ -27,10 +28,14 @@ func checkSqlError(_ sql.Result, err error) error {
 }
 
 func execSql(db *sql.DB, sql string, args ...interface{}) error {
-	return checkSqlError(db.Exec(sql, args...))
+	return checkSqlError(globals.ExecDb(db, sql, args...))
 }
 
 func doMigration(db *sql.DB) error {
+	if globals.SqliteEngine {
+		return doSqliteMigration(db)
+	}
+
 	// v3.10 migration
 
 	// update `quota`, `used` field in `quota` table
@@ -54,3 +59,9 @@ func doMigration(db *sql.DB) error {
 
 	return nil
 }
+
+func doSqliteMigration(db *sql.DB) error {
+	// v3.10 added sqlite support, no migration needed before this version
+
+	return nil
+}

connection/worker.go

@@ -12,7 +12,7 @@ func MysqlWorker(db *sql.DB) {
 	go func() {
 		for {
 			if db == nil || db.Ping() != nil {
-				db = ConnectMySQL()
+				db = ConnectDatabase()
 			}
 
 			time.Sleep(tick)

globals/sql.go

@@ -0,0 +1,77 @@
+package globals
+
+import (
+	"database/sql"
+	"regexp"
+	"strings"
+)
+
+var SqliteEngine = false
+
+type batch struct {
+	Old   string
+	New   string
+	Regex bool
+}
+
+func batchReplace(sql string, batch []batch) string {
+	for _, item := range batch {
+		if item.Regex {
+			sql = regexp.MustCompile(item.Old).ReplaceAllString(sql, item.New)
+			continue
+		}
+
+		sql = strings.ReplaceAll(sql, item.Old, item.New)
+	}
+	return sql
+}
+
+func MultiSql(mysqlSql string, sqliteSql string) string {
+	if SqliteEngine {
+		return sqliteSql
+	}
+	return mysqlSql
+}
+
+func PreflightSql(sql string) string {
+	if SqliteEngine {
+		sql = batchReplace(sql, []batch{
+			// KEYWORD REPLACEMENT
+			{`INT `, `INTEGER `, false},
+			{` AUTO_INCREMENT`, ` AUTOINCREMENT`, false},
+			{`DATETIME`, `TEXT`, false},
+			{`DECIMAL`, `REAL`, false},
+			{`MEDIUMTEXT`, `TEXT`, false},
+			{`VARCHAR`, `TEXT`, false},
+
+			// TEXT(65535) -> TEXT, REAL(10,2) -> REAL
+			{`TEXT\(\d+\)`, `TEXT`, true},
+			{`REAL\(\d+,\d+\)`, `REAL`, true},
+
+			// UNIQUE KEY -> UNIQUE
+			{`UNIQUE KEY`, `UNIQUE`, false},
+		})
+	}
+
+	return sql
+}
+
+func ExecDb(db *sql.DB, sql string, args ...interface{}) (sql.Result, error) {
+	sql = PreflightSql(sql)
+	return db.Exec(sql, args...)
+}
+
+func PrepareDb(db *sql.DB, sql string) (*sql.Stmt, error) {
+	sql = PreflightSql(sql)
+	return db.Prepare(sql)
+}
+
+func QueryDb(db *sql.DB, sql string, args ...interface{}) (*sql.Rows, error) {
+	sql = PreflightSql(sql)
+	return db.Query(sql, args...)
+}
+
+func QueryRowDb(db *sql.DB, sql string, args ...interface{}) *sql.Row {
+	sql = PreflightSql(sql)
+	return db.QueryRow(sql, args...)
+}

go.mod

@@ -53,6 +53,7 @@ require (
 	github.com/leodido/go-urn v1.2.4 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
+	github.com/mattn/go-sqlite3 v1.14.22 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect

go.sum

@@ -403,6 +403,8 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky
 github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
 github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
 github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
+github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
 github.com/miekg/dns v1.1.26/go.mod h1:bPDLeHnStXmXAq1m/Ch/hvfNHr14JKNPMBo3VZKjuso=

manager/broadcast/manage.go

@@ -2,6 +2,7 @@ package broadcast
 
 import (
 	"chat/auth"
+	"chat/globals"
 	"chat/utils"
 	"context"
 	"github.com/gin-gonic/gin"
@@ -11,7 +12,7 @@ func createBroadcast(c *gin.Context, user *auth.User, content string) error {
 	db := utils.GetDBFromContext(c)
 	cache := utils.GetCacheFromContext(c)
 
-	if _, err := db.Exec(`INSERT INTO broadcast (poster_id, content) VALUES (?, ?)`, user.GetID(db), content); err != nil {
+	if _, err := globals.ExecDb(db, `INSERT INTO broadcast (poster_id, content) VALUES (?, ?)`, user.GetID(db), content); err != nil {
 		return err
 	}
 
@@ -24,7 +25,7 @@ func getBroadcastList(c *gin.Context) ([]Info, error) {
 	db := utils.GetDBFromContext(c)
 
 	var broadcastList []Info
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT broadcast.id, broadcast.content, auth.username, broadcast.created_at
 		FROM broadcast
 		INNER JOIN auth ON broadcast.poster_id = auth.id

manager/broadcast/view.go

@@ -1,6 +1,7 @@
 package broadcast
 
 import (
+	"chat/globals"
 	"chat/utils"
 	"context"
 	"github.com/gin-gonic/gin"
@@ -18,7 +19,7 @@ func getLatestBroadcast(c *gin.Context) *Broadcast {
 	}
 
 	var broadcast Broadcast
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT id, content FROM broadcast ORDER BY id DESC LIMIT 1;
 	`).Scan(&broadcast.Index, &broadcast.Content); err != nil {
 		return nil

manager/conversation/mask.go

@@ -26,14 +26,14 @@ func (m *Mask) Save(db *sql.DB, user *auth.User) error {
 	userId := user.GetID(db)
 
 	if m.Id == -1 {
-		_, err := db.Exec(
+		_, err := globals.ExecDb(db,
 			"INSERT INTO mask (mask.user_id, avatar, name, description, context) VALUES (?, ?, ?, ?, ?)",
 			userId, m.Avatar, m.Name, m.Description, utils.Marshal(m.Context),
 		)
 		return err
 	}
 
-	_, err := db.Exec(
+	_, err := globals.ExecDb(db,
 		"UPDATE mask SET avatar = ?, name = ?, description = ?, context = ? WHERE id = ? AND user_id = ?",
 		m.Avatar, m.Name, m.Description, utils.Marshal(m.Context), m.Id, userId,
 	)
@@ -41,12 +41,12 @@ func (m *Mask) Save(db *sql.DB, user *auth.User) error {
 }
 
 func (m *Mask) Delete(db *sql.DB, user *auth.User) error {
-	_, err := db.Exec("DELETE FROM mask WHERE id = ? AND user_id = ?", m.Id, user.GetID(db))
+	_, err := globals.ExecDb(db, "DELETE FROM mask WHERE id = ? AND user_id = ?", m.Id, user.GetID(db))
 	return err
 }
 
 func LoadMask(db *sql.DB, user *auth.User) ([]Mask, error) {
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT id, avatar, name, description, context 
 		FROM mask WHERE user_id = ?
 		ORDER BY id DESC

manager/conversation/shared.go

@@ -49,7 +49,7 @@ func ShareConversation(db *sql.DB, user *auth.User, id int64, refs []int) (strin
 		Refs:           refs,
 	})
 
-	if _, err := db.Exec(`
+	if _, err := globals.ExecDb(db, `
 		INSERT INTO sharing (hash, user_id, conversation_id, refs) VALUES (?, ?, ?, ?)
 		ON DUPLICATE KEY UPDATE refs = ?
 	`, hash, user.GetID(db), id, ref, ref); err != nil {
@@ -86,7 +86,7 @@ func ListSharedConversation(db *sql.DB, user *auth.User) []SharedPreviewForm {
 	}
 
 	id := user.GetID(db)
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 		SELECT conversation.conversation_name, conversation.conversation_id, sharing.updated_at, sharing.hash
 		FROM sharing
 		INNER JOIN conversation 
@@ -120,7 +120,7 @@ func DeleteSharedConversation(db *sql.DB, user *auth.User, hash string) error {
 	}
 
 	id := user.GetID(db)
-	if _, err := db.Exec(`
+	if _, err := globals.ExecDb(db, `
 		DELETE FROM sharing WHERE user_id = ? AND hash = ?
 	`, id, hash); err != nil {
 		return err
@@ -136,7 +136,7 @@ func GetSharedConversation(db *sql.DB, hash string) (*SharedForm, error) {
 		ref     string
 		updated []uint8
 	)
-	if err := db.QueryRow(`
+	if err := globals.QueryRowDb(db, `
 		SELECT auth.username, sharing.refs, sharing.updated_at, conversation.conversation_name,
 		       sharing.user_id, sharing.conversation_id
 		FROM sharing

manager/conversation/storage.go

@@ -20,7 +20,7 @@ func (c *Conversation) SaveConversation(db *sql.DB) bool {
 		ON DUPLICATE KEY UPDATE conversation_name = VALUES(conversation_name), data = VALUES(data)
 	`
 
-	stmt, err := db.Prepare(query)
+	stmt, err := globals.PrepareDb(db, query)
 	if err != nil {
 		return false
 	}
@@ -40,7 +40,7 @@ func (c *Conversation) SaveConversation(db *sql.DB) bool {
 }
 func GetConversationLengthByUserID(db *sql.DB, userId int64) int64 {
 	var length int64
-	err := db.QueryRow("SELECT MAX(conversation_id) FROM conversation WHERE user_id = ?", userId).Scan(&length)
+	err := globals.QueryRowDb(db, "SELECT MAX(conversation_id) FROM conversation WHERE user_id = ?", userId).Scan(&length)
 	if err != nil || length < 0 {
 		return 0
 	}
@@ -57,7 +57,7 @@ func LoadConversation(db *sql.DB, userId int64, conversationId int64) *Conversat
 		data  string
 		model interface{}
 	)
-	err := db.QueryRow(`
+	err := globals.QueryRowDb(db, `
 		SELECT conversation_name, model, data FROM conversation
 		WHERE user_id = ? AND conversation_id = ?
 		`, userId, conversationId).Scan(&conversation.Name, &model, &data)
@@ -81,7 +81,7 @@ func LoadConversation(db *sql.DB, userId int64, conversationId int64) *Conversat
 
 func LoadConversationList(db *sql.DB, userId int64) []Conversation {
 	var conversationList []Conversation
-	rows, err := db.Query(`
+	rows, err := globals.QueryDb(db, `
 			SELECT conversation_id, conversation_name FROM conversation WHERE user_id = ? 
 			ORDER BY conversation_id DESC LIMIT 100
 	`, userId)
@@ -108,7 +108,7 @@ func LoadConversationList(db *sql.DB, userId int64) []Conversation {
 }
 
 func (c *Conversation) DeleteConversation(db *sql.DB) bool {
-	_, err := db.Exec("DELETE FROM conversation WHERE user_id = ? AND conversation_id = ?", c.UserID, c.Id)
+	_, err := globals.ExecDb(db, "DELETE FROM conversation WHERE user_id = ? AND conversation_id = ?", c.UserID, c.Id)
 	if err != nil {
 		return false
 	}
@@ -116,6 +116,6 @@ func (c *Conversation) DeleteConversation(db *sql.DB) bool {
 }
 
 func DeleteAllConversations(db *sql.DB, user auth.User) error {
-	_, err := db.Exec("DELETE FROM conversation WHERE user_id = ?", user.GetID(db))
+	_, err := globals.ExecDb(db, "DELETE FROM conversation WHERE user_id = ?", user.GetID(db))
 	return err
 }