Beispiel #1
0
class UIDStack(object):
    """Thin convenience wrapper around gevent.queue.LifoQueue.
    Each entry in the stack is a pair (uid, metadata), where the metadata may
    be None."""
    def __init__(self):
        self._lifoqueue = LifoQueue()

    def empty(self):
        return self._lifoqueue.empty()

    def get(self):
        return self._lifoqueue.get_nowait()

    def peek(self):
        # This should be LifoQueue.peek_nowait(), which is currently buggy in
        # gevent. Can update with gevent version 1.0.2.
        return self._lifoqueue.queue[-1]

    def put(self, uid, metadata):
        self._lifoqueue.put((uid, metadata))

    def discard(self, objects):
        self._lifoqueue.queue = [
            item for item in self._lifoqueue.queue if item not in objects
        ]

    def qsize(self):
        return self._lifoqueue.qsize()

    def __iter__(self):
        for item in self._lifoqueue.queue:
            yield item
Beispiel #2
0
class UIDStack(object):
    """Thin convenience wrapper around gevent.queue.LifoQueue.
    Each entry in the stack is a pair (uid, metadata), where the metadata may
    be None."""
    def __init__(self):
        self._lifoqueue = LifoQueue()

    def empty(self):
        return self._lifoqueue.empty()

    def get(self):
        return self._lifoqueue.get_nowait()

    def peek(self):
        # This should be LifoQueue.peek_nowait(), which is currently buggy in
        # gevent. Can update with gevent version 1.0.2.
        return self._lifoqueue.queue[-1]

    def put(self, uid, metadata):
        self._lifoqueue.put((uid, metadata))

    def discard(self, objects):
        self._lifoqueue.queue = [item for item in self._lifoqueue.queue if item
                                 not in objects]

    def qsize(self):
        return self._lifoqueue.qsize()

    def __iter__(self):
        for item in self._lifoqueue.queue:
            yield item
Beispiel #3
0
class Stack(object):
    """Thin convenience wrapper around gevent.queue.LifoQueue."""
    def __init__(self, key, initial_elements=None):
        self.key = key
        self._lifoqueue = LifoQueue()
        if initial_elements is not None:
            self._lifoqueue.queue = sorted(list(initial_elements),
                                           key=self.key)

    def empty(self):
        return self._lifoqueue.empty()

    def get(self):
        return self._lifoqueue.get_nowait()

    def peek(self):
        # This should be LifoQueue.peek_nowait(), which is currently buggy in
        # gevent. Can update with gevent version 1.0.2.
        return self._lifoqueue.queue[-1]

    def put(self, obj):
        self._lifoqueue.put(obj)

    def update_from(self, objects):
        for obj in sorted(list(objects), key=self.key):
            self._lifoqueue.put(obj)

    def discard(self, objects):
        self._lifoqueue.queue = [item for item in self._lifoqueue.queue if item
                                 not in objects]

    def qsize(self):
        return self._lifoqueue.qsize()

    def __iter__(self):
        for item in self._lifoqueue.queue:
            yield item
