Fix .tell message collisions

This adds unique IDs to each tell message, to allow for more "proper" ORM access to each message, for marking them read and such

Author: linuxdaemon

plugins/tell.py

@@ -8,30 +8,46 @@
     Boolean,
     Column,
     DateTime,
+    Integer,
     PrimaryKeyConstraint,
     String,
     Table,
     and_,
     not_,
+    update,
 )
 from sqlalchemy.sql import select
 
 from cloudbot import hook
 from cloudbot.event import EventType
+from cloudbot.hook import Priority
 from cloudbot.util import database, timeformat, web
 from cloudbot.util.formatting import gen_markdown_table
 
-table = Table(
-    "tells",
-    database.metadata,
-    Column("connection", String),
-    Column("sender", String),
-    Column("target", String),
-    Column("message", String),
-    Column("is_read", Boolean),
-    Column("time_sent", DateTime),
-    Column("time_read", DateTime),
-)
+
+class TellMessage(database.Base):
+    __tablename__ = "tell_messages"
+
+    msg_id = Column(Integer, primary_key=True, autoincrement=True)
+    conn = Column(String, index=True)
+    sender = Column(String)
+    target = Column(String, index=True)
+    message = Column(String)
+    is_read = Column(Boolean, default=False, index=True)
+    time_sent = Column(DateTime)
+    time_read = Column(DateTime)
+
+    def format_for_message(self):
+        reltime = timeformat.time_since(self.time_sent)
+        return f"{self.sender} sent you a message {reltime} ago: {self.message}"
+
+    def mark_read(self, now=None):
+        if now is None:
+            now = datetime.now()
+
+        self.is_read = True
+        self.time_read = now
+
 
 disable_table = Table(
     "tell_ignores",
@@ -60,12 +76,42 @@
 tell_cache: List[Tuple[str, str]] = []
 
 
+@hook.on_start(priority=Priority.HIGHEST)
+def migrate_tables(db):
+    table = Table(
+        "tells",
+        database.metadata,
+        Column("connection", String),
+        Column("sender", String),
+        Column("target", String),
+        Column("message", String),
+        Column("is_read", Boolean),
+        Column("time_sent", DateTime),
+        Column("time_read", DateTime),
+    )
+
+    if not table.exists(db.bind):
+        return
+
+    if TellMessage.__table__.exists(db.bin):
+        raise Exception(
+            f"Can't migrate table {table.name} to {TellMessage.__table__.name}, destination already exists"
+        )
+
+    data = db.execute(table.select())
+    db.bulk_insert_mappings(TellMessage, data, return_defaults=True)
+    db.commit()
+    table.drop(db.bind)
+
+
 @hook.on_start()
 def load_cache(db):
     new_cache = []
-    for row in db.execute(table.select().where(not_(table.c.is_read))):
-        conn = row["connection"]
-        target = row["target"]
+    for conn, target in db.execute(
+        select(
+            [TellMessage.conn, TellMessage.target], not_(TellMessage.is_read)
+        )
+    ):
         new_cache.append((conn, target))
 
     tell_cache.clear()
@@ -183,48 +229,35 @@ def list_ignores(conn, nick):
         yield mask
 
 
-def get_unread(db, server, target):
+def get_unread(db, server, target) -> List[TellMessage]:
     query = (
-        select([table.c.sender, table.c.message, table.c.time_sent])
-        .where(table.c.connection == server.lower())
-        .where(table.c.target == target.lower())
-        .where(not_(table.c.is_read))
-        .order_by(table.c.time_sent)
+        select(TellMessage)
+        .where(not_(TellMessage.is_read))
+        .where(TellMessage.conn == server)
+        .where(TellMessage.target == target.lower())
+        .order_by(TellMessage.time_sent)
     )
-    return db.execute(query).fetchall()
+
+    return db.execute(query).scalars().all()
 
 
 def count_unread(db, server, target):
     query = (
-        select([sa.func.count()])
-        .select_from(table)
-        .where(table.c.connection == server.lower())
-        .where(table.c.target == target.lower())
-        .where(not_(table.c.is_read))
+        select(sa.func.count(TellMessage.msg_id))
+        .where(TellMessage.conn == server.lower())
+        .where(TellMessage.target == target.lower())
+        .where(not_(TellMessage.is_read))
     )
 
     return db.execute(query).fetchone()[0]
 
 
 def read_all_tells(db, server, target):
     query = (
-        table.update()
-        .where(table.c.connection == server.lower())
-        .where(table.c.target == target.lower())
-        .where(not_(table.c.is_read))
-        .values(is_read=True)
-    )
-    db.execute(query)
-    db.commit()
-    load_cache(db)
-
-
-def read_tell(db, server, target, message):
-    query = (
-        table.update()
-        .where(table.c.connection == server.lower())
-        .where(table.c.target == target.lower())
-        .where(table.c.message == message)
+        update(TellMessage)
+        .where(TellMessage.conn == server.lower())
+        .where(TellMessage.target == target.lower())
+        .where(TellMessage.is_read.is_(False))
         .values(is_read=True)
     )
     db.execute(query)
@@ -233,15 +266,14 @@ def read_tell(db, server, target, message):
 
 
 def add_tell(db, server, sender, target, message):
-    query = table.insert().values(
-        connection=server.lower(),
+    new_tell = TellMessage(
+        conn=server.lower(),
         sender=sender.lower(),
         target=target.lower(),
         message=message,
-        is_read=False,
-        time_sent=datetime.today(),
+        time_sent=datetime.now(),
     )
-    db.execute(query)
+    db.add(new_tell)
     db.commit()
     load_cache(db)
 
@@ -267,20 +299,18 @@ def tellinput(conn, db, nick, notice, content):
     if not tells:
         return
 
-    user_from, message, time_sent = tells[0]
-    reltime = timeformat.time_since(time_sent)
-    reply = "{} sent you a message {} ago: {}".format(
-        user_from, reltime, message
-    )
+    first_tell = tells[0]
+    reply = first_tell.format_for_message()
 
     if len(tells) > 1:
         reply += " (+{} more, {}showtells to view)".format(
             len(tells) - 1, conn.config["command_prefix"][0]
         )
 
-    read_tell(db, conn.name, nick, message)
     notice(reply)
 
+    first_tell.mark_read()
+
 
 @hook.command(autohelp=False)
 def showtells(nick, notice, db, conn):
@@ -293,9 +323,7 @@ def showtells(nick, notice, db, conn):
         return
 
     for tell in tells:
-        sender, message, time_sent = tell
-        past = timeformat.time_since(time_sent)
-        notice("{} sent you a message {} ago: {}".format(sender, past, message))
+        notice(tell.format_for_message())
 
     read_all_tells(db, conn.name, nick)
 

tests/plugin_tests/test_tell.py

@@ -9,7 +9,7 @@
 
 def init_tables(mock_db):
     db_engine = mock_db.engine
-    tell.table.create(db_engine)
+    tell.TellMessage.__table__.create(db_engine)
     tell.disable_table.create(db_engine)
     tell.ignore_table.create(db_engine)
     session = mock_db.session()
@@ -109,8 +109,9 @@ def test_showtells(mock_db, freeze_time):
     event = MagicMock()
     tell.add_tell(mock_db.session(), server, sender, target, message)
 
-    assert mock_db.get_data(tell.table) == [
+    assert mock_db.get_data(tell.TellMessage.__table__) == [
         (
+            1,
             "testconn",
             "foo",
             "other",
@@ -127,8 +128,9 @@ def test_showtells(mock_db, freeze_time):
     assert event.mock_calls == [
         call.notice("foo sent you a message 60 seconds ago: bar")
     ]
-    assert mock_db.get_data(tell.table) == [
+    assert mock_db.get_data(tell.TellMessage.__table__) == [
         (
+            1,
             "testconn",
             "foo",
             "other",
@@ -206,6 +208,75 @@ def test_tellinput(mock_db, freeze_time):
     ]
 
 
+def test_read_tell_spam(mock_db, freeze_time):
+    init_tables(mock_db)
+    db = mock_db.session()
+    conn = MockConn()
+    conn.config["command_prefix"] = "."
+    sender = "foo"
+    nick = "other"
+    message = "bar"
+    message2 = "baraa"
+    event = MagicMock()
+    content = "aa"
+    tell.add_tell(db, conn.name.lower(), sender, nick, message)
+    freeze_time.tick()
+    tell.add_tell(db, conn.name.lower(), sender, nick, message)
+    freeze_time.tick()
+    tell.add_tell(db, conn.name.lower(), sender, nick, message)
+    freeze_time.tick()
+    tell.add_tell(db, conn.name.lower(), sender, nick, message2)
+    assert mock_db.get_data(tell.TellMessage.__table__) == [
+        (
+            1,
+            "testconn",
+            sender,
+            nick,
+            message,
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 36),
+            None,
+        ),
+        (
+            2,
+            "testconn",
+            sender,
+            nick,
+            message,
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 37),
+            None,
+        ),
+        (
+            3,
+            "testconn",
+            sender,
+            nick,
+            message,
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 38),
+            None,
+        ),
+        (
+            4,
+            "testconn",
+            sender,
+            nick,
+            message2,
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 39),
+            None,
+        ),
+    ]
+    res = tell.tellinput(conn, db, nick, event.notice, content)
+    assert res is None
+    assert event.mock_calls == [
+        call.notice(
+            "foo sent you a message 3 seconds ago: bar (+3 more, .showtells to view)"
+        )
+    ]
+
+
 def test_tellinput_multiple(mock_db, freeze_time):
     init_tables(mock_db)
     db = mock_db.session()