Переглянути джерело

db/puppet: add contact_info_set flag

Signed-off-by: Sumner Evans <sumner@beeper.com>
Sumner Evans 2 роки тому
батько
коміт
ae000f98d2

+ 21 - 26
mautrix_instagram/db/puppet.py

@@ -38,6 +38,7 @@ class Puppet:
     photo_mxc: ContentURI | None
     name_set: bool
     avatar_set: bool
+    contact_info_set: bool
 
     is_registered: bool
 
@@ -56,6 +57,7 @@ class Puppet:
             self.photo_mxc,
             self.name_set,
             self.avatar_set,
+            self.contact_info_set,
             self.is_registered,
             self.custom_mxid,
             self.access_token,
@@ -63,21 +65,26 @@ class Puppet:
             str(self.base_url) if self.base_url else None,
         )
 
+    columns: ClassVar[str] = (
+        "pk, name, username, photo_id, photo_mxc, name_set, avatar_set, contact_info_set, "
+        "is_registered, custom_mxid, access_token, next_batch, base_url"
+    )
+
     async def insert(self) -> None:
-        q = (
-            "INSERT INTO puppet (pk, name, username, photo_id, photo_mxc, name_set, avatar_set,"
-            "                    is_registered, custom_mxid, access_token, next_batch, base_url) "
-            "VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)"
-        )
+        q = f"""
+            INSERT INTO puppet ({self.columns})
+            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+        """
         await self.db.execute(q, *self._values)
 
     async def update(self) -> None:
-        q = (
-            "UPDATE puppet SET name=$2, username=$3, photo_id=$4, photo_mxc=$5, name_set=$6,"
-            "                  avatar_set=$7, is_registered=$8, custom_mxid=$9, access_token=$10,"
-            "                  next_batch=$11, base_url=$12 "
-            "WHERE pk=$1"
-        )
+        q = """
+            UPDATE puppet
+            SET name=$2, username=$3, photo_id=$4, photo_mxc=$5, name_set=$6, avatar_set=$7,
+                contact_info_set=$8, is_registered=$9, custom_mxid=$10, access_token=$11,
+                next_batch=$12, base_url=$13
+            WHERE pk=$1
+        """
         await self.db.execute(q, *self._values)
 
     @classmethod
@@ -89,11 +96,7 @@ class Puppet:
 
     @classmethod
     async def get_by_pk(cls, pk: int) -> Puppet | None:
-        q = (
-            "SELECT pk, name, username, photo_id, photo_mxc, name_set, avatar_set, is_registered,"
-            "       custom_mxid, access_token, next_batch, base_url "
-            "FROM puppet WHERE pk=$1"
-        )
+        q = f"SELECT {cls.columns} FROM puppet WHERE pk=$1"
         row = await cls.db.fetchrow(q, pk)
         if not row:
             return None
@@ -101,11 +104,7 @@ class Puppet:
 
     @classmethod
     async def get_by_custom_mxid(cls, mxid: UserID) -> Puppet | None:
-        q = (
-            "SELECT pk, name, username, photo_id, photo_mxc, name_set, avatar_set, is_registered,"
-            "       custom_mxid, access_token, next_batch, base_url "
-            "FROM puppet WHERE custom_mxid=$1"
-        )
+        q = f"SELECT {cls.columns} FROM puppet WHERE custom_mxid=$1"
         row = await cls.db.fetchrow(q, mxid)
         if not row:
             return None
@@ -113,10 +112,6 @@ class Puppet:
 
     @classmethod
     async def all_with_custom_mxid(cls) -> list[Puppet]:
-        q = (
-            "SELECT pk, name, username, photo_id, photo_mxc, name_set, avatar_set, is_registered,"
-            "       custom_mxid, access_token, next_batch, base_url "
-            "FROM puppet WHERE custom_mxid IS NOT NULL"
-        )
+        q = f"SELECT {cls.columns} FROM puppet WHERE custom_mxid IS NOT NULL"
         rows = await cls.db.fetch(q)
         return [cls._from_row(row) for row in rows]

+ 1 - 0
mautrix_instagram/db/upgrade/__init__.py

@@ -31,4 +31,5 @@ from . import (
     v11_per_user_thread_sync_status,
     v12_portal_thread_image_id,
     v13_fix_portal_thread_image_id,
+    v14_puppet_contact_info_set,
 )

+ 25 - 0
mautrix_instagram/db/upgrade/v14_puppet_contact_info_set.py

@@ -0,0 +1,25 @@
+# mautrix-instagram - A Matrix-Instagram puppeting bridge.
+# Copyright (C) 2023 Tulir Asokan, Sumner Evans
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# 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 mautrix.util.async_db import Connection
+
+from . import upgrade_table
+
+
+@upgrade_table.register(description="Add contact_info_set column to puppet table")
+async def upgrade_v14(conn: Connection) -> None:
+    await conn.execute(
+        "ALTER TABLE puppet ADD COLUMN contact_info_set BOOLEAN NOT NULL DEFAULT false"
+    )

+ 2 - 0
mautrix_instagram/puppet.py

@@ -54,6 +54,7 @@ class Puppet(DBPuppet, BasePuppet):
         photo_mxc: ContentURI | None = None,
         name_set: bool = False,
         avatar_set: bool = False,
+        contact_info_set: bool = False,
         is_registered: bool = False,
         custom_mxid: UserID | None = None,
         access_token: str | None = None,
@@ -68,6 +69,7 @@ class Puppet(DBPuppet, BasePuppet):
             name_set=name_set,
             photo_mxc=photo_mxc,
             avatar_set=avatar_set,
+            contact_info_set=contact_info_set,
             is_registered=is_registered,
             custom_mxid=custom_mxid,
             access_token=access_token,