Mercurial > libervia-backend
annotate libervia/backend/memory/migration/versions/fe3a02cb4bec_convert_legacypickle_columns_to_json.py @ 4304:92a886f31581 default tip @
doc (components): new Email gateway documentation:
fix 449
author | Goffi <goffi@goffi.org> |
---|---|
date | Fri, 06 Sep 2024 18:07:44 +0200 |
parents | 0d7bb4df2343 |
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)) |