diff options
Diffstat (limited to 'quarg/main.py')
-rw-r--r-- | quarg/main.py | 193 |
1 files changed, 193 insertions, 0 deletions
diff --git a/quarg/main.py b/quarg/main.py new file mode 100644 index 0000000..eeaf943 --- /dev/null +++ b/quarg/main.py @@ -0,0 +1,193 @@ +import argparse +import configparser +import os +import sys +from abc import ABCMeta, abstractmethod +from timeit import default_timer as timer + +import dateutil.parser +import dateutil.relativedelta +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +import quarg.database.filters as filters +from quarg.database.tables import Backlog, Buffer, Network, QuasselUser, Sender +from quarg.quassel.formatter import format_from +from quarg.quassel.types import BufferType, MessageFlag, MessageType + +def errx(msg): + sys.exit(f'quarg: {msg}') + +# TODO Find out what to do about passing INVALID or GROUP as BufferType +# TODO Find out what to do about passing various message flags + +class ParseEnum(argparse.Action, metaclass=ABCMeta): + def __call__(self, parser, namespace, value, option_string=None): + key = value.upper() + if key not in self.enumclass.__members__: + errx(f'Not a valid {self.enumclass.describe()}: {value}') + + saved = getattr(namespace, self.dest) or [] + saved.append(self.enumclass[key]) + setattr(namespace, self.dest, saved) + + @property + @abstractmethod + def enumclass(self): + pass + +class ParseMessageType(ParseEnum): + @property + def enumclass(self): + return MessageType + +class ParseMessageFlag(ParseEnum): + @property + def enumclass(self): + return MessageFlag + +class ParseBufferType(ParseEnum): + @property + def enumclass(self): + return BufferType + +def parse_isodate(date): + try: + parsed = dateutil.parser.isoparse(date) + except ValueError as err: + errx(f'isoparse: invalid date format \'{date}\', {err}') + except OverflowError as err: + errx(f'isoparse: date overflows: \'{date}\'') + + return parsed + +# FIXME make sure pylint disables are actually still needed everywhere + +class ParseDate(argparse.Action): + # pylint: disable=too-few-public-methods, unsupported-membership-test + def __call__(self, parser, namespace, datespec, option_string=None): + setattr(namespace, self.dest, parse_isodate(datespec)) + +class ParseAround(argparse.Action): + # pylint: disable=too-few-public-methods, unsupported-membership-test + def __call__(self, parser, namespace, aroundspec, option_string=None): + if '/' in aroundspec: + # FIXME / fine here? + datespec, rangespec = aroundspec.split('/', 1) + try: + hour_range = int(rangespec) + except ValueError as err: + errx(err) + else: + datespec, hour_range = (aroundspec, 12) + + date = parse_isodate(datespec) + + offset = dateutil.relativedelta.relativedelta(hours=hour_range) + setattr(namespace, self.dest, (date - offset, date + offset)) + +# TODO Make --after/--before and --around mutually exclusive +# FIXME why need default=None for --joined/--no-joined? + +# pylint: disable=line-too-long +cli = argparse.ArgumentParser() +cli.add_argument('query', nargs='*', help='match messages containing this query') +cli.add_argument('-d', action='store_true', dest='debug', help='print SQL query information') +cli.add_argument('-e', action='store_true', dest='expr', help='interpret query as LIKE expression') +cli.add_argument('-b', action='append', dest='buffer', help='match messages sent to this buffer') +cli.add_argument('-B', action=ParseBufferType, dest='buftype', help='match messages sent to buffers of this type') +cli.add_argument('-n', action='append', dest='nick', help='match messages sent by this nickname') +cli.add_argument('-N', action='append', dest='network', help='match messages sent to this network') +cli.add_argument('-u', action='append', dest='user', help='match messages received by this quassel user') +cli.add_argument('-t', action=ParseMessageType, dest='msgtype', help='match messages of this message type') +cli.add_argument('-f', action=ParseMessageFlag, dest='msgflag', help='match messages with this flag') +cli.add_argument('-p', action='append', dest='prefix', help='match nicks with this prefix') +cli.add_argument('--joined', default=None, action='store_true', dest='joined', help='match messages sent to channels which are currently joined') +cli.add_argument('--no-joined', default=None, action='store_false', dest='joined', help='match messages sent to channels which are not currently joined') +cli.add_argument('--after', action=ParseDate, metavar='DATE', help='match messages sent after this date') +cli.add_argument('--before', action=ParseDate, metavar='DATE', help='match messages sent before this date') +cli.add_argument('--around', action=ParseAround, metavar='DATE', help='match messages sent within 12 hours of this date') +# pylint: enable=line-too-long + +Session = sessionmaker() + +def get_config(): + xdg_config_home = os.path.expanduser('~/.config/') + if 'XDG_CONFIG_HOME' in os.environ: + xdg_config_home = os.environ['XDG_CONFIG_HOME'] + + path = os.path.join(xdg_config_home, 'quarg', 'config') + + config = configparser.ConfigParser() + config.read(path) + + return config + +def collect_predicates(args): + funs = { + 'query': filters.msg_like if args.expr else filters.msg_contains, + 'buffer': filters.buffer, + 'nick': filters.nick, + 'after': filters.time_from, + 'before': filters.time_to, + 'around': filters.time_around, + 'user': filters.user, + 'network': filters.network, + 'msgflag': filters.msgflag, + 'msgtype': filters.msgtype, + 'buftype': filters.buftype, + 'prefix': filters.prefix, + 'joined': filters.joined, + } + + for key, value in vars(args).items(): + # FIXME sadly the 'joined' namespace will contain a falsy value, so + # check against None or [] + if key in funs and value not in [None, []]: + fun = funs[key] + if args.debug: + print(f'{key}: {value}', file=sys.stderr) + if isinstance(value, list): + yield filters.any_filter(fun, value) + else: + yield fun(value) + +def run_query(session, predicates): + start = timer() + + query = session.query(Backlog).join(Sender).join(Buffer).join(Network).join(QuasselUser) + + for predicate in predicates: + query = query.filter(predicate) + + rows = query.order_by(Backlog.time).all() + + end = timer() + + return (rows, end - start) + +def main(): + config = get_config() + + if not config.has_option('Database', 'url'): + errx('No database URL set in config file.') + + args = cli.parse_args() + + engine = create_engine(config.get('Database', 'url'), echo=args.debug) + session = Session(bind=engine) + + predicates = list(collect_predicates(args)) + + if not predicates: + errx('Nothing to match.') + + rows, time = run_query(session, predicates) + + for row in rows: + print(format_from(row)) + + print(f'Query returned {len(rows)} lines in {time:.4f} seconds.', file=sys.stderr) + +if __name__ == "__main__": + main() |