Beispiel #4
0
class SQLPlugin(Plugin):
    global_plugin = True

    def load(self, ctx):
        self.models = ctx.get('models', {})
        self.backfills = {}
        self.user_updates = LifoQueue(maxsize=4096)
        super(SQLPlugin, self).load(ctx)

    def unload(self, ctx):
        ctx['models'] = self.models
        super(SQLPlugin, self).unload(ctx)

    @Plugin.schedule(15, init=False)
    def update_users(self):
        already_updated = set()

        while True:
            # Only update so many at a time
            if len(already_updated) > 10000:
                return

            try:
                user_id, data = self.user_updates.get_nowait()
            except Empty:
                return

            if user_id in already_updated:
                continue

            already_updated.add(user_id)

            try:
                User.update(**data).where(User.user_id == user_id).execute()
            except:
                self.log.exception('Failed to update user %s: ', user_id)

    @Plugin.listen('VoiceStateUpdate', priority=Priority.BEFORE)
    def on_voice_state_update(self, event):
        pre_state = self.state.voice_states.get(event.session_id)
        GuildVoiceSession.create_or_update(pre_state, event.state)

    @Plugin.listen('PresenceUpdate')
    def on_presence_update(self, event):
        updates = {}

        if event.user.avatar != UNSET:
            updates['avatar'] = event.user.avatar

        if event.user.username != UNSET:
            updates['username'] = event.user.username

        if event.user.discriminator != UNSET:
            updates['discriminator'] = int(event.user.discriminator)

        if not updates:
            return

        self.user_updates.put((event.user.id, updates))

    @Plugin.listen('MessageCreate')
    def on_message_create(self, event):
        Message.from_disco_message(event.message)

    @Plugin.listen('MessageUpdate')
    def on_message_update(self, event):
        Message.from_disco_message_update(event.message)

    @Plugin.listen('MessageDelete')
    def on_message_delete(self, event):
        Message.update(deleted=True).where(Message.id == event.id).execute()

    @Plugin.listen('MessageDeleteBulk')
    def on_message_delete_bulk(self, event):
        Message.update(deleted=True).where((Message.id << event.ids)).execute()

    @Plugin.listen('MessageReactionAdd', priority=Priority.BEFORE)
    def on_message_reaction_add(self, event):
        Reaction.from_disco_reaction(event)

    @Plugin.listen('MessageReactionRemove', priority=Priority.BEFORE)
    def on_message_reaction_remove(self, event):
        Reaction.delete().where(
            (Reaction.message_id == event.message_id) &
            (Reaction.user_id == event.user_id) &
            (Reaction.emoji_id == (event.emoji.id or None)) &
            (Reaction.emoji_name == (event.emoji.name or None))).execute()

    @Plugin.listen('MessageReactionRemoveAll')
    def on_message_reaction_remove_all(self, event):
        Reaction.delete().where((Reaction.message_id == event.message_id)).execute()

    @Plugin.listen('GuildEmojisUpdate', priority=Priority.BEFORE)
    def on_guild_emojis_update(self, event):
        ids = []

        for emoji in event.emojis:
            GuildEmoji.from_disco_guild_emoji(emoji, event.guild_id)
            ids.append(emoji.id)

        GuildEmoji.update(deleted=True).where(
            (GuildEmoji.guild_id == event.guild_id) &
            (~(GuildEmoji.emoji_id << ids))
        ).execute()

    @Plugin.listen('GuildCreate')
    def on_guild_create(self, event):
        for channel in list(event.channels.values()):
            Channel.from_disco_channel(channel)

        for emoji in list(event.emojis.values()):
            GuildEmoji.from_disco_guild_emoji(emoji, guild_id=event.guild.id)

    @Plugin.listen('GuildDelete')
    def on_guild_delete(self, event):
        if event.deleted:
            Channel.update(deleted=True).where(
                Channel.guild_id == event.id
            ).execute()

    @Plugin.listen('ChannelCreate')
    def on_channel_create(self, event):
        Channel.from_disco_channel(event.channel)

    @Plugin.listen('ChannelUpdate')
    def on_channel_update(self, event):
        Channel.from_disco_channel(event.channel)

    @Plugin.listen('ChannelDelete')
    def on_channel_delete(self, event):
        Channel.update(deleted=True).where(Channel.channel_id == event.channel.id).execute()

    @Plugin.command('sql', level=-1, global_=True)
    def command_sql(self, event):
        conn = database.obj.get_conn()

        try:
            tbl = MessageTable(codeblock=False)

            with conn.cursor() as cur:
                start = time.time()
                cur.execute(event.codeblock.format(e=event))
                dur = time.time() - start

                if not cur.description:
                    return event.msg.reply('_Query took {}ms - no result._'.format(int(dur * 1000)))

                tbl.set_header(*[desc[0] for desc in cur.description])

                for row in cur.fetchall():
                    tbl.add(*row)

                result = tbl.compile()
                if len(result) > 1900:
                    return event.msg.reply(
                        '_Query took {}ms_'.format(int(dur * 1000)),
                        attachments=[('sql_result_{}.txt'.format(event.msg.id), result)]
                    )

                event.msg.reply(u'```{}```_Query took {}ms_'.format(result, int(dur * 1000)))
        except psycopg2.Error as e:
            event.msg.reply(u'```{}```'.format(e.pgerror))

    @Plugin.command('init', '<entity:user|channel>', level=-1, group='markov', global_=True)
    def command_markov(self, event, entity):
        if isinstance(entity, DiscoUser):
            q = Message.select().where(Message.author_id == entity.id).limit(500000)
        else:
            q = Message.select().where(Message.channel_id == entity.id).limit(500000)

        text = [msg.content for msg in q]
        self.models[entity.id] = markovify.NewlineText('\n'.join(text))
        event.msg.reply(u':ok_hand: created markov model for {} using {} messages'.format(entity, len(text)))

    @Plugin.command('one', '<entity:user|channel>', level=-1, group='markov', global_=True)
    def command_markov_one(self, event, entity):
        if entity.id not in self.models:
            return event.msg.reply(':warning: no model created yet for {}'.format(entity))

        sentence = self.models[entity.id].make_sentence(max_overlap_ratio=1, max_overlap_total=500)
        if not sentence:
            event.msg.reply(':warning: not enough data :(')
            return
        event.msg.reply(u'{}: {}'.format(entity, sentence))

    @Plugin.command('many', '<entity:user|channel> [count:int]', level=-1, group='markov', global_=True)
    def command_markov_many(self, event, entity, count=5):
        if entity.id not in self.models:
            return event.msg.reply(':warning: no model created yet for {}'.format(entity))

        for _ in range(int(count)):
            sentence = self.models[entity.id].make_sentence(max_overlap_total=500)
            if not sentence:
                event.msg.reply(':warning: not enough data :(')
                return
            event.msg.reply(u'{}: {}'.format(entity, sentence))

    @Plugin.command('list', level=-1, group='markov', global_=True)
    def command_markov_list(self, event):
        event.msg.reply(u'`{}`'.format(', '.join(map(str, self.models.keys()))))

    @Plugin.command('delete', '<oid:snowflake>', level=-1, group='markov', global_=True)
    def command_markov_delete(self, event, oid):
        if oid not in self.models:
            return event.msg.reply(':warning: no model with that ID')

        del self.models[oid]
        event.msg.reply(':ok_hand: deleted model')

    @Plugin.command('clear', level=-1, group='markov', global_=True)
    def command_markov_clear(self, event):
        self.models = {}
        event.msg.reply(':ok_hand: cleared models')

    @Plugin.command('message', '<channel:snowflake> <message:snowflake>', level=-1, group='backfill', global_=True)
    def command_backfill_message(self, event, channel, message):
        channel = self.state.channels.get(channel)
        Message.from_disco_message(channel.get_message(message))
        return event.msg.reply(':ok_hand: backfilled')

    @Plugin.command('reactions', '<message:snowflake>', level=-1, group='backfill', global_=True)
    def command_sql_reactions(self, event, message):
        try:
            message = Message.get(id=message)
        except Message.DoesNotExist:
            return event.msg.reply(':warning: no message found')

        message = self.state.channels.get(message.channel_id).get_message(message.id)
        for reaction in message.reactions:
            for users in message.get_reactors(reaction.emoji, bulk=True):
                Reaction.from_disco_reactors(message.id, reaction, (i.id for i in users))

    @Plugin.command('global', '<duration:str> [pool:int]', level=-1, global_=True, context={'mode': 'global'}, group='recover')
    @Plugin.command('here', '<duration:str> [pool:int]', level=-1, global_=True, context={'mode': 'here'}, group='recover')
    def command_recover(self, event, duration, pool=4, mode=None):
        if mode == 'global':
            channels = list(self.state.channels.values())
        else:
            channels = list(event.guild.channels.values())

        start_at = parse_duration(duration, negative=True)

        pool = Pool(pool)

        total = len(channels)
        msg = event.msg.reply('Recovery Status: 0/{}'.format(total))
        recoveries = []

        def updater():
            last = len(recoveries)

            while True:
                if last != len(recoveries):
                    last = len(recoveries)
                    msg.edit('Recovery Status: {}/{}'.format(len(recoveries), total))
                gevent.sleep(5)

        u = self.spawn(updater)

        try:
            for channel in channels:
                pool.wait_available()
                r = Recovery(self.log, channel, start_at)
                pool.spawn(r.run)
                recoveries.append(r)
        finally:
            pool.join()
            u.kill()

        msg.edit('RECOVERY COMPLETED ({} total messages)'.format(
            sum([i._recovered for i in recoveries])
        ))

    @Plugin.command('backfill channel', '[channel:snowflake]', level=-1, global_=True)
    def command_backfill_channel(self, event, channel=None):
        channel = self.state.channels.get(channel) if channel else event.channel
        backfill_channel.queue(channel.id)
        event.msg.reply(':ok_hand: enqueued channel to be backfilled')

    @Plugin.command('backfill guild', '[guild:guild] [concurrency:int]', level=-1, global_=True)
    def command_backfill_guild(self, event, guild=None, concurrency=1):
        guild = guild or event.guild
        backfill_guild.queue(guild.id)
        event.msg.reply(':ok_hand: enqueued guild to be backfilled')

    @Plugin.command('usage', '<word:str> [unit:str] [amount:int]', level=-1, group='words')
    def words_usage(self, event, word, unit='days', amount=7):
        sql = '''
            SELECT date, coalesce(count, 0) AS count
            FROM
                generate_series(
                    NOW() - interval %s,
                    NOW(),
                    %s
                ) AS date
            LEFT OUTER JOIN (
                SELECT date_trunc(%s, timestamp) AS dt, count(*) AS count
                FROM messages
                WHERE
                    timestamp >= (NOW() - interval %s) AND
                    timestamp < (NOW()) AND
                    guild_id=%s AND
                    (SELECT count(*) FROM regexp_matches(content, %s)) >= 1
                GROUP BY dt
            ) results
            ON (date_trunc(%s, date) = results.dt);
        '''

        msg = event.msg.reply(':alarm_clock: One moment pls...')

        start = time.time()
        tuples = list(Message.raw(
            sql,
            '{} {}'.format(amount, unit),
            '1 {}'.format(unit),
            unit,
            '{} {}'.format(amount, unit),
            event.guild.id,
            '\s?{}\s?'.format(word),
            unit
        ).tuples())
        sql_duration = time.time() - start

        start = time.time()
        chart = pygal.Line()
        chart.title = 'Usage of {} Over {} {}'.format(
            word, amount, unit,
        )

        if unit == 'days':
            chart.x_labels = [i[0].strftime('%a %d') for i in tuples]
        elif unit == 'minutes':
            chart.x_labels = [i[0].strftime('%X') for i in tuples]
        else:
            chart.x_labels = [i[0].strftime('%x %X') for i in tuples]

        chart.x_labels = [i[0] for i in tuples]
        chart.add(word, [i[1] for i in tuples])

        pngdata = cairosvg.svg2png(
            bytestring=chart.render(),
            dpi=72)
        chart_duration = time.time() - start

        event.msg.reply(
            '_SQL: {}ms_ - _Chart: {}ms_'.format(
                int(sql_duration * 1000),
                int(chart_duration * 1000),
            ),
            attachments=[('chart.png', pngdata)])
        msg.delete()

    @Plugin.command('top', '<target:user|channel|guild>', level=-1, group='words')
    def words_top(self, event, target):
        if isinstance(target, DiscoUser):
            q = 'author_id'
        elif isinstance(target, DiscoChannel):
            q = 'channel_id'
        elif isinstance(target, DiscoGuild):
            q = 'guild_id'
        else:
            raise Exception("You should not be here")

        sql = """
            SELECT word, count(*)
            FROM (
                SELECT regexp_split_to_table(content, '\s') as word
                FROM messages
                WHERE {}=%s
                LIMIT 3000000
            ) t
            GROUP BY word
            ORDER BY 2 DESC
            LIMIT 30
        """.format(q)

        t = MessageTable()
        t.set_header('Word', 'Count')

        for word, count in Message.raw(sql, target.id).tuples():
            if '```' in word:
                continue
            t.add(word, count)

        event.msg.reply(t.compile())
