Data migration for UserUniqueLogin/IpAddress (#19124)

* Data migration for UserUniqueLogin/IpAddress

* Update the model as well

* Linting

* Only do index if it does not exist

* Seasoning

* Update tests

---------

Co-authored-by: Ee Durbin <ewdurbin@gmail.com>

Author: di

tests/common/db/accounts.py

@@ -19,6 +19,7 @@
 
 from ...common.constants import REMOTE_ADDR
 from .base import WarehouseFactory
+from .ip_addresses import IpAddressFactory
 
 fake = faker.Faker()
 
@@ -140,3 +141,10 @@ class Meta:
 
     user = factory.SubFactory(UserFactory)
     ip_address = REMOTE_ADDR
+    ip_address_id = factory.LazyAttribute(
+        lambda o: (
+            IpAddressFactory.create(ip_address=o.ip_address).id
+            if o.ip_address is not None
+            else None
+        )
+    )

tests/unit/accounts/test_services.py

@@ -2095,10 +2095,12 @@ def test_device_is_pending_not_expired(self, user_service, monkeypatch):
 
     def test_device_is_pending_and_expired(self, user_service, monkeypatch):
         user = UserFactory.create(with_verified_primary_email=True)
+        ip_address = IpAddressFactory.create(ip_address=REMOTE_ADDR)
         UserUniqueLoginFactory.create(
             user=user,
-            ip_address=REMOTE_ADDR,
             status="pending",
+            ip_address=str(ip_address.ip_address),
+            ip_address_id=ip_address.id,
             created=datetime.datetime(1970, 1, 1),
             expires=datetime.datetime(1970, 1, 1),
         )
@@ -2115,7 +2117,7 @@ def test_device_is_pending_and_expired(self, user_service, monkeypatch):
                 )
             },
             find_service=lambda *a, **kw: token_service,
-            ip_address=IpAddressFactory.create(ip_address=REMOTE_ADDR),
+            ip_address=ip_address,
         )
 
         assert not user_service.device_is_known(user.id, user_service.request)

tests/unit/accounts/test_views.py

