def sqlite_format(items, filepath):
db = SqliteDatabase(filepath.as_posix())
db_proxy.initialize(db)
db.create_tables((Hotel, ))
with db.atomic():
Hotel.delete().execute()
for batch in chunked(items, 900):
Hotel.insert_many(batch).execute()
def _db_store_enc_stream_in(self, file_entry: Files, stream_in: BinaryIO,
db: SqliteDatabase) -> None:
""" Store stream as encrypted chunks in Chunks table."""
n_chunks = 0
with db.atomic():
for chunk_group in self._chunked(
self._generate_enc_chunks(stream_in, file_entry.file_id),
20):
Chunks.bulk_create(chunk_group)
n_chunks += len(chunk_group)
file_entry.n_chunks = n_chunks
file_entry.save()
def load_tags_db():
'''
load user tags data from uploaded db file
Args:
file: io.BufferedRandom -> uploaded db file stream
'''
db_name = get_data_path('uploaded.db')
try:
db_upload = SqliteDatabase(db_name)
db_upload.get_tables()
except DatabaseError:
raise DBError()
db_is_old = False
tag_data = []
missed_fanhaos = []
tag_file_added = 0
sql_old = '''select item_rate.rate_value, item.fanhao
from item_rate inner
join item on item_rate.item_id = item.id
where item_rate.rate_type=1 '''
sql_new = '''select item_rate.rate_value, item.fanhao
from item_rate inner
join item on item_rate.item_id = item.fanhao
where item_rate.rate_type=1 '''
cursor = db_upload.execute_sql(sql_old)
res = cursor.fetchone()
if res:
db_is_old = True
if db_is_old:
cursor = db_upload.execute_sql(sql_old)
else:
cursor = db_upload.execute_sql(sql_new)
for row in cursor.fetchall():
tag_data.append(row)
with db_upload.atomic():
for rate_value, fanhao in tag_data:
item_rate = ItemRate.saveit(RATE_TYPE.USER_RATE, rate_value,
fanhao)
if item_rate:
tag_file_added += 1
if not Item.get_by_fanhao(fanhao):
# add to get from spider
missed_fanhaos.append(fanhao)
logger.debug(tag_data)
logger.info(f'added user tag rate: {tag_file_added}')
logger.info(f'added fanhao to download: {len(missed_fanhaos)}')
return tag_file_added, missed_fanhaos
def create_bulk_ratings(rows):
fields = [Rating.udemy_id, Rating.created, Rating.rating_score]
count = len(rows)
print('creating {} rows'.format(count))
database = SqliteDatabase('ratings.db')
with database.atomic():
step_range = 250 #scale down as needed
for step in range(0, count, step_range):
Rating.insert_many(
rows[step:step+step_range],
fields=fields
).execute()
print('done creating rows')
class DatabaseHandler(Handler):
TIME_FORMAT = "%Y.%m.%d. %H:%M:%S"
_persons_select: Select = None
_unknown_persons_select: Select = None
_known_persons_select: Select = None
_images_select: Select = None
_encodings_select: Select = None
valid = False
database: SqliteDatabase = None
def __init__(self, app, db_location):
super().__init__(app)
DBModel._handler = self
self.database = SqliteDatabase(db_location,
pragmas=(('foreign_keys', 'on'), ))
self.database.connect()
self.init_tables([UserEvent, Encoding, Person, Image, User])
# db.close()
self.refresh()
def init_tables(self, tables: List[DBModel]):
# for some reason, changing the meta database on the DBModel doesn't inherit to the model implementations
# so we have to set it on each table manually
for table in tables:
table._meta.database = self.database
self.database.create_tables(tables)
def add_person(self,
name: str,
unknown: bool = True,
thumbnail: Image = None) -> Person:
new_person = Person()
new_person.name = name
new_person.first_seen = datetime.now()
new_person.last_seen = new_person.first_seen
new_person.unknown = unknown
new_person.thumbnail = thumbnail
with self.database.atomic():
new_person.save()
self.refresh()
logging.info("A new {} person was added with the name {}".format(
'unkown' if unknown else 'known', name))
return new_person
def get_person_by_name(self, name: str) -> Person:
return Person.get(Person.name == name)
def get_user_by_name(self, name: str) -> User:
return User.get(User.name == name)
def get_image_by_name(self, name: str) -> Image:
return Image.get(Image.name == name)
def get_all_events(self) -> List:
return UserEvent.select()
def log_event(self, text: str):
new_event = UserEvent.create(datetime=datetime.now(), event=text)
new_event.save()
def invalidate(self):
self.valid = False
def refresh(self):
self._persons_select = Person.select()
# peewee doesn't work with an equality check in the format of Person.unknown is False
# instead we should use Person.unknown == False
self._known_persons_select = Person.select().where(
Person.unknown == False) # NOQA
self._unknown_persons_select = Person.select().where(
Person.unknown == True) # NOQA
self._images_select = Image.select()
self._encodings_select = Encoding.select()
self.valid = True
def get_persons(self) -> List[Person]:
"""
Returns all persons, known and unknown both
"""
logging.info("Getting all persons")
if not self.valid:
self.refresh()
return prefetch(self._persons_select, self._images_select,
self._encodings_select)
def get_known_persons(self) -> List[Person]:
logging.info("Getting known persons")
if not self.valid:
self.refresh()
return prefetch(self._known_persons_select, self._images_select,
self._encodings_select)
def get_unknown_persons(self) -> List[Person]:
logging.info("Getting unknown persons")
if not self.valid:
self.refresh()
return prefetch(self._unknown_persons_select, self._images_select,
self._encodings_select)
# loads the face_recognition_settings table from the database into a dictionary
def load_face_recognition_settings(self):
with open(Path("Data/FaceRecSettings.json")) as json_file:
sett = json.load(json_file)
return sett
# loads the notification_settings table from the database into a dictionary
def load_notification_settings(self):
with open(Path("Data/NotificationSettings.json")) as json_file:
sett = json.load(json_file)
return sett
def chaperone_internal(item, fun, db, db_fun, save, force=False):
fname = fun.__name__
meta = None
if db_fun == 'sqlite':
db = SqliteDatabase(db)
database_proxy.initialize(db)
if db_fun == 'postgres':
db = PostgresqlDatabase(db['dbname'],
user=db['user'],
password=db['password'],
host=db['host'],
port=db['port'],
autorollback=True,
autocommit=True)
database_proxy.initialize(db)
db.connect()
try:
with db.atomic():
# Here we create an instance instead of a record
# If save=False, we need not create a record in the
# db
if save is False:
check = QueueCheck(item=item, complete=False, function=fname)
else:
check = QueueCheck.create(item=item,
complete=False,
function=fname)
except IntegrityError:
check = QueueCheck.get(QueueCheck.item == item,
QueueCheck.function == fname)
try:
if check.complete is False or force is True:
meta = fun(item)
if save is True:
# Data is expected to be in a pandas data frame or
# a dictionary, in both cases ready to be converted
# to json. We use json to take advantage of postgres
check.meta = dumps(meta)
check.complete = True
check.error = '' # Reset error message for posterity
check.save()
else:
return meta
except Exception as e:
if save is True:
check.item = item
check.complete = False
check.error = str(e)
check.save()
else:
return e
db.close()
class MsgDb:
def __init__(self, name, sql_dir = 'db/'):
self.name = str(name)
self.SQL_DIR = sql_dir
# 连接数据库
db_name = '{:s}.db'.format(self.name)
path = os.path.join(self.SQL_DIR, db_name)
self.db = SqliteDatabase(path, pragmas=(
('cache_size', -1024 * 64), # 64MB page-cache.
('journal_mode', 'wal'))) # Use WAL-mode (you should always use this!).
self.db.connect()
self.db.bind(MODELS)
self.db.create_tables(MODELS)
def close(self):
self.db.close()
def __user_medal(self, msg):
'''
格式化信息中的牌子
'''
if msg['cmd'] == 'DANMU_MSG':
if msg['info'][3]:
if msg['info'][3][11] == 1:
return {
'n': msg['info'][3][1], #牌子名
'l': msg['info'][3][0], #牌子等级
'c': msg['info'][3][10], #舰长类别
}
else:
if msg['data']['medal_info']:
if msg['data']['medal_info']['is_lighted'] == 1:
return {
'n': msg['data']['medal_info']['medal_name'], #牌子名
'l': msg['data']['medal_info']['medal_level'], #牌子等级
'c': msg['data']['medal_info']['guard_level'], #舰长类别
}
return {
'n': None, #牌子名
'l': None, #牌子等级
'c': None #舰长类别
}
def __insert_msg(self, msg, insert_list):
'''
向列表插入一条信息
识别并格式化
'''
if msg['cmd'] == 'DANMU_MSG':
medal = self.__user_medal(msg)
insert_list[msg['cmd']].append({
'captain': medal['c'], #舰长类别
'content': msg['info'][1], #弹幕文本
'medal_level': medal['l'], #牌子等级
'medal_name': medal['n'], #牌子名
'timestamp': msg['info'][0][4], #时间戳
'uid': msg['info'][2][0], #用户id
'username': msg['info'][2][1], #用户名
})
elif msg['cmd'] == 'SEND_GIFT':
medal = self.__user_medal(msg)
insert_list[msg['cmd']].append({
'captain': medal['c'], #舰长类别
'gift_coin_type': msg['data']['coin_type'], #货币类型
'gift_name': msg['data']['giftName'], #礼物名
'gift_num': msg['data']['num'], #礼物数量
'gift_price': msg['data']['price'], #礼物单价
'gift_total_price': msg['data']['total_coin'], #礼物总价
'medal_level': medal['l'], #牌子等级
'medal_name': medal['n'], #牌子名
'timestamp': msg['data']['timestamp']*1000, #时间戳
'uid': msg['data']['uid'], #用户id
'username': msg['data']['uname'], #用户名
})
elif msg['cmd'] == 'USER_TOAST_MSG':
insert_list[msg['cmd']].append({
'captain': msg['data']['guard_level'], #大航海类别
'captain_num': msg['data']['num'], #数量
'captain_total_price': msg['data']['price'], #真实花费价格（RMB）
'timestamp': msg['data']['start_time']*1000, #时间戳
'uid': msg['data']['uid'], #用户id
'username': msg['data']['username'], #用户名
})
elif msg['cmd'] == 'SUPER_CHAT_MESSAGE':
medal = self.__user_medal(msg)
insert_list[msg['cmd']].append({
'captain': medal['c'], #舰长类别
'medal_level': medal['l'], #牌子等级
'medal_name': medal['n'], #牌子名
'superchat_content': msg['data']['message'], #醒目留言文本
'superchat_price': msg['data']['price'], #价格（RMB）
'timestamp': msg['data']['ts']*1000, #时间戳
'uid': msg['data']['uid'], #用户id
'username': msg['data']['user_info']['uname'], #用户名
})
elif msg['cmd'] == 'VIEW':
insert_list[msg['cmd']].append({
'timestamp': msg['timestamp'], #时间戳
'view': msg['data'], #人气值
})
elif msg['cmd'] == 'WATCHED_CHANGE':
insert_list[msg['cmd']].append({
'timestamp': msg['timestamp'], #时间戳
'watched': msg['data']['num'], #看过的人数
})
else:
return None
def commit_msgs(self, msgs):
'''
向数据库提交信息列表
'''
# 绑定model
self.db.bind(MODELS)
# 列表分类
insert_list = {}
for t in MSG_TYPE:
insert_list[t] = []
# 创建插入列表
for msg in msgs:
self.__insert_msg(msg, insert_list)
# 插入数据
with self.db.atomic():
Danmu.insert_many(insert_list['DANMU_MSG']).execute()
Gift.insert_many(insert_list['SEND_GIFT']).execute()
Captain.insert_many(insert_list['USER_TOAST_MSG']).execute()
SuperChat.insert_many(insert_list['SUPER_CHAT_MESSAGE']).execute()
View.insert_many(insert_list['VIEW']).execute()
WatchedChange.insert_many(insert_list['WATCHED_CHANGE']).execute()
def live_action(self, timestamp, action):
'''
向数据库提交直播状态（开始、结束）
'''
# 绑定model
self.db.bind(MODELS)
# 插入数据
LiveStatus.insert(timestamp=timestamp, action=action).execute()
def get_live_time(self, date):
'''
date为datetime.date对象
获取当天直播开始和结束时间列表
将北京时间第一天4:00到第二天4:00作为分界线
若不存在，则返回[[None, None]]
'''
t = datetime.time(4, 0, 0, 0)
start_t = datetime.datetime.combine(date, t)
dt = datetime.timedelta(days=1)
time_range = [int(start_t.timestamp()*1000), int((start_t+dt).timestamp()*1000)]
data = (LiveStatus
.select()
.where(LiveStatus.timestamp>=time_range[0], LiveStatus.timestamp<=time_range[1])
.order_by(LiveStatus.timestamp))
if len(data) == 0:
return [[None, None]]
# 核心逻辑：
# 第一个判断是start还是end，end就把开头设为start
# 最后一个判断是start还是end，是start且边界在当前时间之前（说明不是没播完）就把结尾设为end
# 上一个是start，下一个就去找end，如果是start就无视掉，反之同理
action_list = list(data.dicts())
if action_list[0]["action"] == "end":
action_list.insert(0, {"action": "start", "timestamp": time_range[0]})
if action_list[-1]["action"] == "start" and time_range[1]<time.time()*1000:
action_list.append({"action": "end", "timestamp": time_range[1]})
rst = []
last = "end"
for d in action_list:
if d["action"] == last:
continue
if d["action"] == "start":
rst.append([d["timestamp"], None])
last = "start"
elif d["action"] == "end":
rst[-1][1] = d["timestamp"]
last = "end"
# 筛选每个时间段 太短（小于1min）不要
rst_filtered = []
for x in rst:
if x[-1] != None:
if x[1]-x[0]<60*1000:
continue
rst_filtered.append(x)
return rst_filtered
def query_table_on_live(self, model, date, idx):
'''
取出当天在开播时间段内的一个表格
若没有结束时间就到最新的为止
date为datetime.date对象
idx为直播次序编号
'''
start, end = self.get_live_time(date)[idx]
# 如果没有结束时间就持续到下一天4点
if end == None:
dt = datetime.timedelta(days=1)
end = int(datetime.datetime.combine(date+dt, datetime.time(4, 0, 0, 0)).timestamp()*1000)
# 绑定model
self.db.bind(MODELS)
data = None
if start:
if end:
data = model.select().where(model.timestamp >= start, model.timestamp <= end).dicts()
else:
data = model.select().where(model.timestamp >= start).dicts()
return list(data)
def query_msg(self, date, idx):
'''
取出本数据库在开播时间段内的所有弹幕和sc
date为datetime.date对象
idx为今天直播次序编号 从0开始
'''
return (
self.query_table_on_live(Danmu, date, idx),
self.query_table_on_live(SuperChat, date, idx)
)
#!/usr/bin/env python
import os
from datetime import datetime
from peewee import DateTimeField, IntegerField, SqliteDatabase
from playhouse.migrate import SqliteMigrator, migrate
db = SqliteDatabase(os.getenv("DATABASE_FILE", "omegabot.db"))
migrator = SqliteMigrator(db)
with db.atomic():
migrate(
migrator.add_column("user", "xp", IntegerField(default=0)),
migrator.add_column("user", "xp_last_update_time",
DateTimeField(default=datetime.now)),
)