crud.py 6.56 KB
Newer Older
1 2 3 4 5 6
from .c_db import UPostgresDB

TABLE_CHAT = """
DROP TABLE IF EXISTS "chat";
CREATE TABLE chat (
    chat_id varchar(1000) PRIMARY KEY,
7
    user_id varchar(1000),
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108
    info text,
    create_time timestamp(6) DEFAULT current_timestamp,
    deleted int2
);
COMMENT ON COLUMN "chat"."chat_id" IS '会话id';
COMMENT ON COLUMN "chat"."user_id" IS '会话创建用户id';
COMMENT ON COLUMN "chat"."info" IS '会话简介';
COMMENT ON COLUMN "chat"."create_time" IS '会话创建时间,默认为当前时间';
COMMENT ON COLUMN "chat"."deleted" IS '是否删除:0=否,1=是';
COMMENT ON TABLE "chat" IS '会话信息表';
DROP SEQUENCE IF EXISTS "chat_seq";
CREATE SEQUENCE "chat_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE chat ALTER COLUMN chat_id SET DEFAULT nextval('chat_seq'::regclass);
"""

TABLE_TURN_QA = """
DROP TABLE IF EXISTS "turn_qa";
CREATE TABLE turn_qa (
    turn_id varchar(1000) PRIMARY KEY,
    chat_id varchar(1000),
    question text,
    answer text,
    create_time timestamp(6) DEFAULT current_timestamp,
    turn_number int,
    is_last int2
);
COMMENT ON COLUMN "turn_qa"."turn_id" IS '会话轮次id';
COMMENT ON COLUMN "turn_qa"."chat_id" IS '会话id';
COMMENT ON COLUMN "turn_qa"."question" IS '该轮会话问题';
COMMENT ON COLUMN "turn_qa"."answer" IS '该轮会话答案';
COMMENT ON COLUMN "turn_qa"."create_time" IS '该轮会话创建时间,默认为当前时间';
COMMENT ON COLUMN "turn_qa"."turn_number" IS '会话轮数';
COMMENT ON COLUMN "turn_qa"."is_last" IS '是否为最后一轮对话:0=否,1=是';
COMMENT ON TABLE "turn_qa" IS '会话轮次信息表';
DROP SEQUENCE IF EXISTS "turn_qa_seq";
CREATE SEQUENCE "turn_qa_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE turn_qa ALTER COLUMN turn_id SET DEFAULT nextval('turn_qa_seq'::regclass);
"""

TABLE_USER = """
DROP TABLE IF EXISTS "c_user";
CREATE TABLE c_user (
    user_id varchar(1000) PRIMARY KEY,
    account varchar(20) NOT NULL,
    password varchar(50) NOT NULL
);
COMMENT ON COLUMN "c_user"."user_id" IS '用户id';
COMMENT ON COLUMN "c_user"."account" IS '用户帐户';
COMMENT ON COLUMN "c_user"."password" IS '用户密码';
COMMENT ON TABLE "c_user" IS '用户表';
DROP SEQUENCE IF EXISTS "c_user_seq";
CREATE SEQUENCE "c_user_seq"
INCREMENT 1
MINVALUE 1
MAXVALUE 2147483647
START 1
CACHE 1;
ALTER TABLE c_user ALTER COLUMN user_id SET DEFAULT nextval('c_user_seq'::regclass);
"""


class CRUD:
    def __init__(self, _db: UPostgresDB):
        self.db = _db

    def create_table(self):
        self.db.execute(TABLE_CHAT)
        self.db.execute(TABLE_TURN_QA)
        self.db.execute(TABLE_USER)

    def get_history(self, _chat_id):
        query = f'SELECT question,answer FROM turn_qa WHERE chat_id=(%s) ORDER BY turn_number ASC'
        self.db.execute_args(query, (_chat_id,))
        ans = self.db.fetchall()
        return ans

    def insert_turn_qa(self, chat_id, question, answer, turn_number, is_last):
        query = f'INSERT INTO turn_qa(chat_id, question, answer, turn_number, is_last) VALUES (%s,%s,%s,%s,%s)'
        self.db.execute_args(query, (chat_id, question, answer, turn_number, is_last))

    def insert_c_user(self, account, password):
        query = f'INSERT INTO c_user(account, password) VALUES (%s,%s)'
        self.db.execute_args(query, (account, password))

    def insert_chat(self, user_id, info, deleted):
        query = f'INSERT INTO chat(user_id, info, deleted) VALUES (%s,%s,%s)'
        self.db.execute_args(query, (user_id, info, deleted))

    def update_last(self, chat_id):
        query = f'UPDATE turn_qa SET is_last = 0 WHERE chat_id = (%s) AND is_last = 1'
        self.db.execute_args(query, (chat_id,))
