From 67cb512eb4c2fd6fb9d3cb9a4018ce32af685dd3 Mon Sep 17 00:00:00 2001 From: Zhang Minghan Date: Fri, 1 Mar 2024 18:46:00 +0800 Subject: [PATCH] feat alpha: support sqlite and connection stack restruct --- admin/analysis.go | 8 +- admin/invitation.go | 7 +- admin/redeem.go | 5 +- admin/user.go | 25 +-- app/src/api/connection.ts | 147 ++++++++++++++++-- app/src/api/conversation.ts | 19 ++- app/src/assets/main.less | 6 +- app/src/assets/pages/home.less | 2 +- .../admin/assemblies/ModelUsageChart.tsx | 2 +- .../components/home/ConversationSegment.tsx | 18 ++- app/src/store/chat.ts | 27 ++++ auth/apikey.go | 7 +- auth/auth.go | 20 +-- auth/invitation.go | 7 +- auth/package.go | 11 +- auth/quota.go | 29 ++-- auth/redeem.go | 7 +- auth/struct.go | 20 +-- auth/subscription.go | 11 +- chatnio.db | Bin 0 -> 110592 bytes cli/admin.go | 2 +- cli/invite.go | 2 +- cli/token.go | 2 +- connection/database.go | 54 ++++--- connection/db_migration.go | 13 +- connection/worker.go | 2 +- globals/sql.go | 77 +++++++++ go.mod | 1 + go.sum | 2 + manager/broadcast/manage.go | 5 +- manager/broadcast/view.go | 3 +- manager/conversation/mask.go | 8 +- manager/conversation/shared.go | 8 +- manager/conversation/storage.go | 12 +- 34 files changed, 429 insertions(+), 140 deletions(-) create mode 100644 chatnio.db create mode 100644 globals/sql.go diff --git a/admin/analysis.go b/admin/analysis.go index 2712a78b..65a4df8f 100644 --- a/admin/analysis.go +++ b/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 <= ?) diff --git a/admin/invitation.go b/admin/invitation.go index b4f9c848..532958d1 100644 --- a/admin/invitation.go +++ b/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) diff --git a/admin/redeem.go b/admin/redeem.go index d2b33361..95809e17 100644 --- a/admin/redeem.go +++ b/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) diff --git a/admin/user.go b/admin/user.go index 9462015d..16d19e69 100644 --- a/admin/user.go +++ b/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 diff --git a/app/src/api/connection.ts b/app/src/api/connection.ts index e7dd99d1..dbf9bd85 100644 --- a/app/src/api/connection.ts +++ b/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 = { - message: data.message, - endpoint: endpoint, - }; + const trace = + this.stack || + JSON.stringify( + { + message: data.message, + endpoint: endpoint, + }, + null, + 2, + ); t && this.triggerCallback({ - message: ` -${t("request-failed")} -\`\`\`json -${JSON.stringify(trace, null, 2)} -\`\`\` - `, + message: `${t("request-failed")}\n\`\`\`json\n${trace}\n\`\`\`\n`, 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.setId(message.conversation); - } - this.callback && this.callback(message); + this.callback && this.callback(this.id, 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); + } +} diff --git a/app/src/api/conversation.ts b/app/src/api/conversation.ts index bf3dd479..501a7f58 100644 --- a/app/src/api/conversation.ts +++ b/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); diff --git a/app/src/assets/main.less b/app/src/assets/main.less index 41f15388..ac34a356 100644 --- a/app/src/assets/main.less +++ b/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; -} \ No newline at end of file +} diff --git a/app/src/assets/pages/home.less b/app/src/assets/pages/home.less index 62643ef1..70618994 100644 --- a/app/src/assets/pages/home.less +++ b/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; diff --git a/app/src/components/admin/assemblies/ModelUsageChart.tsx b/app/src/components/admin/assemblies/ModelUsageChart.tsx index 08869992..8ecd6aa3 100644 --- a/app/src/components/admin/assemblies/ModelUsageChart.tsx +++ b/app/src/components/admin/assemblies/ModelUsageChart.tsx @@ -65,7 +65,7 @@ function ModelUsageChart({ labels, datasets }: ModelChartProps) {
diff --git a/app/src/components/home/ConversationSegment.tsx b/app/src/components/home/ConversationSegment.tsx index 459d3dc3..08d33945 100644 --- a/app/src/components/home/ConversationSegment.tsx +++ b/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 (
{filterMessage(conversation.name)}
-
{conversation.id}
+
+ {loading ? ( + + ) : ( + conversation.id + )} +
{ diff --git a/app/src/store/chat.ts b/app/src/store/chat.ts index 13c6bbef..be08559c 100644 --- a/app/src/store/chat.ts +++ b/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; 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 => 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)); diff --git a/auth/apikey.go b/auth/apikey.go index 1f7bd00c..4ab9ef5e 100644 --- a/auth/apikey.go +++ b/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 diff --git a/auth/auth.go b/auth/auth.go index 344a33ab..f1c99244 100644 --- a/auth/auth.go +++ b/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 } } diff --git a/auth/invitation.go b/auth/invitation.go index e81db5ed..78df2d8a 100644 --- a/auth/invitation.go +++ b/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 diff --git a/auth/package.go b/auth/package.go index 5fdcc930..eb97b4f3 100644 --- a/auth/package.go +++ b/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 } diff --git a/auth/quota.go b/auth/quota.go index b4123fdf..0170e69b 100644 --- a/auth/quota.go +++ b/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("a); err != nil { + if err := globals.QueryRowDb(db, "SELECT quota FROM quota WHERE user_id = ?", u.GetID(db)).Scan("a); 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("a); err != nil { + if err := globals.QueryRowDb(db, "SELECT used FROM quota WHERE user_id = ?", u.GetID(db)).Scan("a); 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 } diff --git a/auth/redeem.go b/auth/redeem.go index 8be84920..b2b22928 100644 --- a/auth/redeem.go +++ b/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 diff --git a/auth/struct.go b/auth/struct.go index 5a844834..d194c7ee 100644 --- a/auth/struct.go +++ b/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 diff --git a/auth/subscription.go b/auth/subscription.go index e78a92fb..5409c9b6 100644 --- a/auth/subscription.go +++ b/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 } diff --git a/chatnio.db b/chatnio.db new file mode 100644 index 0000000000000000000000000000000000000000..e761b6ad26c9c26c8b902164d8bb11e8977a379b GIT binary patch literal 110592 zcmeI)-%{IF9tUuXKmG#&hY+H$TvVA(20Y-5WgCc_=`;=%l*X9En3-m0XEevw0W@IQ zv1CFn(herM*mu~gz34@!574)0Z~7L!Z!bF2J<<_If|^*{1(wX$gDvUkNJskobp9xE zyz#uCTU6SrH*3U_GNGFxK?psTq);d{!roKtef%0`GXuv9Y$|x>yUmV-9_;^f&|Z7f z$ubrHHvVQbdgJSnABP{u{umyM{v7)$@;dt6;P;Wo17z@Scz(RbxqkQ z=hvROyK88@WoJj4n7tboKbsUnx}nhn^W`49tW?sn>dvyt&t8@Jbh7>Cc0{~)Pk5&| zYPzwnTg1}q#>w1-PwyoCVg;?&4bI&gs12_2VDb$vN|^7RUjovMTkM=9tmqvH$-c8M%2EVN6?hQxxgGoxz>J6M@JVR7YNw_u%MtLc+@F*DKQM)JuJ7Wvc;|2WSiNfkS6YD7<~-1k-8IosUsp_;^Dfw9`1RQh2(KbaqZ!@OB+ehvH2QSW#NN zb{IS!fus|oVewP3o3wLQs_abShX=zEacW99T5)U}q`FJCsXMvh^K}yg*9?y^KbKP- zM?koqC3s<;e!#ARZa-i*?8zHpF+bH!$vs zYKNFR($W|C(o80u@C{wBr2qFY+==4xI<+BCxv5iH6&c zq7iXoLU{Ya0jv>o*M7%*baqN$w=+6qpID^nVK7L|J66$Xr#tXuVx|^7;Nv~67x2=& zNbGho?DS}FRY%VofA4(6r|Dhu*HyiH^Idbk6{o3LBs|77Hb&obmZjo)St@Q83Z2pS z^mfBaW**EZtyZ&APo@^+%=CYC6`!f|z5{f}b&HYnx+C8fvUe%^Y%na&&INIVX-MszoN-oayonC7c`Q~R?`bWvik+{g@t*T4RqssMc#*CPi}E7D+8sX(c;f zC8=CBlcU)R$<3#=O!i@xELJOvGR;2Jat{|$*(^y{7YN-V*{u|nb3^RBg!u0v_Ja-UKM}LPPr&d0e+k8ZxyaJ6ZwNpD0uX=z1Rwwb2tWV=5P$##E=yp%Kk3}_ zd;kA{{r~^_*v)_ZSNsdWWyJyIKmY;|fB*y_009U<00Izz00b^lAlBDE8RNet9TC{` zfBT03?fZZCxA`x!XzUvT5P$##AOHafKmY;|fB*y_0D&tMu)qJ0`~NEp0E&hH1Rwwb z2tWV=5P$##AOHafT!H|e|6c+!ihuwFAOHafKmY;|fB*y_009VGu>kJ>uQ&iG9s&@6 z00bZa0SG_<0uX=z1R!t;0{H&_B_N{+2tWV=5P$##AOHafKmY;|fWQ?C;Q9X*2LQ!G z00Izz00bZa0SG_<0uX=z1TH}U&;Kt08AU(<0uX=z1Rwwb2tWV=5P$##u2=xi|F1Xz zC>{b3fB*y_009U<00Izz00ba#2?BWje+kGa0s;_#00bZa0SG_<0uX=z1R!w50(ky^ z#Q{L^5P$##AOHafKmY;|fB*y_0D(&o!1Mo0Kt>S|fB*y_009U<00Izz00bZafh!iU zpZ^Q-e}~u~UJ!r)1Rwwb2tWV=5P$##AOHafT$#W>gm5T!IG!Q);o!LsK(%EDxlT68y zoccu0ev(bw&;Nz^?;-Yw7X%;x0SG_<0uX=z1Rwwb2tWV=AD}>ge=NrD|9=R@fA|1T z5T^tI2tWV=5P$##AOHafKmY;|xaI;Y1Bc_W|G`%R@cVz){6fItLjVF0fB*y_009U< z00Izz00epv!0-R{fP^d%fB*y_009U<00Izz00bZafom>+`~PbmKpZ{RP_^LVBvakn)@5_0=NVa!o0g=bW64x@pm-$~WLi)w*F(!;;F%S8k?ilM;(+D&dLJ z=gM+^vrv|nHcKV8C$+q~rfih+YtQ%&me)(l>Ppcma;8&2LMkcC3d_HwY)GVK?ab&} zf*nJm9Tp?vlQxcYok;RpG$RksK`qhSr;$4KaH+EPggAgmgor2ji|=sOiSOZV^kb8z*xUKE0Flixsq9H#m21pf4G{$)S zSac#y8u~7M%|CDXBFxhQN0*)HaRcDA7~WdKcB!(`&Q_m*_24wEU0_$*+41x00Y6JM z&8N3#Z}pH_=9$Nj#^Mq2{(a#cbtuhN#jG~oms44X50x z8(rn;K|^O%{kYBL8w6AjST`)zBJQa**8K0BIpP#Q0M;+iW}~T_^t?tgP_!3shQ+7% zyHzAFzAAf{#@Y`?BBCe?Zy!0qbQ*N$W7HQu-p~i_7tgB{o~|REos$;4U5C@5cvCMn zWPzj;qhaw=v75AWR;uhw;)e&r5piltI9hRR8>G5Rwy8V0;q!G91J?|XFu%4=bsPcV zdY0gYdHMmn3cCG(-3upggvI<+HzoIMRoR^xX~$z+s>Pd>9n#c|ZIe%id{S;=V5uEq z?np~t3jT_JspjBj&FCj``^9l)!FhbjUujNYle$keYX_qR~!w;K}5ShY$Y3c=#f*+r_Zc zqrFufJ#YNI^AVq>cg zFrTzq%}PC)T97l-yb=e5JyYj>2k737EJn`jj(l6l-lgcX!LT?xEAXq>9le|MWs4eB zdTKH3i^i#yGq%(m+cMFX1YR>fc4BMaGuiJx-_~|Y;22ieErlDZGebVR)^vj9Yi+Fh zE;al?az_O 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...) +} diff --git a/go.mod b/go.mod index 1c584fcb..46da86ac 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index d73051b5..68aa742f 100644 --- a/go.sum +++ b/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= diff --git a/manager/broadcast/manage.go b/manager/broadcast/manage.go index 6436fa90..79737a1e 100644 --- a/manager/broadcast/manage.go +++ b/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 diff --git a/manager/broadcast/view.go b/manager/broadcast/view.go index d2d76682..672f8f84 100644 --- a/manager/broadcast/view.go +++ b/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 diff --git a/manager/conversation/mask.go b/manager/conversation/mask.go index 336af3fb..97c6d24f 100644 --- a/manager/conversation/mask.go +++ b/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 diff --git a/manager/conversation/shared.go b/manager/conversation/shared.go index 7cf365cb..cf249ddf 100644 --- a/manager/conversation/shared.go +++ b/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 diff --git a/manager/conversation/storage.go b/manager/conversation/storage.go index e245b7f1..64c947e7 100644 --- a/manager/conversation/storage.go +++ b/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 }