diff --git a/core/bot.py b/core/bot.py index 60f273a..4f46c72 100644 --- a/core/bot.py +++ b/core/bot.py @@ -5,6 +5,9 @@ import os import Queue import collections +from sqlalchemy.orm import scoped_session, sessionmaker +from sqlalchemy import create_engine + from core import config, irc, main, loader @@ -97,6 +100,11 @@ class Bot(object): self.config = config.Config(self.logger) self.logger.debug("Config object created.") + # db + engine = create_engine('sqlite:///cloudbot.db') + db_factory = sessionmaker(bind=engine) + self.db_session = scoped_session(db_factory) + def connect(self): """connect to all the networks defined in the bot config""" diff --git a/core/main.py b/core/main.py index 5973e65..654933a 100755 --- a/core/main.py +++ b/core/main.py @@ -3,6 +3,7 @@ import traceback import Queue import re +from sqlalchemy.orm import scoped_session thread.stack_size(1024 * 512) # reduce vm size @@ -55,12 +56,15 @@ class Input(dict): def run(func, input): args = func._args + uses_db = 'db' in args and 'db' not in input + if 'inp' not in input: input.inp = input.paraml if args: - if 'db' in args and 'db' not in input: - input.db = get_db_connection(input.conn) + if uses_db: + # create SQLAlchemy session + input.db = input.bot.db_session() if 'input' in args: input.input = input if 0 in args: @@ -73,6 +77,10 @@ def run(func, input): if out is not None: input.reply(unicode(out)) + if uses_db: + # close SQLAlchemy session + input.db.close() + def do_sieve(sieve, bot, input, func, type, args): try: @@ -94,7 +102,6 @@ class Handler(object): def start(self): uses_db = 'db' in self.func._args - db_conns = {} while True: input = self.input_queue.get() @@ -102,11 +109,7 @@ class Handler(object): break if uses_db: - db = db_conns.get(input.conn) - if db is None: - db = self.bot.get_db_connection(input.conn) - db_conns[input.conn] = db - input.db = db + input.db = self.bot.db try: run(self.func, input) @@ -165,7 +168,6 @@ def main(bot, conn, out): prefix = '^(?:[{}]?|'.format(command_prefix) else: prefix = '^(?:[{}]|'.format(command_prefix) - command_re = prefix + inp.conn.nick command_re += r'[,;:]+\s+)(\w+)(?:$|\s+)(.*)' diff --git a/plugins/dbtest.py b/plugins/dbtest.py new file mode 100644 index 0000000..e4d47e0 --- /dev/null +++ b/plugins/dbtest.py @@ -0,0 +1,20 @@ +from util import hook + +from sqlalchemy import Table, Column, Integer, String, MetaData, ForeignKey +from sqlalchemy.orm import scoped_session, sessionmaker + +def create_tables(metadata): + users_table = Table('users', metadata, + Column('id', Integer, primary_key=True), + Column('name', String), + Column('fullname', String), + Column('password', String) + ) + users_table.create() + + +@hook.command +def dbtest(inp, db=None): + metadata = MetaData(db) + create_tables(metadata) + print metadata \ No newline at end of file