From c755f392defef09e0baf10aadd13ea23ba970d2f Mon Sep 17 00:00:00 2001 From: BennyThink Date: Wed, 29 Dec 2021 15:54:25 +0800 Subject: [PATCH] working after switch to mysql --- db.py | 73 +++++++++++++++++++++------------------ limit.py | 10 +++--- requirements.txt | 1 + tools/migrate_to_mysql.py | 26 ++++++++++++++ utils.py | 14 ++++---- ytdl.py | 8 +++-- 6 files changed, 84 insertions(+), 48 deletions(-) create mode 100644 tools/migrate_to_mysql.py diff --git a/db.py b/db.py index bd8a208..c60885e 100644 --- a/db.py +++ b/db.py @@ -10,16 +10,16 @@ __author__ = "Benny " import logging import os import re -import sqlite3 import subprocess import time from io import BytesIO import fakeredis +import pymysql import redis from beautifultable import BeautifulTable -from config import QUOTA, REDIS +from config import MYSQL_HOST, MYSQL_PASS, MYSQL_USER, QUOTA, REDIS class Redis: @@ -76,7 +76,7 @@ class Redis: def show_usage(self): from downloader import sizeof_fmt - db = SQLite() + db = MySQL() db.cur.execute("select * from VIP") data = db.cur.fetchall() fd = [] @@ -134,38 +134,45 @@ class Redis: return file -class SQLite: - def __init__(self): - super(SQLite, self).__init__() - self.con = sqlite3.connect("vip.sqlite", check_same_thread=False) - self.cur = self.con.cursor() - SQL = """ - create table if not exists VIP - ( - user_id integer - constraint VIP - primary key, - username varchar(100), - payment_amount integer, - payment_id varchar(100), - level integer default 1, - quota int default %s - ); - """ % QUOTA - self.cur.execute(SQL) - SQL = """ - create table if not exists settings - ( - user_id integer - constraint settings_pk - primary key, - resolution varchar(64), - method varchar(64) - ); +class MySQL: + vip_sql = """ + create table if not exists VIP + ( + user_id bigint not null, + username varchar(256) null, + payment_amount int null, + payment_id varchar(256) null, + level int default 1 null, + quota bigint default %s null, + constraint VIP_pk + primary key (user_id) + ); + """ % QUOTA - """ - self.cur.execute(SQL) + settings_sql = """ + create table if not exists settings + ( + user_id bigint not null, + resolution varchar(128) null, + method varchar(64) null, + constraint settings_pk + primary key (user_id) + ); + """ + + def __init__(self): + self.con = pymysql.connect(host=MYSQL_HOST, user=MYSQL_USER, passwd=MYSQL_PASS, db="vip", charset="utf8mb4") + self.cur = self.con.cursor() + self.init_db() + + def init_db(self): + self.cur.execute(self.vip_sql) + self.cur.execute(self.settings_sql) self.con.commit() def __del__(self): self.con.close() + + +if __name__ == '__main__': + db = MySQL() diff --git a/limit.py b/limit.py index 6841bdf..2d59a4a 100644 --- a/limit.py +++ b/limit.py @@ -17,7 +17,7 @@ import requests from config import (AFD_TOKEN, AFD_USER_ID, COFFEE_TOKEN, ENABLE_VIP, EX, MULTIPLY, OWNER, QUOTA, USD2CNY) -from db import Redis, SQLite +from db import MySQL, Redis from utils import apply_log_formatter apply_log_formatter() @@ -32,17 +32,17 @@ def get_username(chat_id): return data -class VIP(Redis, SQLite): +class VIP(Redis, MySQL): def check_vip(self, user_id: "int") -> "tuple": - self.cur.execute("SELECT * FROM VIP WHERE user_id=?", (user_id,)) + self.cur.execute("SELECT * FROM VIP WHERE user_id=%s", (user_id,)) data = self.cur.fetchone() return data def add_vip(self, user_data: "dict") -> ("bool", "str"): - sql = "INSERT INTO VIP VALUES (?,?,?,?,?,?);" + sql = "INSERT INTO VIP VALUES (%s,%s,%s,%s,%s,%s);" # first select - self.cur.execute("SELECT * FROM VIP WHERE payment_id=?", (user_data["payment_id"],)) + self.cur.execute("SELECT * FROM VIP WHERE payment_id=%s", (user_data["payment_id"],)) is_exist = self.cur.fetchone() if is_exist: return "Failed. {} is being used by user {}".format(user_data["payment_id"], is_exist[0]) diff --git a/requirements.txt b/requirements.txt index ed07849..050c661 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,7 @@ youtube-dl==2021.6.6 APScheduler==3.7.0 beautifultable==1.0.1 ffmpeg-python==0.2.0 +PyMySQL==1.0.2 supervisor tgbot-ping diff --git a/tools/migrate_to_mysql.py b/tools/migrate_to_mysql.py new file mode 100644 index 0000000..30b5f4f --- /dev/null +++ b/tools/migrate_to_mysql.py @@ -0,0 +1,26 @@ +#!/usr/local/bin/python3 +# coding: utf-8 + +# ytdlbot - migrate_to_mysql.py +# 12/29/21 15:28 +# + +__author__ = "Benny " + +import sqlite3 +import pymysql + +mysql_con = pymysql.connect(host='localhost', user='root', passwd='root', db='vip', charset='utf8mb4') +sqlite_con = sqlite3.connect('vip.sqlite') + +vips = sqlite_con.execute('SELECT * FROM VIP').fetchall() + +for vip in vips: + mysql_con.cursor().execute('INSERT INTO vip VALUES (%s, %s, %s, %s, %s, %s)', vip) + +settings = sqlite_con.execute('SELECT * FROM settings').fetchall() + +for setting in settings: + mysql_con.cursor().execute("INSERT INTO settings VALUES (%s,%s,%s)", setting) + +mysql_con.commit() diff --git a/utils.py b/utils.py index 437e070..8f7446a 100644 --- a/utils.py +++ b/utils.py @@ -9,7 +9,7 @@ __author__ = "Benny " import logging -from db import SQLite +from db import MySQL def apply_log_formatter(): @@ -27,9 +27,9 @@ def customize_logger(logger: "list"): def get_user_settings(user_id: "str") -> "tuple": - db = SQLite() + db = MySQL() cur = db.cur - cur.execute("SELECT * FROM settings WHERE user_id = ?", (user_id,)) + cur.execute("SELECT * FROM settings WHERE user_id = %s", (user_id,)) data = cur.fetchone() if data is None: return 100, "high", "video" @@ -37,9 +37,9 @@ def get_user_settings(user_id: "str") -> "tuple": def set_user_settings(user_id: int, field: "str", value: "str"): - db = SQLite() + db = MySQL() cur = db.cur - cur.execute("SELECT * FROM settings WHERE user_id = ?", (user_id,)) + cur.execute("SELECT * FROM settings WHERE user_id = %s", (user_id,)) data = cur.fetchone() if data is None: resolution = method = "" @@ -49,9 +49,9 @@ def set_user_settings(user_id: int, field: "str", value: "str"): if field == "method": method = value resolution = "high" - cur.execute("INSERT INTO settings VALUES (?,?,?)", (user_id, resolution, method)) + cur.execute("INSERT INTO settings VALUES (%s,%s,%s)", (user_id, resolution, method)) else: - cur.execute(f"UPDATE settings SET {field} = ? WHERE user_id = ?", (value, user_id)) + cur.execute(f"UPDATE settings SET {field} =%s WHERE user_id = %s", (value, user_id)) db.con.commit() diff --git a/ytdl.py b/ytdl.py index 3868c21..e919e9a 100644 --- a/ytdl.py +++ b/ytdl.py @@ -24,7 +24,7 @@ from tgbot_ping import get_runtime from config import (APP_HASH, APP_ID, AUTHORIZED_USER, ENABLE_VIP, OWNER, REQUIRED_MEMBERSHIP, TOKEN, WORKERS) from constant import BotText -from db import Redis, SQLite +from db import MySQL, Redis from downloader import convert_flac, sizeof_fmt, upload_hook, ytdl_download from limit import verify_payment from utils import customize_logger, get_user_settings, set_user_settings @@ -32,7 +32,9 @@ from utils import customize_logger, get_user_settings, set_user_settings def create_app(session="ytdl", workers=WORKERS): _app = Client(session, APP_ID, APP_HASH, - bot_token=TOKEN, workers=workers) + bot_token=TOKEN, workers=workers, + proxy={'hostname': '127.0.0.1', 'port': 1086} + ) return _app @@ -290,7 +292,7 @@ def audio_callback(client: "Client", callback_query: types.CallbackQuery): if __name__ == '__main__': - SQLite() + MySQL() scheduler = BackgroundScheduler() scheduler.add_job(Redis().reset_today, 'cron', hour=0, minute=0) scheduler.start()