From da5b8556f28f53fb6939a6e94d33cfc9e571e359 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Wed, 6 Apr 2022 12:47:35 +0300 Subject: [PATCH] Add phone number field for puppets --- mautrix_telegram/db/puppet.py | 15 +++++++----- mautrix_telegram/db/upgrade/__init__.py | 1 + .../db/upgrade/v00_latest_revision.py | 3 ++- .../db/upgrade/v01_initial_revision.py | 10 ++++---- .../db/upgrade/v07_puppet_phone_number.py | 23 +++++++++++++++++++ mautrix_telegram/puppet.py | 8 +++++++ 6 files changed, 48 insertions(+), 12 deletions(-) create mode 100644 mautrix_telegram/db/upgrade/v07_puppet_phone_number.py diff --git a/mautrix_telegram/db/puppet.py b/mautrix_telegram/db/puppet.py index 8fde9c66..37b7a0f7 100644 --- a/mautrix_telegram/db/puppet.py +++ b/mautrix_telegram/db/puppet.py @@ -43,6 +43,7 @@ class Puppet: displayname_quality: int disable_updates: bool username: str | None + phone: str | None photo_id: str | None avatar_url: ContentURI | None name_set: bool @@ -65,7 +66,7 @@ class Puppet: columns: ClassVar[str] = ( "id, is_registered, displayname, displayname_source, displayname_contact, " - "displayname_quality, disable_updates, username, photo_id, avatar_url, " + "displayname_quality, disable_updates, username, phone, photo_id, avatar_url, " "name_set, avatar_set, is_bot, is_channel, custom_mxid, access_token, next_batch, base_url" ) @@ -105,6 +106,7 @@ class Puppet: self.displayname_quality, self.disable_updates, self.username, + self.phone, self.photo_id, self.avatar_url, self.name_set, @@ -121,9 +123,9 @@ class Puppet: q = """ UPDATE puppet SET is_registered=$2, displayname=$3, displayname_source=$4, displayname_contact=$5, - displayname_quality=$6, disable_updates=$7, username=$8, photo_id=$9, - avatar_url=$10, name_set=$11, avatar_set=$12, is_bot=$13, is_channel=$14, - custom_mxid=$15, access_token=$16, next_batch=$17, base_url=$18 + displayname_quality=$6, disable_updates=$7, username=$8, phone=$9, photo_id=$10, + avatar_url=$11, name_set=$12, avatar_set=$13, is_bot=$14, is_channel=$15, + custom_mxid=$16, access_token=$17, next_batch=$18, base_url=$19 WHERE id=$1 """ await self.db.execute(q, *self._values) @@ -132,8 +134,9 @@ class Puppet: q = """ INSERT INTO puppet ( id, is_registered, displayname, displayname_source, displayname_contact, - displayname_quality, disable_updates, username, photo_id, avatar_url, name_set, + displayname_quality, disable_updates, username, phone, photo_id, avatar_url, name_set, avatar_set, is_bot, is_channel, custom_mxid, access_token, next_batch, base_url - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18) + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, + $19) """ await self.db.execute(q, *self._values) diff --git a/mautrix_telegram/db/upgrade/__init__.py b/mautrix_telegram/db/upgrade/__init__.py index 6ebc258b..2f3fe794 100644 --- a/mautrix_telegram/db/upgrade/__init__.py +++ b/mautrix_telegram/db/upgrade/__init__.py @@ -9,4 +9,5 @@ from . import ( v04_disappearing_messages, v05_channel_ghosts, v06_puppet_avatar_url, + v07_puppet_phone_number, ) diff --git a/mautrix_telegram/db/upgrade/v00_latest_revision.py b/mautrix_telegram/db/upgrade/v00_latest_revision.py index f03f9a9b..4e88f737 100644 --- a/mautrix_telegram/db/upgrade/v00_latest_revision.py +++ b/mautrix_telegram/db/upgrade/v00_latest_revision.py @@ -16,7 +16,7 @@ from mautrix.util.async_db import Connection -async def create_v6_tables(conn: Connection) -> int: +async def create_v7_tables(conn: Connection) -> int: await conn.execute( """CREATE TABLE "user" ( mxid TEXT PRIMARY KEY, @@ -98,6 +98,7 @@ async def create_v6_tables(conn: Connection) -> int: displayname_quality INTEGER NOT NULL DEFAULT 0, disable_updates BOOLEAN NOT NULL DEFAULT false, username TEXT, + phone TEXT, photo_id TEXT, avatar_url TEXT, name_set BOOLEAN NOT NULL DEFAULT false, diff --git a/mautrix_telegram/db/upgrade/v01_initial_revision.py b/mautrix_telegram/db/upgrade/v01_initial_revision.py index bcbd034b..a3592bd7 100644 --- a/mautrix_telegram/db/upgrade/v01_initial_revision.py +++ b/mautrix_telegram/db/upgrade/v01_initial_revision.py @@ -18,7 +18,7 @@ from __future__ import annotations from mautrix.util.async_db import Connection, Scheme from . import upgrade_table -from .v00_latest_revision import create_v6_tables +from .v00_latest_revision import create_v7_tables legacy_version_query = "SELECT version_num FROM alembic_version" last_legacy_version = "bfc0a39bfe02" @@ -34,9 +34,9 @@ def table_exists(scheme: str, name: str) -> str: async def first_upgrade_target(conn: Connection, scheme: str) -> int: is_legacy = await conn.fetchval(table_exists(scheme, "alembic_version")) - # If it's a legacy db, the upgrade process will go to v1 and run each migration up to v6. - # If it's a new db, we'll create the v6 tables directly (see the create_v6_tables call). - return 1 if is_legacy else 6 + # If it's a legacy db, the upgrade process will go to v1 and run each migration up to v7. + # If it's a new db, we'll create the v7 tables directly (see the create_v7_tables call). + return 1 if is_legacy else 7 @upgrade_table.register(description="Initial asyncpg revision", upgrades_to=first_upgrade_target) @@ -46,7 +46,7 @@ async def upgrade_v1(conn: Connection, scheme: str) -> int: await migrate_legacy_to_v1(conn, scheme) return 1 else: - return await create_v6_tables(conn) + return await create_v7_tables(conn) async def drop_constraints(conn: Connection, table: str, contype: str) -> None: diff --git a/mautrix_telegram/db/upgrade/v07_puppet_phone_number.py b/mautrix_telegram/db/upgrade/v07_puppet_phone_number.py new file mode 100644 index 00000000..87dd4b9c --- /dev/null +++ b/mautrix_telegram/db/upgrade/v07_puppet_phone_number.py @@ -0,0 +1,23 @@ +# mautrix-telegram - A Matrix-Telegram puppeting bridge +# Copyright (C) 2022 Tulir Asokan +# +# 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 . +from mautrix.util.async_db import Connection + +from . import upgrade_table + + +@upgrade_table.register(description="Store phone number in puppet table") +async def upgrade_v7(conn: Connection) -> None: + await conn.execute("ALTER TABLE puppet ADD COLUMN phone TEXT") diff --git a/mautrix_telegram/puppet.py b/mautrix_telegram/puppet.py index da338884..e94b9bab 100644 --- a/mautrix_telegram/puppet.py +++ b/mautrix_telegram/puppet.py @@ -73,6 +73,7 @@ class Puppet(DBPuppet, BasePuppet): displayname_quality: int = 0, disable_updates: bool = False, username: str | None = None, + phone: str | None = None, photo_id: str | None = None, avatar_url: ContentURI | None = None, name_set: bool = False, @@ -93,6 +94,7 @@ class Puppet(DBPuppet, BasePuppet): displayname_quality=displayname_quality, disable_updates=disable_updates, username=username, + phone=phone, photo_id=photo_id, avatar_url=avatar_url, name_set=name_set, @@ -134,7 +136,9 @@ class Puppet(DBPuppet, BasePuppet): return { "name": self.displayname, "username": self.username, + "phone": f"+{self.phone.lstrip('+')}" if self.phone else None, "is_bot": self.is_bot, + "avatar_url": self.avatar_url, } @property @@ -261,6 +265,10 @@ class Puppet(DBPuppet, BasePuppet): self.username = info.username changed = True + if getattr(info, "phone", None) and self.phone != info.phone: + self.phone = info.phone + changed = True + if not self.disable_updates: try: changed = await self.update_displayname(source, info) or changed