1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
"""several node keys per account: one per machine the person runs a node on
Revision ID: f8a9b0c1d2e3
Revises: e7f8a9b0c1d2
"""
import uuid
from collections.abc import Sequence
from datetime import UTC, datetime
import sqlalchemy as sa
from alembic import op
revision: str = "f8a9b0c1d2e3"
down_revision: str | Sequence[str] | None = "e7f8a9b0c1d2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
node_keys = op.create_table(
"node_keys",
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("user_id", sa.String(36), sa.ForeignKey("users.id"), nullable=False),
sa.Column("pk_node_ed25519", sa.String(64), nullable=False),
sa.Column("linked_at", sa.DateTime(timezone=True)),
)
op.create_index("uq_node_keys_user_key", "node_keys",
["user_id", "pk_node_ed25519"], unique=True)
# The key each account had linked keeps working: it becomes that account's
# first row, so no node has to be linked again after the upgrade.
conn = op.get_bind()
linked = conn.execute(sa.text(
"SELECT id, pk_node_ed25519 FROM users WHERE pk_node_ed25519 IS NOT NULL "
"AND pk_node_ed25519 <> ''")).fetchall()
now = datetime.now(UTC)
if linked:
op.bulk_insert(node_keys, [
{"id": str(uuid.uuid4()), "user_id": uid, "pk_node_ed25519": pk,
"linked_at": now}
for uid, pk in linked])
with op.batch_alter_table("users") as batch:
batch.drop_column("pk_node_ed25519")
def downgrade() -> None:
with op.batch_alter_table("users") as batch:
batch.add_column(sa.Column("pk_node_ed25519", sa.String(64), nullable=True))
# One column holds one key: the most recently linked one is kept, which is
# the node the account would have had linked last before this revision.
conn = op.get_bind()
rows = conn.execute(sa.text(
"SELECT user_id, pk_node_ed25519 FROM node_keys ORDER BY linked_at")).fetchall()
latest = {uid: pk for uid, pk in rows}
for uid, pk in latest.items():
conn.execute(sa.text("UPDATE users SET pk_node_ed25519 = :pk WHERE id = :uid"),
{"pk": pk, "uid": uid})
op.drop_index("uq_node_keys_user_key", table_name="node_keys")
op.drop_table("node_keys")
|