(allura) 01/03: Fix convert plaintext to encrypted script to support nested fields

[br...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/allura.git

commit 72a31893d6a3176b47afe9b421bbddf599070c22
Author: Carlos Cruz <ca...@slashdotmedia.com>
AuthorDate: Wed Apr 15 21:00:59 2026 +0000

    Fix convert plaintext to encrypted script to support nested fields
---
 scripts/convert_encrypted_field.py | 105 +++++++++++++++++++++++++++++--------
 1 file changed, 82 insertions(+), 23 deletions(-)

diff --git a/scripts/convert_encrypted_field.py b/scripts/convert_encrypted_field.py
index 04a5b3bfe..46e2f61ba 100644
--- a/scripts/convert_encrypted_field.py
+++ b/scripts/convert_encrypted_field.py
@@ -21,23 +21,62 @@
 
 import defopt
 
-from allura.lib.utils import chunked_find
 from ming import schema
-from ming.odm import session
+from ming.odm import mapper, session
 from ming.odm.declarative import MappedClass
 from ming.odm.property import FieldProperty, DecryptedProperty
 
 
-def _encrypted_field_name(plain_field_name: str) -> str:
-    path = plain_field_name.split('.')
-    path[-1] = f'{path[-1]}_encrypted'
-    return '.'.join(path)
+CHUNK_SIZE = 1000
 
 
-def _get_field_value(record, field_name: str):
-    value = record
-    for part in field_name.split('.'):
-        value = value[part]
+class MissingFieldPathError(Exception):
+    pass
+
+
+def _default_encrypted_field_name(plain_field_name: str) -> str:
+    parent, sep, leaf = plain_field_name.rpartition('.')
+    return f'{parent}{sep}{leaf}_encrypted'
+
+
+def _split_field_path(field_name: str) -> list[str]:
+    field_path = field_name.split('.')
+    assert field_path and all(field_path), f'Invalid dotted field path: {field_name!r}'
+    return field_path
+
+
+def _schema_for_field_path(Model: type[MappedClass], field_name: str):
+    field_path = _split_field_path(field_name)
+    top_level_name = field_path[0]
+    try:
+        top_level_prop = getattr(Model, top_level_name)
+    except AttributeError as e:
+        raise MissingFieldPathError(
+            f'Missing field path {field_name!r}; no top-level field {top_level_name!r}') from e
+    assert isinstance(top_level_prop, FieldProperty)
+
+    current_schema = top_level_prop.field.schema
+    for i, path_part in enumerate(field_path[1:], start=1):
+        traversed = '.'.join(field_path[:i])
+        if isinstance(current_schema, schema.Array):
+            raise AssertionError(
+                f'Nested array paths are not supported: {field_name!r} (array at {traversed!r})')
+        if not isinstance(current_schema, schema.Object):
+            raise AssertionError(
+                f'Invalid nested field path {field_name!r}; {traversed!r} is not an object field')
+        if path_part not in current_schema.fields:
+            raise MissingFieldPathError(
+                f'Invalid nested field path {field_name!r}; missing key {traversed + "." + path_part!r}')
+
+        current_schema = current_schema.fields[path_part]
+
+    return current_schema
+
+
+def _get_nested_value(rec: dict, field_name: str):
+    value = rec
+    for path_part in _split_field_path(field_name):
+        value = value[path_part]
     return value
 
 
@@ -45,23 +84,31 @@ def main(class_name: str, plain_field_name: str,
          *, remove_unencrypted: bool = False, redo_all: bool = False):
     """
     :param class_name: full class name, e.g. allura.model.user.User
-    :param plain_field_name: name of the unencrypted field, e.g. display_name or author.username
+    :param plain_field_name: name of the unencrypted field, e.g. display_name
     :param remove_unencrypted: WARNING only run this after your codebase is already on the latest code
     :param redo_all: re-encrypt records that already have encrypted values (in case they changed since last run)
     """
-    encrypted_field_name = _encrypted_field_name(plain_field_name)
-    is_nested_field = '.' in plain_field_name
+    encrypted_field_name = _default_encrypted_field_name(plain_field_name)
 
     module_name, class_basename = class_name.rsplit('.', 1)
     module = import_module(module_name)
     Model: type[MappedClass] = getattr(module, class_basename)
 
     # sanity checks that the fields are correct and ready
-    if not is_nested_field:
-        encr_prop = getattr(Model, encrypted_field_name)
-        assert isinstance(encr_prop, FieldProperty)
-        assert encr_prop.field.type == schema.Binary
-        if remove_unencrypted:
+    try:
+        encr_schema = _schema_for_field_path(Model, encrypted_field_name)
+    except MissingFieldPathError:
+        # pre-migration support: encrypted field might not yet be present in
+        # model schema, but we can still create it with raw MongoDB updates
+        encr_schema = None
+    if encr_schema is not None:
+        assert isinstance(encr_schema, schema.Binary)
+    elif remove_unencrypted:
+        raise AssertionError(
+            f'Cannot use --remove-unencrypted because {encrypted_field_name!r} '
+            'is not present in model schema yet')
+    if remove_unencrypted:
+        if '.' not in plain_field_name:
             plain_prop = Model.__dict__[plain_field_name]  # getattr() better but needs Ming fix released
             assert isinstance(plain_prop, DecryptedProperty)
             assert plain_prop.encrypted_field == encrypted_field_name
@@ -90,16 +137,28 @@ def main(class_name: str, plain_field_name: str,
     }
     if not redo_all:
         q[encrypted_field_name] = None
+
+    m = mapper(Model)
+    raw_collection = sess.impl.db[m.collection.m.collection_name]
+
     count = 0
-    for chunk in chunked_find(Model, q):
-        for rec in chunk:
-            val = _get_field_value(rec, plain_field_name)
+    projection = {'_id': 1, plain_field_name: 1}
+    last_id = None
+    while True:
+        chunk_q = {'$and': [q, {'_id': {'$gt': last_id}}]} if last_id is not None else q
+        docs = list(raw_collection.find(chunk_q, projection).sort('_id', 1).limit(CHUNK_SIZE))
+        if not docs:
+            break
+
+        for rec_doc in docs:
+            val = _get_nested_value(rec_doc, plain_field_name)
             encr_val = Model.encr(val)
-            Model.query.update({'_id': rec['_id']}, {
+            raw_collection.update_one({'_id': rec_doc['_id']}, {
                 '$set': {encrypted_field_name: encr_val},
             })
             count += 1
-            sess.expunge(rec)
+
+        last_id = docs[-1]['_id']
         print(f'Converted {count} so far...')
     if not count:
         print(f'Did not find any {class_name} records with {plain_field_name} values to encrypt')