1 // Copyright 2021, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 use anyhow::{anyhow, Context, Result};
16 use rusqlite::{params, OptionalExtension, Transaction};
17
create_or_get_version(tx: &Transaction, current_version: u32) -> Result<u32>18 fn create_or_get_version(tx: &Transaction, current_version: u32) -> Result<u32> {
19 tx.execute(
20 "CREATE TABLE IF NOT EXISTS persistent.version (
21 id INTEGER PRIMARY KEY,
22 version INTEGER);",
23 [],
24 )
25 .context("In create_or_get_version: Failed to create version table.")?;
26
27 let version = tx
28 .query_row("SELECT version FROM persistent.version WHERE id = 0;", [], |row| row.get(0))
29 .optional()
30 .context("In create_or_get_version: Failed to read version.")?;
31
32 let version = if let Some(version) = version {
33 version
34 } else {
35 // If no version table existed it could mean one of two things:
36 // 1) This database is completely new. In this case the version has to be set
37 // to the current version and the current version which also needs to be
38 // returned.
39 // 2) The database predates db versioning. In this case the version needs to be
40 // set to 0, and 0 needs to be returned.
41 let version = if tx
42 .query_row(
43 "SELECT name FROM persistent.sqlite_master
44 WHERE type = 'table' AND name = 'keyentry';",
45 [],
46 |_| Ok(()),
47 )
48 .optional()
49 .context("In create_or_get_version: Failed to check for keyentry table.")?
50 .is_none()
51 {
52 current_version
53 } else {
54 0
55 };
56
57 tx.execute("INSERT INTO persistent.version (id, version) VALUES(0, ?);", params![version])
58 .context("In create_or_get_version: Failed to insert initial version.")?;
59 version
60 };
61 Ok(version)
62 }
63
update_version(tx: &Transaction, new_version: u32) -> Result<()>64 pub(crate) fn update_version(tx: &Transaction, new_version: u32) -> Result<()> {
65 let updated = tx
66 .execute("UPDATE persistent.version SET version = ? WHERE id = 0;", params![new_version])
67 .context("In update_version: Failed to update row.")?;
68 if updated == 1 {
69 Ok(())
70 } else {
71 Err(anyhow!("In update_version: No rows were updated."))
72 }
73 }
74
upgrade_database<F>(tx: &Transaction, current_version: u32, upgraders: &[F]) -> Result<()> where F: Fn(&Transaction) -> Result<u32> + 'static,75 pub fn upgrade_database<F>(tx: &Transaction, current_version: u32, upgraders: &[F]) -> Result<()>
76 where
77 F: Fn(&Transaction) -> Result<u32> + 'static,
78 {
79 if upgraders.len() < current_version as usize {
80 return Err(anyhow!("In upgrade_database: Insufficient upgraders provided."));
81 }
82 let mut db_version = create_or_get_version(tx, current_version)
83 .context("In upgrade_database: Failed to get database version.")?;
84 while db_version < current_version {
85 log::info!("Current DB version={db_version}, perform upgrade");
86 db_version = upgraders[db_version as usize](tx).with_context(|| {
87 format!("In upgrade_database: Trying to upgrade from db version {}.", db_version)
88 })?;
89 log::info!("DB upgrade successful, current DB version now={db_version}");
90 }
91 update_version(tx, db_version).context("In upgrade_database.")
92 }
93
94 #[cfg(test)]
95 mod test {
96 use super::*;
97 use rusqlite::{Connection, TransactionBehavior};
98
99 #[test]
upgrade_database_test()100 fn upgrade_database_test() {
101 let mut conn = Connection::open_in_memory().unwrap();
102 conn.execute("ATTACH DATABASE 'file::memory:' as persistent;", []).unwrap();
103
104 let upgraders: Vec<_> = (0..30_u32)
105 .map(move |i| {
106 move |tx: &Transaction| {
107 tx.execute(
108 "INSERT INTO persistent.test (test_field) VALUES(?);",
109 params![i + 1],
110 )
111 .with_context(|| format!("In upgrade_from_{}_to_{}.", i, i + 1))?;
112 Ok(i + 1)
113 }
114 })
115 .collect();
116
117 for legacy in &[false, true] {
118 if *legacy {
119 conn.execute(
120 "CREATE TABLE IF NOT EXISTS persistent.keyentry (
121 id INTEGER UNIQUE,
122 key_type INTEGER,
123 domain INTEGER,
124 namespace INTEGER,
125 alias BLOB,
126 state INTEGER,
127 km_uuid BLOB);",
128 [],
129 )
130 .unwrap();
131 }
132 for from in 1..29 {
133 for to in from..30 {
134 conn.execute("DROP TABLE IF EXISTS persistent.version;", []).unwrap();
135 conn.execute("DROP TABLE IF EXISTS persistent.test;", []).unwrap();
136 conn.execute(
137 "CREATE TABLE IF NOT EXISTS persistent.test (
138 id INTEGER PRIMARY KEY,
139 test_field INTEGER);",
140 [],
141 )
142 .unwrap();
143
144 {
145 let tx =
146 conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
147 create_or_get_version(&tx, from).unwrap();
148 tx.commit().unwrap();
149 }
150 {
151 let tx =
152 conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
153 upgrade_database(&tx, to, &upgraders).unwrap();
154 tx.commit().unwrap();
155 }
156
157 // In the legacy database case all upgraders starting from 0 have to run. So
158 // after the upgrade step, the expectations need to be adjusted.
159 let from = if *legacy { 0 } else { from };
160
161 // There must be exactly to - from rows.
162 assert_eq!(
163 to - from,
164 conn.query_row(
165 "SELECT COUNT(test_field) FROM persistent.test;",
166 [],
167 |row| row.get(0)
168 )
169 .unwrap()
170 );
171 // Each row must have the correct relation between id and test_field. If this
172 // is not the case, the upgraders were not executed in the correct order.
173 assert_eq!(
174 to - from,
175 conn.query_row(
176 "SELECT COUNT(test_field) FROM persistent.test
177 WHERE id = test_field - ?;",
178 params![from],
179 |row| row.get(0)
180 )
181 .unwrap()
182 );
183 }
184 }
185 }
186 }
187
188 #[test]
create_or_get_version_new_database()189 fn create_or_get_version_new_database() {
190 let mut conn = Connection::open_in_memory().unwrap();
191 conn.execute("ATTACH DATABASE 'file::memory:' as persistent;", []).unwrap();
192 {
193 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
194 let version = create_or_get_version(&tx, 3).unwrap();
195 tx.commit().unwrap();
196 assert_eq!(version, 3);
197 }
198
199 // Was the version table created as expected?
200 assert_eq!(
201 Ok("version".to_owned()),
202 conn.query_row(
203 "SELECT name FROM persistent.sqlite_master
204 WHERE type = 'table' AND name = 'version';",
205 [],
206 |row| row.get(0),
207 )
208 );
209
210 // There is exactly one row in the version table.
211 assert_eq!(
212 Ok(1),
213 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
214 );
215
216 // The version must be set to 3
217 assert_eq!(
218 Ok(3),
219 conn.query_row("SELECT version from persistent.version WHERE id = 0;", [], |row| row
220 .get(0))
221 );
222
223 // Will subsequent calls to create_or_get_version still return the same version even
224 // if the current version changes.
225 {
226 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
227 let version = create_or_get_version(&tx, 5).unwrap();
228 tx.commit().unwrap();
229 assert_eq!(version, 3);
230 }
231
232 // There is still exactly one row in the version table.
233 assert_eq!(
234 Ok(1),
235 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
236 );
237
238 // Bump the version.
239 {
240 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
241 update_version(&tx, 5).unwrap();
242 tx.commit().unwrap();
243 }
244
245 // Now the version should have changed.
246 {
247 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
248 let version = create_or_get_version(&tx, 7).unwrap();
249 tx.commit().unwrap();
250 assert_eq!(version, 5);
251 }
252
253 // There is still exactly one row in the version table.
254 assert_eq!(
255 Ok(1),
256 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
257 );
258
259 // The version must be set to 5
260 assert_eq!(
261 Ok(5),
262 conn.query_row("SELECT version from persistent.version WHERE id = 0;", [], |row| row
263 .get(0))
264 );
265 }
266
267 #[test]
create_or_get_version_legacy_database()268 fn create_or_get_version_legacy_database() {
269 let mut conn = Connection::open_in_memory().unwrap();
270 conn.execute("ATTACH DATABASE 'file::memory:' as persistent;", []).unwrap();
271 // A legacy (version 0) database is detected if the keyentry table exists but no
272 // version table.
273 conn.execute(
274 "CREATE TABLE IF NOT EXISTS persistent.keyentry (
275 id INTEGER UNIQUE,
276 key_type INTEGER,
277 domain INTEGER,
278 namespace INTEGER,
279 alias BLOB,
280 state INTEGER,
281 km_uuid BLOB);",
282 [],
283 )
284 .unwrap();
285
286 {
287 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
288 let version = create_or_get_version(&tx, 3).unwrap();
289 tx.commit().unwrap();
290 // In the legacy case, version 0 must be returned.
291 assert_eq!(version, 0);
292 }
293
294 // Was the version table created as expected?
295 assert_eq!(
296 Ok("version".to_owned()),
297 conn.query_row(
298 "SELECT name FROM persistent.sqlite_master
299 WHERE type = 'table' AND name = 'version';",
300 [],
301 |row| row.get(0),
302 )
303 );
304
305 // There is exactly one row in the version table.
306 assert_eq!(
307 Ok(1),
308 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
309 );
310
311 // The version must be set to 0
312 assert_eq!(
313 Ok(0),
314 conn.query_row("SELECT version from persistent.version WHERE id = 0;", [], |row| row
315 .get(0))
316 );
317
318 // Will subsequent calls to create_or_get_version still return the same version even
319 // if the current version changes.
320 {
321 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
322 let version = create_or_get_version(&tx, 5).unwrap();
323 tx.commit().unwrap();
324 assert_eq!(version, 0);
325 }
326
327 // There is still exactly one row in the version table.
328 assert_eq!(
329 Ok(1),
330 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
331 );
332
333 // Bump the version.
334 {
335 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
336 update_version(&tx, 5).unwrap();
337 tx.commit().unwrap();
338 }
339
340 // Now the version should have changed.
341 {
342 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate).unwrap();
343 let version = create_or_get_version(&tx, 7).unwrap();
344 tx.commit().unwrap();
345 assert_eq!(version, 5);
346 }
347
348 // There is still exactly one row in the version table.
349 assert_eq!(
350 Ok(1),
351 conn.query_row("SELECT COUNT(id) from persistent.version;", [], |row| row.get(0))
352 );
353
354 // The version must be set to 5
355 assert_eq!(
356 Ok(5),
357 conn.query_row("SELECT version from persistent.version WHERE id = 0;", [], |row| row
358 .get(0))
359 );
360 }
361 }
362