import pymysql

from vector_test.vector_util import MyVector


class Database:
    def __init__(self, host, user, password, database):
        self.connection = pymysql.connect(
            host=host,
            user=user,
            password=password,
            database=database
        )

    def query(self, sql):
        with self.connection.cursor() as cursor:
            cursor.execute(sql)
            result = cursor.fetchall()
        return result

    def execute(self, sql):
        with self.connection.cursor() as cursor:
            cursor.execute(sql)
        self.connection.commit()


db = Database('www.tianyunperfect.cn', 'memos', 's3ijTsaH5cciP8s8', 'memos')
client = MyVector("101.43.195.48", port=6333, grpc_port=6334, api_key="tianyunperfect123456", embedding_path='/Users/alvin/IdeaProjects/python-base/vector_test/models--shibing624--text2vec-base-chinese')
# 查询
result = db.query("SELECT id,content,resource_name FROM memo where row_status = 'NORMAL' and creator_id=1")
# 遍历result
for row in result:
    client.create_document(collection_name="memo_collection", doc_id=row[0], vector=client.model.encode(row[1]), payload={"content": row[1], "resource_name": row[2]})