Mercurial > libervia-backend
annotate libervia/backend/memory/migration/versions/fe3a02cb4bec_convert_legacypickle_columns_to_json.py @ 4228:79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Those plugins are using optional modules, and if their importation fails, there is
probably no data to migrate either.
| author | Goffi <goffi@goffi.org> |
|---|---|
| date | Sat, 09 Mar 2024 15:13:14 +0100 |
| parents | 1a7a3e4b52a4 |
| children | 0d7bb4df2343 |
| rev | line source |
|---|---|
| 4212 | 1 """convert LegacyPickle columns to JSON |
| 2 | |
| 3 Revision ID: fe3a02cb4bec | |
| 4 Revises: 610345f77e75 | |
| 5 Create Date: 2024-02-22 14:55:59.993983 | |
| 6 | |
| 7 """ | |
| 8 from alembic import op | |
| 9 import sqlalchemy as sa | |
| 10 import pickle | |
| 11 import json | |
|
4228
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
12 try: |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
13 from libervia.backend.plugins.plugin_xep_0373 import PublicKeyMetadata |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
14 except Exception: |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
15 PublicKeyMetadata = None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
16 print( |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
17 "Warning: Can't import XEP-0373, its data won't be updated. It's probably not " |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
18 "used on this installation." |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
19 ) |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
20 try: |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
21 from libervia.backend.plugins.plugin_xep_0384 import TrustMessageCacheEntry |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
22 except Exception: |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
23 TrustMessageCacheEntry = None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
24 print( |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
25 "Warning: Can't import XEP-0384, its data won't be updated. It's probably not " |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
26 "used on this installation." |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
27 ) |
| 4212 | 28 |
| 29 # revision identifiers, used by Alembic. | |
| 30 revision = "fe3a02cb4bec" | |
| 31 down_revision = "610345f77e75" | |
| 32 branch_labels = None | |
| 33 depends_on = None | |
| 34 | |
| 35 | |
| 36 def convert_pickle_to_json(value, table, primary_keys): | |
| 37 """Convert pickled data to JSON, handling potential errors.""" | |
| 38 if value is None: | |
| 39 return None | |
| 40 try: | |
| 41 # some values are converted to bytes with LegacyPickle | |
| 42 if isinstance(value, str): | |
| 43 value = value.encode() | |
| 44 try: | |
| 45 deserialized = pickle.loads(value, encoding="utf-8") | |
| 46 except ModuleNotFoundError: | |
| 47 deserialized = pickle.loads( | |
| 48 value.replace(b"sat.plugins", b"libervia.backend.plugins"), | |
| 49 encoding="utf-8", | |
| 50 ) | |
| 51 if ( | |
|
4228
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
52 PublicKeyMetadata is not None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
53 and table == "private_ind_bin" |
| 4212 | 54 and primary_keys[0] == "XEP-0373" |
| 55 and not primary_keys[1].startswith("/trust") | |
| 56 and isinstance(deserialized, set) | |
| 57 and deserialized | |
| 58 and isinstance(next(iter(deserialized)), PublicKeyMetadata) | |
| 59 ): | |
| 60 # XEP-0373 plugin was pickling an internal class, this can't be converted | |
| 61 # directly to JSON, so we do a special treatment with the add `to_dict` and | |
| 62 # `from_dict` methods. | |
| 63 deserialized = [pkm.to_dict() for pkm in deserialized] | |
| 64 | |
| 4216 | 65 elif ( |
|
4228
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
66 TrustMessageCacheEntry is not None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
67 and table == "private_ind_bin" |
| 4216 | 68 and primary_keys[0] == "XEP-0384/TM" |
| 69 and primary_keys[1] == "cache" | |
| 70 ): | |
| 71 # Same issue and solution as for XEP-0373 | |
| 72 try: | |
| 73 deserialized = [tm.to_dict() for tm in deserialized] | |
| 74 except Exception as e: | |
| 75 print( | |
| 76 "Warning: Failed to convert Trust Management cache with value " | |
| 77 f" {deserialized!r}, using empty array instead: {e}" | |
| 78 ) | |
| 79 deserialized=[] | |
| 80 | |
| 4212 | 81 ret = json.dumps(deserialized, ensure_ascii=False, default=str) |
| 82 if table == 'history' and ret == "{}": | |
| 83 # For history, we can remove empty data, but for other tables it may be | |
| 84 # significant. | |
| 85 ret = None | |
| 86 return ret | |
| 87 except Exception as e: | |
| 88 print( | |
| 89 f"Warning: Failed to convert pickle to JSON, using NULL instead. Error: {e}" | |
| 90 ) | |
| 91 return None | |
| 92 | |
| 93 | |
| 94 def upgrade(): | |
| 95 print( | |
| 96 "This migration may take very long, please be patient and don't stop the process." | |
| 97 ) | |
| 98 connection = op.get_bind() | |
| 99 | |
| 100 tables_and_columns = [ | |
| 101 ("history", "extra", "uid"), | |
| 102 ("private_gen_bin", "value", "namespace", "key"), | |
| 103 ("private_ind_bin", "value", "namespace", "key", "profile_id"), | |
| 104 ] | |
| 105 | |
| 106 for table, column, *primary_keys in tables_and_columns: | |
| 107 primary_key_clause = " AND ".join(f"{pk} = :{pk}" for pk in primary_keys) | |
| 108 select_stmt = sa.text(f"SELECT {', '.join(primary_keys)}, {column} FROM {table}") | |
| 109 update_stmt = sa.text( | |
| 110 f"UPDATE {table} SET {column} = :{column} WHERE {primary_key_clause}" | |
| 111 ) | |
| 112 | |
| 113 result = connection.execute(select_stmt) | |
| 114 for row in result: | |
| 115 value = row[-1] | |
| 116 if value is None: | |
| 117 continue | |
| 118 data = {pk: row[idx] for idx, pk in enumerate(primary_keys)} | |
| 119 data[column] = convert_pickle_to_json(value, table, row[:-1]) | |
| 120 connection.execute(update_stmt.bindparams(**data)) | |
| 121 | |
| 122 | |
| 123 def convert_json_to_pickle(value, table, primary_keys): | |
| 124 """Convert JSON data back to pickled data, handling potential errors.""" | |
| 125 if value is None: | |
| 126 return None | |
| 127 try: | |
| 128 deserialized = json.loads(value) | |
| 129 # Check for the specific table and primary key conditions that require special | |
| 130 # handling | |
| 131 if ( | |
|
4228
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
132 PublicKeyMetadata is not None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
133 and table == "private_ind_bin" |
| 4212 | 134 and primary_keys[0] == "XEP-0373" |
| 135 and not primary_keys[1].startswith("/trust") | |
| 136 ): | |
| 137 # Convert list of dicts back to set of PublicKeyMetadata objects | |
| 138 if isinstance(deserialized, list): | |
| 139 deserialized = {PublicKeyMetadata.from_dict(d) for d in deserialized} | |
| 4216 | 140 elif ( |
|
4228
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
141 TrustMessageCacheEntry is not None |
|
79a4870cfbdf
migration: fix migration when XEP-0373 or XEP-0384 can't be imported:
Goffi <goffi@goffi.org>
parents:
4216
diff
changeset
|
142 and table == "private_ind_bin" |
| 4216 | 143 and primary_keys[0] == "XEP-0384/TM" |
| 144 and primary_keys[1] == "cache" | |
| 145 ): | |
| 146 # Convert list of dicts back to set of TrustMessageCacheEntry objects | |
| 147 if isinstance(deserialized, list): | |
| 148 deserialized = {TrustMessageCacheEntry.from_dict(d) for d in deserialized} | |
| 4212 | 149 return pickle.dumps(deserialized, 0) |
| 150 except Exception as e: | |
| 151 print( | |
| 152 f"Warning: Failed to convert JSON to pickle, using NULL instead. Error: {e}" | |
| 153 ) | |
| 154 return None | |
| 155 | |
| 156 | |
| 157 def downgrade(): | |
| 158 print( | |
| 159 "Reverting JSON columns to LegacyPickle format. This may take a while, please be " | |
| 160 "patient." | |
| 161 ) | |
| 162 connection = op.get_bind() | |
| 163 | |
| 164 tables_and_columns = [ | |
| 165 ("history", "extra", "uid"), | |
| 166 ("private_gen_bin", "value", "namespace", "key"), | |
| 167 ("private_ind_bin", "value", "namespace", "key", "profile_id"), | |
| 168 ] | |
| 169 | |
| 170 for table, column, *primary_keys in tables_and_columns: | |
| 171 primary_key_clause = " AND ".join(f"{pk} = :{pk}" for pk in primary_keys) | |
| 172 select_stmt = sa.text(f"SELECT {', '.join(primary_keys)}, {column} FROM {table}") | |
| 173 update_stmt = sa.text( | |
| 174 f"UPDATE {table} SET {column} = :{column} WHERE {primary_key_clause}" | |
| 175 ) | |
| 176 | |
| 177 result = connection.execute(select_stmt) | |
| 178 for row in result: | |
| 179 value = row[-1] | |
| 180 if value is None: | |
| 181 continue | |
| 182 data = {pk: row[idx] for idx, pk in enumerate(primary_keys)} | |
| 183 data[column] = convert_json_to_pickle(value, table, row[:-1]) | |
| 184 connection.execute(update_stmt.bindparams(**data)) |