@@ -641,7 +641,8 @@ def test_login_with_remembered_device_confirms_unique_login(
 
         UserUniqueLoginFactory.create(
             user=user,
-            ip_address=db_request.remote_addr,
+            ip_address=str(db_request.ip_address.ip_address),
+            ip_address_id=db_request.ip_address.id,
             status=UniqueLoginStatus.PENDING,
         )
 
@@ -699,7 +700,8 @@ def test_login_updates_last_used(self, monkeypatch, db_request, pyramid_services
         past_timestamp = datetime.datetime(1970, 1, 1)
         UserUniqueLoginFactory.create(
             user=user,
-            ip_address=db_request.remote_addr,
+            ip_address=str(db_request.ip_address.ip_address),
+            ip_address_id=db_request.ip_address.id,
             status=UniqueLoginStatus.CONFIRMED,
             last_used=past_timestamp,
         )
@@ -5507,7 +5509,9 @@ def test_ip_address_mismatch(self, db_request):
     def test_success(self, monkeypatch, db_request):
         user = UserFactory.create(last_login=datetime.datetime.now(datetime.UTC))
         unique_login = UserUniqueLoginFactory.create(
-            user=user, ip_address=db_request.remote_addr
+            user=user,
+            ip_address=str(db_request.ip_address.ip_address),
+            ip_address_id=db_request.ip_address.id,
         )
         db_request.user = None
         db_request.params = {"token": "foo"}

tests/unit/admin/views/test_ipaddresses.py

@@ -64,12 +64,15 @@ def test_ip_address_found_no_unique_logins(self, db_request):
         assert result == {"ip_address": ip_address, "unique_logins": []}
 
     def test_ip_address_found_with_unique_logins(self, db_request):
-        ip_address = IpAddressFactory()
         unique_login = UserUniqueLoginFactory.create(
-            ip_address=str(ip_address.ip_address)
+            ip_address=str(db_request.ip_address.ip_address),
+            ip_address_id=str(db_request.ip_address.id),
         )
-        db_request.matchdict["ip_address"] = str(ip_address.ip_address)
+        db_request.matchdict["ip_address"] = str(db_request.ip_address.ip_address)
 
         result = ip_views.ip_address_detail(db_request)
 
-        assert result == {"ip_address": ip_address, "unique_logins": [unique_login]}
+        assert result == {
+            "ip_address": db_request.ip_address,
+            "unique_logins": [unique_login],
+        }

warehouse/accounts/models.py

@@ -514,7 +514,7 @@ class UserUniqueLogin(db.Model):
 
     ip_address_id: Mapped[int] = mapped_column(
         ForeignKey("ip_addresses.id", onupdate="CASCADE", ondelete="CASCADE"),
-        nullable=True,
+        nullable=False,
         index=True,
     )
     ip_address: Mapped[str] = mapped_column(String, nullable=False)

warehouse/migrations/versions/be443e514e3e_data_migration_for_useruniquelogin_ip_.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Data migration for UserUniqueLogin.ip_address_id
+
+Revision ID: be443e514e3e
+Revises: df52c3746740
+Create Date: 2025-12-02 17:32:29.770684
+"""
+
+
+import os
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "be443e514e3e"
+down_revision = "df52c3746740"
+
+
+def _get_remaining_logins_to_update(conn):
+    return conn.execute(
+        sa.text("SELECT COUNT(*) FROM user_unique_logins WHERE ip_address_id IS NULL")
+    ).scalar_one()
+
+
+def _get_remaining_ips_to_insert(conn):
+    return conn.execute(
+        sa.text(
+            """
+SELECT COUNT(DISTINCT user_unique_logins.ip_address)
+FROM user_unique_logins
+LEFT JOIN ip_addresses ON user_unique_logins.ip_address::inet = ip_addresses.ip_address
+WHERE ip_addresses.id IS NULL AND user_unique_logins.ip_address IS NOT NULL
+"""
+        )
+    ).scalar_one()
+
+
+def upgrade():
+    op.execute("SET statement_timeout = 120000")
+    op.execute("SET lock_timeout = 120000")
+
+    op.create_index(
+        "ix_user_unique_logins_ip_address_migration",
+        "user_unique_logins",
+        ["ip_address"],
+        unique=False,
+        if_not_exists=True,
+    )
+
+    bind = op.get_bind()
+    batch_size = 1000
+    salt = os.environ.get("WAREHOUSE_IP_SALT")
+
+    while _get_remaining_ips_to_insert(bind) > 0:
+        bind.execute(
+            sa.text(
+                f"""
+INSERT INTO ip_addresses (ip_address, hashed_ip_address)
+SELECT
+    DISTINCT user_unique_logins.ip_address::inet,
+    encode(
+        digest(
+            CONCAT(user_unique_logins.ip_address, '{salt}'),
+            'sha256'
+        ),
+        'hex'
+    )
+FROM user_unique_logins
+LEFT JOIN ip_addresses ON user_unique_logins.ip_address::inet = ip_addresses.ip_address
+WHERE ip_addresses.id IS NULL AND user_unique_logins.ip_address IS NOT NULL
+LIMIT :batch_size
+"""
+            ),
+            {"batch_size": batch_size},
+        )
+        bind.commit()
+
+    while _get_remaining_logins_to_update(bind) > 0:
+        bind.execute(
+            sa.text(
+                """
+UPDATE user_unique_logins
+SET ip_address_id = ip_addresses.id
+FROM ip_addresses
+WHERE
+    user_unique_logins.ip_address::inet = ip_addresses.ip_address AND
+    user_unique_logins.ip_address_id IS NULL AND
+    user_unique_logins.id IN (
+        SELECT id
+        FROM user_unique_logins
+        WHERE ip_address_id IS NULL
+        ORDER BY id
+        LIMIT :batch_size
+    )
+"""
+            ),
+            {"batch_size": batch_size},
+        )
+        bind.commit()
+
+    # Finally make the ip_address_id column non-nullable
+    op.alter_column("user_unique_logins", "ip_address_id", nullable=False)
+
+    op.drop_index(
+        "ix_user_unique_logins_ip_address_migration",
+        "user_unique_logins",
+        if_exists=True,
+    )
+
+
+def downgrade():
+    op.alter_column("user_unique_logins", "ip_address_id", nullable=True)