2 files changed, 128 insertions, 0 deletions

diff --git a/src/mailman/database/alembic/versions/47294d3a604_pendable_indexes.py b/src/mailman/database/alembic/versions/47294d3a604_pendable_indexes.py
index 6ad5c7bd5..b6e64c8bb 100644
--- a/src/mailman/database/alembic/versions/47294d3a604_pendable_indexes.py
+++ b/src/mailman/database/alembic/versions/47294d3a604_pendable_indexes.py
@@ -13,9 +13,30 @@ Create Date: 2015-12-02 11:46:47.295174
 revision = '47294d3a604'
 down_revision = '33bc0099223'
 
+import json
+import sqlalchemy as sa
 from alembic import op
 
 
+TYPE_CLUES = {
+    'member_id': 'probe',
+    'token_owner': 'subscription',
+    '_mod_message_id': 'data',
+}
+
+pended_table = sa.sql.table(
+    'pended',
+    sa.sql.column('id', sa.Integer),
+    )
+
+keyvalue_table = sa.sql.table(
+    'pendedkeyvalue',
+    sa.sql.column('id', sa.Integer),
+    sa.sql.column('key', sa.Unicode),
+    sa.sql.column('value', sa.Unicode),
+    sa.sql.column('pended_id', sa.Integer),
+    )
+
 def upgrade():
     op.create_index(
         op.f('ix_pended_expiration_date'), 'pended', ['expiration_date'],
@@ -26,6 +47,35 @@ def upgrade():
     op.create_index(
         op.f('ix_pendedkeyvalue_value'), 'pendedkeyvalue', ['value'],
         unique=False)
+    # Data migration.
+    connection = op.get_bind()
+    for pended_result in connection.execute(pended_table.select()).fetchall():
+        pended_id = pended_result['id']
+        keyvalues = connection.execute(keyvalue_table.select().where(
+            keyvalue_table.c.pended_id == pended_id
+            )).fetchall()
+        kv_type = [kv for kv in keyvalues if kv['key'] == 'type']
+        if kv_type:
+            # Convert existing type keys from JSON to plain text.
+            kv_type = kv_type[0] # the (pended_id, key) tuple is unique.
+            try:
+                new_value = json.loads(kv_type['value'])
+            except ValueError:
+                pass # New-style entry (or already converted).
+            else:
+                connection.execute(keyvalue_table.update().where(
+                    keyvalue_table.c.id == kv_type['id']
+                ).values(value=new_value))
+        else:
+            # Detect the type and add the corresponding type key.
+            keys = [kv['key'] for kv in keyvalues]
+            for clue_key, clue_type in TYPE_CLUES.items():
+                if clue_key not in keys:
+                    continue
+                # We found the type, update the DB.
+                connection.execute(keyvalue_table.insert().values(
+                    key='type', value=clue_type, pended_id=pended_id))
+                break
 
 
 def downgrade():
@@ -33,3 +83,17 @@ def downgrade():
     op.drop_index(op.f('ix_pendedkeyvalue_key'), table_name='pendedkeyvalue')
     op.drop_index(op.f('ix_pended_token'), table_name='pended')
     op.drop_index(op.f('ix_pended_expiration_date'), table_name='pended')
+    # Data migration.
+    connection = op.get_bind()
+    # Remove the introduced type keys.
+    connection.execute(keyvalue_table.delete().where(sa.and_(
+        keyvalue_table.c.key == 'type',
+        keyvalue_table.c.value.in_(TYPE_CLUES.values())
+        )))
+    # Convert the other type keys to JSON.
+    keyvalues = connection.execute(keyvalue_table.select().where(
+        keyvalue_table.c.key == 'type')).fetchall()
+    for keyvalue in keyvalues:
+        connection.execute(keyvalue_table.update().where(
+            keyvalue_table.c.id == keyvalue['id']
+            ).values(value=json.dumps(keyvalue['value'])))
diff --git a/src/mailman/database/tests/test_migrations.py b/src/mailman/database/tests/test_migrations.py
index 91f19bfb2..3140ed1ab 100644
--- a/src/mailman/database/tests/test_migrations.py
+++ b/src/mailman/database/tests/test_migrations.py
@@ -100,3 +100,67 @@ class TestMigrations(unittest.TestCase):
             header_match_table.select()).fetchall()
         self.assertEqual(results,
             [(1, hm[0], hm[1]) for hm in test_header_matches])
+
+    def test_47294d3a604_pendable_keyvalues(self):
+        # We have 5 pended items:
+        # - one is a probe request
+        # - one is a subscription request
+        # - one is a moderation request
+        # - one is a held message
+        # - one is a registration request in the new format
+        # The first three used to have no 'type' key and must be properly
+        # typed, the held message used to have a type key, but in JSON, and
+        # must be converted.
+        pended_table = sa.sql.table(
+            'pended',
+            sa.sql.column('id', sa.Integer),
+            )
+        keyvalue_table = sa.sql.table(
+            'pendedkeyvalue',
+            sa.sql.column('id', sa.Integer),
+            sa.sql.column('key', sa.Unicode),
+            sa.sql.column('value', sa.Unicode),
+            sa.sql.column('pended_id', sa.Integer),
+            )
+        def get_from_db():
+            results = {}
+            for i in range(1, 6):
+                query = sa.sql.select(
+                    [keyvalue_table.c.key, keyvalue_table.c.value]
+                ).where(
+                    keyvalue_table.c.pended_id == i
+                )
+                results[i] = dict([
+                    (r['key'], r['value']) for r in
+                    config.db.store.execute(query).fetchall()
+                    ])
+            return results
+        # Start at the previous revision
+        alembic.command.downgrade(alembic_cfg, '33bc0099223')
+        for i in range(1, 6):
+            config.db.store.execute(pended_table.insert().values(id=i))
+        config.db.store.execute(keyvalue_table.insert().values([
+            {'pended_id': 1, 'key': 'member_id', 'value': 'test-value'},
+            {'pended_id': 2, 'key': 'token_owner', 'value': 'test-value'},
+            {'pended_id': 3, 'key': '_mod_message_id', 'value': 'test-value'},
+            {'pended_id': 4, 'key': 'type', 'value': '"held message"'},
+            {'pended_id': 5, 'key': 'type', 'value': 'registration'},
+            ]))
+        config.db.store.commit()
+        # Upgrading.
+        alembic.command.upgrade(alembic_cfg, '47294d3a604')
+        results = get_from_db()
+        for i in range(1, 5):
+            self.assertIn('type', results[i])
+        self.assertEqual(results[1]['type'], 'probe')
+        self.assertEqual(results[2]['type'], 'subscription')
+        self.assertEqual(results[3]['type'], 'data')
+        self.assertEqual(results[4]['type'], 'held message')
+        self.assertEqual(results[5]['type'], 'registration')
+        # Downgrading.
+        alembic.command.downgrade(alembic_cfg, '33bc0099223')
+        results = get_from_db()
+        for i in range(1, 4):
+            self.assertNotIn('type', results[i])
+        self.assertEqual(results[4]['type'], '"held message"')
+        self.assertEqual(results[5]['type'], '"registration"')