Mercurial > libervia-backend
annotate libervia/backend/memory/migration/versions/fe3a02cb4bec_convert_legacypickle_columns_to_json.py @ 4268:51d004e50786
docker (backend): set `+use_local_shared_tmp` in conf.
author | Goffi <goffi@goffi.org> |
---|---|
date | Thu, 13 Jun 2024 13:22:41 +0200 |
parents | 79a4870cfbdf |
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)) |