# mautrix-telegram - A Matrix-Telegram puppeting bridge # Copyright (C) 2021 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 __future__ import annotations from mautrix.util.async_db import Connection, Scheme from . import upgrade_table from .v00_latest_revision import create_v6_tables legacy_version_query = "SELECT version_num FROM alembic_version" last_legacy_version = "bfc0a39bfe02" def table_exists(scheme: str, name: str) -> str: if scheme == Scheme.SQLITE: return f"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='{name}')" elif scheme in (Scheme.POSTGRES, Scheme.COCKROACH): return f"SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_name='{name}')" raise RuntimeError("unsupported database scheme") 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 @upgrade_table.register(description="Initial asyncpg revision", upgrades_to=first_upgrade_target) async def upgrade_v1(conn: Connection, scheme: str) -> int: is_legacy = await conn.fetchval(table_exists(scheme, "alembic_version")) if is_legacy: await migrate_legacy_to_v1(conn, scheme) return 1 else: return await create_v6_tables(conn) async def drop_constraints(conn: Connection, table: str, contype: str) -> None: q = ( "SELECT conname FROM pg_constraint con INNER JOIN pg_class rel ON rel.oid=con.conrelid " f"WHERE rel.relname='{table}' AND contype='{contype}'" ) names = [row["conname"] for row in await conn.fetch(q)] drops = ", ".join(f"DROP CONSTRAINT {name}" for name in names) await conn.execute(f"ALTER TABLE {table} {drops}") async def migrate_legacy_to_v1(conn: Connection, scheme: str) -> None: legacy_version = await conn.fetchval(legacy_version_query) if legacy_version != last_legacy_version: raise RuntimeError( "Legacy database is not on last version. " "Please upgrade the old database with alembic or drop it completely first." ) if scheme != "sqlite": await drop_constraints(conn, "contact", contype="f") await conn.execute( """ ALTER TABLE contact ADD CONSTRAINT contact_user_fkey FOREIGN KEY (contact) REFERENCES puppet(id) ON DELETE CASCADE ON UPDATE CASCADE, ADD CONSTRAINT contact_contact_fkey FOREIGN KEY ("user") REFERENCES "user"(tgid) ON DELETE CASCADE ON UPDATE CASCADE """ ) await drop_constraints(conn, "telethon_sessions", contype="p") await conn.execute( """ ALTER TABLE telethon_sessions ADD CONSTRAINT telethon_sessions_pkey PRIMARY KEY (session_id) """ ) await drop_constraints(conn, "telegram_file", contype="f") await conn.execute( """ ALTER TABLE telegram_file ADD CONSTRAINT fk_file_thumbnail FOREIGN KEY (thumbnail) REFERENCES telegram_file(id) ON UPDATE CASCADE ON DELETE SET NULL """ ) await conn.execute("ALTER TABLE puppet ALTER COLUMN id DROP IDENTITY IF EXISTS") await conn.execute("ALTER TABLE puppet ALTER COLUMN id DROP DEFAULT") await conn.execute("DROP SEQUENCE IF EXISTS puppet_id_seq") await conn.execute("ALTER TABLE bot_chat ALTER COLUMN id DROP IDENTITY IF EXISTS") await conn.execute("ALTER TABLE bot_chat ALTER COLUMN id DROP DEFAULT") await conn.execute("DROP SEQUENCE IF EXISTS bot_chat_id_seq") await conn.execute("ALTER TABLE portal ALTER COLUMN config TYPE jsonb USING config::jsonb") await conn.execute( "ALTER TABLE telegram_file ALTER COLUMN decryption_info TYPE jsonb " "USING decryption_info::jsonb" ) await varchar_to_text(conn) else: await conn.execute( """CREATE TABLE telethon_sessions_new ( session_id TEXT PRIMARY KEY, dc_id INTEGER, server_address TEXT, port INTEGER, auth_key bytea )""" ) await conn.execute( """ INSERT INTO telethon_sessions_new (session_id, dc_id, server_address, port, auth_key) SELECT session_id, dc_id, server_address, port, auth_key FROM telethon_sessions """ ) await conn.execute("DROP TABLE telethon_sessions") await conn.execute("ALTER TABLE telethon_sessions_new RENAME TO telethon_sessions") await update_state_store(conn, scheme) await conn.execute('ALTER TABLE "user" ADD COLUMN is_bot BOOLEAN NOT NULL DEFAULT false') await conn.execute("ALTER TABLE puppet RENAME COLUMN matrix_registered TO is_registered") await conn.execute("DROP TABLE telethon_version") await conn.execute("DROP TABLE alembic_version") async def update_state_store(conn: Connection, scheme: str) -> None: # The Matrix state store already has more or less the correct schema, so set the version await conn.execute("CREATE TABLE mx_version (version INTEGER PRIMARY KEY)") await conn.execute("INSERT INTO mx_version (version) VALUES (2)") await conn.execute("UPDATE mx_user_profile SET membership='LEAVE' WHERE membership='LEFT'") if scheme != "sqlite": # Also add the membership type on postgres await conn.execute( "CREATE TYPE membership AS ENUM ('join', 'leave', 'invite', 'ban', 'knock')" ) await conn.execute( "ALTER TABLE mx_user_profile ALTER COLUMN membership TYPE membership " "USING LOWER(membership)::membership" ) else: # On SQLite there's no custom type, but we still want to lowercase everything await conn.execute("UPDATE mx_user_profile SET membership=LOWER(membership)") async def varchar_to_text(conn: Connection) -> None: columns_to_adjust = { "user": ("mxid", "tg_username", "tg_phone"), "portal": ( "peer_type", "mxid", "username", "title", "about", "photo_id", "avatar_url", "config", ), "message": ("mxid", "mx_room"), "puppet": ( "displayname", "username", "photo_id", "access_token", "custom_mxid", "next_batch", "base_url", ), "bot_chat": ("type",), "telegram_file": ("id", "mxc", "mime_type", "thumbnail"), # Phone is a bigint in the old schema, which is safe, but we don't do math on it, # so let's change it to a string "telethon_entities": ("session_id", "username", "name", "phone"), "telethon_sent_files": ("session_id",), "telethon_sessions": ("session_id", "server_address"), "telethon_update_state": ("session_id",), "mx_room_state": ("room_id",), "mx_user_profile": ("room_id", "user_id", "displayname", "avatar_url"), } for table, columns in columns_to_adjust.items(): for column in columns: await conn.execute(f'ALTER TABLE "{table}" ALTER COLUMN {column} TYPE TEXT')