From 3438179d48de6b01e56c55dbaf7f137c6a31817b Mon Sep 17 00:00:00 2001
From: Andrew Ferrazzutti <fair@miscworks.net>
Date: Sun, 28 Feb 2021 01:50:14 -0500
Subject: [PATCH] Stuff for multi-user support

---
 matrix_appservice_line/db/upgrade.py |  9 ++++++
 matrix_appservice_line/db/user.py    | 45 ++++++++++++++++++++--------
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/matrix_appservice_line/db/upgrade.py b/matrix_appservice_line/db/upgrade.py
index 0596239..27b44eb 100644
--- a/matrix_appservice_line/db/upgrade.py
+++ b/matrix_appservice_line/db/upgrade.py
@@ -46,3 +46,12 @@ async def upgrade_v1(conn: Connection) -> None:
 
         UNIQUE (mxid, mx_room)
     )""")
+
+
+@upgrade_table.register(description="Multi-user support")
+async def upgrade_v2(conn: Connection) -> None:
+    mid_exists = await conn.fetchval(
+        'SELECT EXISTS(SELECT FROM information_schema.columns '
+        'WHERE table_name="user" AND column_name="mid"')
+    if not mid_exists:
+        await conn.execute("ALTER TABLE 'user' ADD COLUMN mid TEXT")
\ No newline at end of file
diff --git a/matrix_appservice_line/db/user.py b/matrix_appservice_line/db/user.py
index f3b2fe9..f6065f6 100644
--- a/matrix_appservice_line/db/user.py
+++ b/matrix_appservice_line/db/user.py
@@ -13,8 +13,9 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-from typing import Optional, ClassVar, TYPE_CHECKING
+from typing import Optional, List, TYPE_CHECKING, ClassVar
 
+from asyncpg import Record
 from attr import dataclass
 
 from mautrix.types import UserID, RoomID
@@ -28,22 +29,40 @@ class User:
     db: ClassVar[Database] = fake_db
 
     mxid: UserID
+    mid: Optional[str]
     notice_room: Optional[RoomID]
 
-    async def insert(self) -> None:
-        q = ('INSERT INTO "user" (mxid, notice_room) '
-             'VALUES ($1, $2)')
-        await self.db.execute(q, self.mxid, self.notice_room)
+    @classmethod
+    def _from_row(cls, row: Optional[Record]) -> Optional['User']:
+        if row is None:
+            return None
+        return cls(**row)
 
-    async def update(self) -> None:
-        await self.db.execute('UPDATE "user" SET notice_room=$2 WHERE mxid=$1',
-                              self.mxid, self.notice_room)
+    @classmethod
+    async def all_logged_in(cls) -> List['User']:
+        rows = await cls.db.fetch('SELECT mxid, mid, notice_room FROM "user" '
+                                  "WHERE mid IS NOT NULL")
+        return [cls._from_row(row) for row in rows]
+
+    @classmethod
+    async def get_by_mid(cls, mid: str) -> Optional['User']:
+        q = 'SELECT mxid, mid, notice_room FROM "user" WHERE mid=$1'
+        row = await cls.db.fetchrow(q, mid)
+        return cls._from_row(row)
 
     @classmethod
     async def get_by_mxid(cls, mxid: UserID) -> Optional['User']:
-        q = ("SELECT mxid, notice_room "
-             'FROM "user" WHERE mxid=$1')
+        q = 'SELECT mxid, mid, notice_room FROM "user" WHERE mxid=$1'
         row = await cls.db.fetchrow(q, mxid)
-        if not row:
-            return None
-        return cls(**row)
+        return cls._from_row(row)
+
+    async def insert(self) -> None:
+        q = 'INSERT INTO "user" (mxid, mid, notice_room) VALUES ($1, $2, $3)'
+        await self.db.execute(q, self.mxid, self.mid, self.notice_room)
+
+    async def delete(self) -> None:
+        await self.db.execute('DELETE FROM "user" WHERE mxid=$1', self.mxid)
+
+    async def update(self) -> None:
+        await self.db.execute('UPDATE "user" SET mid=$2, notice_room=$3 WHERE mxid=$1',
+                              self.mxid, self.mid, self.notice_room)