83 lines
2.5 KiB
Python
83 lines
2.5 KiB
Python
# import sqlite3
|
|
import pymysql
|
|
import click
|
|
from flask import current_app, g
|
|
from werkzeug.security import generate_password_hash
|
|
import os
|
|
|
|
def get_db():
|
|
if 'db' not in g:
|
|
# g.db = sqlite3.connect(
|
|
# current_app.config['DATABASE'],
|
|
# detect_types=sqlite3.PARSE_DECLTYPES
|
|
# )
|
|
# g.db.row_factory = sqlite3.Row
|
|
g.db = pymysql.connect(
|
|
host=current_app.config['DATABASE_HOST'],
|
|
user=current_app.config['DATABASE_USER'],
|
|
password=current_app.config['DATABASE_PASS'],
|
|
database=current_app.config['DATABASE'], cursorclass=pymysql.cursors.DictCursor)
|
|
|
|
return g.db
|
|
|
|
|
|
def close_db(e=None):
|
|
db = g.pop('db', None)
|
|
|
|
if db is not None:
|
|
db.close()
|
|
|
|
def parse_sql(filename):
|
|
data = open(filename, 'r').readlines()
|
|
stmts = []
|
|
DELIMITER = ';'
|
|
stmt = ''
|
|
for lineno, line in enumerate(data):
|
|
if not line.strip():
|
|
continue
|
|
if line.startswith('--'):
|
|
continue
|
|
if 'DELIMITER' in line:
|
|
DELIMITER = line.split()[1]
|
|
continue
|
|
if (DELIMITER not in line):
|
|
stmt += line.replace(DELIMITER, ';')
|
|
continue
|
|
if stmt:
|
|
stmt += line
|
|
stmts.append(stmt.strip())
|
|
stmt = ''
|
|
else:
|
|
stmts.append(line.strip())
|
|
return stmts
|
|
|
|
@click.command('init-db')
|
|
def init_db():
|
|
scriptpath = os.path.split(os.path.realpath(__file__))[0]
|
|
db = get_db()
|
|
cur = db.cursor()
|
|
click.echo("PyMySQL do not support script execution...")
|
|
click.echo("Initializing database schema...")
|
|
sqls = parse_sql(os.path.join(scriptpath,"./sql/db_create.sql"))
|
|
for sql in sqls:
|
|
cur.execute(sql)
|
|
db.commit()
|
|
# with current_app.open_resource('sql/db_create_sqlite.sql') as f:
|
|
# db.executescript(f.read().decode('utf8'))
|
|
click.echo("Initializing database triggers...")
|
|
sqls = parse_sql(os.path.join(scriptpath,"./sql/trigger.sql"))
|
|
for sql in sqls:
|
|
cur.execute(sql)
|
|
db.commit()
|
|
# with current_app.open_resource('sql/trigger_sqlite.sql') as f:
|
|
# db.executescript(f.read().decode('utf8'))
|
|
click.echo("Initializing admin passwd...")
|
|
print(len(generate_password_hash("lolicon")))
|
|
cur.execute("insert into admin (passwd) values (%s)", (generate_password_hash(current_app.config["ADMIN_PASS"])))
|
|
db.commit()
|
|
cur.close()
|
|
click.echo("Initialized database.")
|
|
|
|
def init_app(app):
|
|
app.teardown_appcontext(close_db)
|
|
app.cli.add_command(init_db) |