陈正乐 committed
109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141

    def user_exist_id(self, _user_id):
        query = f'SELECT * FROM c_user WHERE user_id = (%s)'
        self.db.execute_args(query, (_user_id,))
        return self.db.fetchone()

    def user_exist_account(self, _account):
        query = f'SELECT * FROM c_user WHERE account = (%s)'
        self.db.execute_args(query, (_account,))
        return self.db.fetchone()

    def user_exist_account_password(self, _account, _password):
        query = f'SELECT user_id FROM c_user WHERE account = (%s) AND password = (%s)'
        self.db.execute_args(query, (_account, _password))
        return self.db.fetchone()

    def chat_exist_chatid_userid(self, _chat_id, _user_id):
        query = f'SELECT * FROM chat WHERE chat_id = (%s) AND user_id = (%s)'
        self.db.execute_args(query, (_chat_id, _user_id))
        return self.db.fetchone()

    def get_chat_list_userid(self, _user_id):
        query = f'SELECT info FROM chat WHERE user_id = (%s) AND deleted = 0 order by create_time desc'
        self.db.execute_args(query, (_user_id,))
        return self.db.fetchall()

    def get_chatinfo_from_chatid(self, _chat_id):
        query = f'SELECT info FROM chat WHERE chat_id = (%s)'
        self.db.execute_args(query, (_chat_id,))
        return self.db.fetchone()

    def delete_chat(self, _chat_id):
        query = f'UPDATE chat SET deleted = 1 WHERE chat_id = (%s)'
陈正乐 committed
142 143 144 145 146 147
        self.db.execute_args(query, (_chat_id,))

    def get_last_question(self, _chat_id):
        query = f'SELECT question FROM turn_qa WHERE chat_id = (%s) AND turn_number = 1'
        self.db.execute_args(query, (_chat_id,))
        return self.db.fetchone()[0]
148 149 150 151 152 153 154

    def get_users(self):
        query = f'SELECT account FROM c_user'
        self.db.execute(query)
        return self.db.fetchall()

    def get_chats(self, account):
陈正乐 committed
155
        query = f'SELECT chat.chat_id,chat.info FROM chat JOIN c_user ON chat.user_id = c_user.user_id WHERE c_user.account = (%s) ORDER BY chat.create_time DESC;'
156 157
        self.db.execute_args(query, (account,))
        return self.db.fetchall()
陈正乐 committed
158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182

    def create_chat(self, user_id, info, deleted):
        query = f'INSERT INTO chat(user_id, info, deleted) VALUES (%s,%s,%s) RETURNING chat_id'
        self.db.execute_args(query, (user_id, info, deleted))
        ans = self.db.fetchall()[0][0]
        return ans

    def get_uersid_from_account(self, account):
        query = f'SELECT user_id FROM c_user WHERE account = (%s)'
        self.db.execute_args(query, (account, ))
        ans = self.db.fetchall()[0][0]
        print(ans)
        return ans

    def get_chat_info(self, chat_id):
        query = f'SELECT info FROM chat WHERE chat_id = (%s)'
        self.db.execute_args(query, (chat_id,))
        ans = self.db.fetchall()[0][0]
        print(ans)
        return ans

    def set_info(self, chat_id, info):
        query = f'UPDATE chat SET info = (%s) WHERE chat_id = (%s)'
        self.db.execute_args(query, (info, chat_id))