working after switch to mysql

This commit is contained in:
BennyThink 2021-12-29 15:54:25 +08:00
parent 5f19a51a9b
commit c755f392de
No known key found for this signature in database
GPG Key ID: 6CD0DBDA5235D481
6 changed files with 84 additions and 48 deletions

73
db.py
View File

@ -10,16 +10,16 @@ __author__ = "Benny <benny.think@gmail.com>"
import logging import logging
import os import os
import re import re
import sqlite3
import subprocess import subprocess
import time import time
from io import BytesIO from io import BytesIO
import fakeredis import fakeredis
import pymysql
import redis import redis
from beautifultable import BeautifulTable from beautifultable import BeautifulTable
from config import QUOTA, REDIS from config import MYSQL_HOST, MYSQL_PASS, MYSQL_USER, QUOTA, REDIS
class Redis: class Redis:
@ -76,7 +76,7 @@ class Redis:
def show_usage(self): def show_usage(self):
from downloader import sizeof_fmt from downloader import sizeof_fmt
db = SQLite() db = MySQL()
db.cur.execute("select * from VIP") db.cur.execute("select * from VIP")
data = db.cur.fetchall() data = db.cur.fetchall()
fd = [] fd = []
@ -134,38 +134,45 @@ class Redis:
return file return file
class SQLite: class MySQL:
def __init__(self): vip_sql = """
super(SQLite, self).__init__() create table if not exists VIP
self.con = sqlite3.connect("vip.sqlite", check_same_thread=False) (
self.cur = self.con.cursor() user_id bigint not null,
SQL = """ username varchar(256) null,
create table if not exists VIP payment_amount int null,
( payment_id varchar(256) null,
user_id integer level int default 1 null,
constraint VIP quota bigint default %s null,
primary key, constraint VIP_pk
username varchar(100), primary key (user_id)
payment_amount integer, );
payment_id varchar(100), """ % QUOTA
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)
);
""" settings_sql = """
self.cur.execute(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() self.con.commit()
def __del__(self): def __del__(self):
self.con.close() self.con.close()
if __name__ == '__main__':
db = MySQL()

View File

@ -17,7 +17,7 @@ import requests
from config import (AFD_TOKEN, AFD_USER_ID, COFFEE_TOKEN, ENABLE_VIP, EX, from config import (AFD_TOKEN, AFD_USER_ID, COFFEE_TOKEN, ENABLE_VIP, EX,
MULTIPLY, OWNER, QUOTA, USD2CNY) MULTIPLY, OWNER, QUOTA, USD2CNY)
from db import Redis, SQLite from db import MySQL, Redis
from utils import apply_log_formatter from utils import apply_log_formatter
apply_log_formatter() apply_log_formatter()
@ -32,17 +32,17 @@ def get_username(chat_id):
return data return data
class VIP(Redis, SQLite): class VIP(Redis, MySQL):
def check_vip(self, user_id: "int") -> "tuple": 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() data = self.cur.fetchone()
return data return data
def add_vip(self, user_data: "dict") -> ("bool", "str"): 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 # 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() is_exist = self.cur.fetchone()
if is_exist: if is_exist:
return "Failed. {} is being used by user {}".format(user_data["payment_id"], is_exist[0]) return "Failed. {} is being used by user {}".format(user_data["payment_id"], is_exist[0])

View File

@ -5,6 +5,7 @@ youtube-dl==2021.6.6
APScheduler==3.7.0 APScheduler==3.7.0
beautifultable==1.0.1 beautifultable==1.0.1
ffmpeg-python==0.2.0 ffmpeg-python==0.2.0
PyMySQL==1.0.2
supervisor supervisor
tgbot-ping tgbot-ping

26
tools/migrate_to_mysql.py Normal file
View File

@ -0,0 +1,26 @@
#!/usr/local/bin/python3
# coding: utf-8
# ytdlbot - migrate_to_mysql.py
# 12/29/21 15:28
#
__author__ = "Benny <benny.think@gmail.com>"
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()

View File

@ -9,7 +9,7 @@ __author__ = "Benny <benny.think@gmail.com>"
import logging import logging
from db import SQLite from db import MySQL
def apply_log_formatter(): def apply_log_formatter():
@ -27,9 +27,9 @@ def customize_logger(logger: "list"):
def get_user_settings(user_id: "str") -> "tuple": def get_user_settings(user_id: "str") -> "tuple":
db = SQLite() db = MySQL()
cur = db.cur 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() data = cur.fetchone()
if data is None: if data is None:
return 100, "high", "video" 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"): def set_user_settings(user_id: int, field: "str", value: "str"):
db = SQLite() db = MySQL()
cur = db.cur 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() data = cur.fetchone()
if data is None: if data is None:
resolution = method = "" resolution = method = ""
@ -49,9 +49,9 @@ def set_user_settings(user_id: int, field: "str", value: "str"):
if field == "method": if field == "method":
method = value method = value
resolution = "high" 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: 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() db.con.commit()

View File

@ -24,7 +24,7 @@ from tgbot_ping import get_runtime
from config import (APP_HASH, APP_ID, AUTHORIZED_USER, ENABLE_VIP, OWNER, from config import (APP_HASH, APP_ID, AUTHORIZED_USER, ENABLE_VIP, OWNER,
REQUIRED_MEMBERSHIP, TOKEN, WORKERS) REQUIRED_MEMBERSHIP, TOKEN, WORKERS)
from constant import BotText 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 downloader import convert_flac, sizeof_fmt, upload_hook, ytdl_download
from limit import verify_payment from limit import verify_payment
from utils import customize_logger, get_user_settings, set_user_settings 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): def create_app(session="ytdl", workers=WORKERS):
_app = Client(session, APP_ID, APP_HASH, _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 return _app
@ -290,7 +292,7 @@ def audio_callback(client: "Client", callback_query: types.CallbackQuery):
if __name__ == '__main__': if __name__ == '__main__':
SQLite() MySQL()
scheduler = BackgroundScheduler() scheduler = BackgroundScheduler()
scheduler.add_job(Redis().reset_today, 'cron', hour=0, minute=0) scheduler.add_job(Redis().reset_today, 'cron', hour=0, minute=0)
scheduler.start() scheduler.start()