diff --git a/partybot/pnutbot.py b/partybot/pnutbot.py index 0528c89..13f9a6d 100644 --- a/partybot/pnutbot.py +++ b/partybot/pnutbot.py @@ -9,6 +9,7 @@ import json import random import re import argparse +import os from sqlalchemy import create_engine, and_ from sqlalchemy.orm import scoped_session, sessionmaker @@ -20,6 +21,12 @@ _shutdown = threading.Event() _connected = threading.Event() _error = threading.Event() +logger = logging.getLogger() +config = {} +db_session = None + +SNAP_DATA = os.environ.get('SNAP_DATA') + def subscribe(connection_id): url = f"https://api.pnut.io/v1/channels/{config['CHANNEL']}/messages" url += "?connection_id=" + connection_id @@ -353,7 +360,8 @@ def on_open(ws): t.start() def main(): - logger = logging.getLogger() + global config + global db_session a_parser = argparse.ArgumentParser() a_parser.add_argument( '-d', action='store_true', dest='debug', @@ -363,6 +371,10 @@ def main(): '-c', '--config', default="config.yaml", help="configuration file" ) + a_parser.add_argument( + '-s', '--store', default="store.db", + help="database store url" + ) args = a_parser.parse_args() if args.debug: @@ -370,10 +382,17 @@ def main(): else: logging.basicConfig(level=logging.INFO) - with open(args.config, 'rb') as config_file: + if SNAP_DATA is None: + filename = args.config + db_url = args.store + else: + filename = SNAP_DATA + '/config.yaml' + db_url = 'sqlite:///' + SNAP_DATA + '/store.db' + + with open(filename, 'rb') as config_file: config = yaml.load(config_file, Loader=yaml.SafeLoader) - engine = create_engine(config['SERVICE_DB']) + engine = create_engine(db_url) db_session = scoped_session(sessionmaker(bind=engine)) Base.query = db_session.query_property() Base.metadata.create_all(bind=engine) diff --git a/setup.py b/setup.py index aba9ad5..253f41e 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,7 @@ from setuptools import setup setup( name='partybot', + version='0.2.0', py_modules=['partybot.models','partybot.pnutbot'], install_requires=[ 'PyYAML',