Show global header message when bot loses connection with server

This commit is contained in:
Jeremy Zhang
2017-06-13 01:39:49 +00:00
parent 926d4d14a9
commit 2647753007
8 changed files with 78 additions and 2 deletions

View File

@ -17,6 +17,9 @@ class Titan(discord.Client):
self.http.user_agent += ' TitanEmbeds-Bot'
self.database = DatabaseInterface(self)
self.command = Commands(self, self.database)
self.database_connected = False
self.loop.create_task(self.send_webserver_heartbeat())
def _cleanup(self):
try:
@ -31,6 +34,14 @@ class Titan(discord.Client):
gathered.exception()
except: # Can be ignored
pass
async def send_webserver_heartbeat(self):
await self.wait_until_ready()
while not self.database_connected:
await asyncio.sleep(1) # Wait until db is connected
while not self.is_closed:
await self.database.send_webserver_heartbeat()
await asyncio.sleep(60)
def run(self):
try:
@ -57,6 +68,7 @@ class Titan(discord.Client):
try:
await self.database.connect(config["database-uri"] + "?charset=utf8mb4")
self.database_connected = True
except Exception:
self.logger.error("Unable to connect to specified database!")
traceback.print_exc()

View File

@ -7,6 +7,7 @@ from sqlalchemy.ext.declarative import declarative_base
import json
import discord
import time
Base = declarative_base()
@ -15,6 +16,7 @@ from titanembeds.database.messages import Messages
from titanembeds.database.guild_members import GuildMembers
from titanembeds.database.unauthenticated_users import UnauthenticatedUsers
from titanembeds.database.unauthenticated_bans import UnauthenticatedBans
from titanembeds.database.keyvalue_properties import KeyValueProperties
class DatabaseInterface(object):
# Courtesy of https://github.com/SunDwarf/Jokusoramame
@ -363,3 +365,15 @@ class DatabaseInterface(object):
dbuser.revoked = True
session.commit()
return "Successfully kicked **{}#{}**!".format(dbuser.username, dbuser.discriminator)
async def send_webserver_heartbeat(self):
async with threadpool():
with self.get_session() as session:
key = "bot_heartbeat"
q = session.query(KeyValueProperties).filter(KeyValueProperties.key == key)
if q.count() == 0:
session.add(KeyValueProperties(key=key, value=time.time()))
else:
firstobj = q.first()
firstobj.value = time.time()
session.commit()

View File

@ -0,0 +1,17 @@
from titanembeds.database import db, Base
import datetime
class KeyValueProperties(Base):
__tablename__ = "keyvalue_properties"
id = db.Column(db.Integer, primary_key=True) # Auto incremented id
key = db.Column(db.String(255), nullable=False) # Property Key
value = db.Column(db.Text()) # Property value
expiration = db.Column(db.TIMESTAMP) # Suggested Expiration for value (None = no expire) in secs
def __init__(self, key, value, expiration=None):
self.key = key
self.value = value
if expiration:
self.expiration = datetime.now() + timedelta(seconds = expiration)
else:
self.expiration = None