class AbstractDatabaseConnectionPool(object):
    def __init__(self, maxsize=100, maxwait=1.0, expires=None, cleanup=None):
        """
        The pool manages opened connections to the database. The main strategy is to keep the smallest number
        of alive connections which are required for best web service performance.
        In most cases connections are taken from the pool. In case of views-peeks, pool creates some
        extra resources preventing service gone unavailable. In time of low traffic (night) unnecessary
        connections are released.

        Parameters
        ----------
        maxsize : int
                  Soft limit of the number of created connections. After reaching this limit
                  taking the next connection first waits `maxwait` time for any returned slot.
        maxwait : float
                  The time in seconds which is to be wait before creating new connection after the pool gets empty.
                  It may be 0 then immediate connections are created til `maxoverflow` is reached.
        expires : float
                  The time in seconds indicates how long connection should stay alive.
                  It is also used to close unneeded slots.
        """
        if not isinstance(maxsize, integer_types):
            raise TypeError('Expected integer, got %r' % (maxsize, ))

        self._maxsize = maxsize
        self._maxwait = maxwait
        self._expires = expires
        self._cleanup = cleanup
        self._created_at = {}
        self._latest_use = {}
        self._pool = LifoQueue()
        self._size = 0
        self._latest_cleanup = 0 if self._expires or self._cleanup else 0xffffffffffffffff
        self._interval_cleanup = min(
            self._expires or self._cleanup, self._cleanup
            or self._expires) if self._expires or self._cleanup else 0
        self._cleanup_lock = Semaphore(value=1)

    def create_connection(self):
        raise NotImplementedError()

    def close_connection(self, item):
        try:
            self._size -= 1
            self._created_at.pop(id(item), None)
            self._latest_use.pop(id(item), None)
            item.close()
        except Exception:
            pass

    def cleanup(self):
        self._cleanup_queue(time.time())

    def _cleanup_queue(self, now):

        if self._latest_cleanup > now:
            return

        with self._cleanup_lock:

            if self._latest_cleanup > now:
                return

            self._latest_cleanup = now + self._interval_cleanup

            cleanup = now - self._cleanup if self._cleanup else None
            expires = now - self._expires if self._expires else None

            # Instead of creating new LIFO for self._pool, the ole one is reused,
            # beacuse some othere might wait for connetion on it.
            fresh_slots = []

            try:
                # try to fill self._pool ASAP, preventing creation of new connections.
                # because after this loop LIFO will be in reversed order
                while not self._pool.empty():
                    item = self._pool.get_nowait()
                    if cleanup and self._latest_use.get(id(item), 0) < cleanup:
                        self.close_connection(item)
                    elif expires and self._created_at.get(id(item),
                                                          0) < expires:
                        self.close_connection(item)
                    else:
                        fresh_slots.append(item)
            except Empty:
                pass

            # Reverse order back (frestest connections should be at the begining)
            for conn in reversed(fresh_slots):
                self._pool.put_nowait(conn)

    def get(self):

        try:
            return self._pool.get_nowait()
        except Empty:
            pass

        if self._size >= self._maxsize:
            try:
                return self._pool.get(timeout=self._maxwait)
            except Empty:
                pass

        # It is posiible that after waiting self._maxwait time, non connection has been returned
        # because of cleaning up old ones on put(), so there is not connection but also LIFO is not full.
        # In that case new connection shouls be created, otherwise exception is risen.
        if self._size >= self._maxsize:
            raise OperationalError(
                "Too many connections created: {} (maxsize is {})".format(
                    self._size, self._maxsize))

        try:
            self._size += 1
            conn = self.create_connection()
        except:
            self._size -= 1
            raise

        now = time.time()
        self._created_at[id(conn)] = now
        self._latest_use[id(conn)] = now
        return conn

    def put(self, conn):
        now = time.time()
        self._pool.put(conn)
        self._latest_use[id(conn)] = now

        self._cleanup_queue(now)

    def closeall(self):
        while not self._pool.empty():
            conn = self._pool.get_nowait()
            try:
                conn.close()
            except Exception:
                pass
        self._size = 0

    @contextlib.contextmanager
    def connection(self, isolation_level=None):
        conn = self.get()
        try:
            if isolation_level is not None:
                if conn.isolation_level == isolation_level:
                    isolation_level = None
                else:
                    conn.set_isolation_level(isolation_level)
            yield conn
        except:
            if conn.closed:
                conn = None
                self.closeall()
            else:
                conn = self._rollback(conn)
            raise
        else:
            if conn.closed:
                raise OperationalError(
                    "Cannot commit because connection was closed: %r" %
                    (conn, ))
            conn.commit()
        finally:
            if conn is not None and not conn.closed:
                if isolation_level is not None:
                    conn.set_isolation_level(isolation_level)
                self.put(conn)

    @contextlib.contextmanager
    def cursor(self, *args, **kwargs):
        isolation_level = kwargs.pop('isolation_level', None)
        with self.connection(isolation_level) as conn:
            yield conn.cursor(*args, **kwargs)

    def _rollback(self, conn):
        try:
            conn.rollback()
        except:
            gevent.get_hub().handle_error(conn, *sys.exc_info())
            return
        return conn

    def execute(self, *args, **kwargs):
        with self.cursor(**kwargs) as cursor:
            cursor.execute(*args)
            return cursor.rowcount

    def fetchone(self, *args, **kwargs):
        with self.cursor(**kwargs) as cursor:
            cursor.execute(*args)
            return cursor.fetchone()

    def fetchall(self, *args, **kwargs):
        with self.cursor(**kwargs) as cursor:
            cursor.execute(*args)
            return cursor.fetchall()

    def fetchiter(self, *args, **kwargs):
        with self.cursor(**kwargs) as cursor:
            cursor.execute(*args)
            while True:
                items = cursor.fetchmany()
                if not items:
                    break
                for item in items:
                    yield item