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