xtaodada/ytdlbot
1
0
Fork 0
mirror of https://github.com/omg-xtao/ytdlbot.git synced 2025-01-30 09:48:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

working after switch to mysql

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
db.py
View File

@ -10,16 +10,16 @@ __author__ = "Benny <benny.think@gmail.com>"
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()

10
limit.py
View File

@ -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])

1
requirements.txt
View File

@ -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

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()

14
utils.py
View File

@ -9,7 +9,7 @@ __author__ = "Benny <benny.think@gmail.com>"
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()

8
ytdl.py
View File

@ -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()