relevancy/
schema.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
4 */
5
6use rusqlite::{Connection, Transaction};
7use sql_support::open_database::{self, ConnectionInitializer};
8
9/// The current database schema version.
10///
11/// For any changes to the schema [`SQL`], please make sure to:
12///
13///  1. Bump this version.
14///  2. Add a migration from the old version to the new version in
15///     [`RelevancyConnectionInitializer::upgrade_from`].
16pub const VERSION: u32 = 15;
17
18/// The current database schema.
19pub const SQL: &str = "
20    CREATE TABLE url_interest(
21        url_hash BLOB NOT NULL,
22        interest_code INTEGER NOT NULL,
23        PRIMARY KEY (url_hash, interest_code)
24    ) WITHOUT ROWID;
25
26    -- Stores user interest vectors.  The `kind` field stores the raw code from the `InterestVectorKind` enum.
27    CREATE TABLE user_interest(
28        kind INTEGER NOT NULL,
29        interest_code INTEGER NOT NULL,
30        count INTEGER NOT NULL,
31        PRIMARY KEY (kind, interest_code)
32    ) WITHOUT ROWID;
33    CREATE TABLE multi_armed_bandit(
34        bandit TEXT NOT NULL,
35        arm TEXT NOT NULL,
36        alpha INTEGER NOT NULL,
37        beta INTEGER NOT NULL,
38        clicks INTEGER NOT NULL,
39        impressions INTEGER NOT NULL,
40        PRIMARY KEY (bandit, arm)
41    ) WITHOUT ROWID;
42";
43
44/// Initializes an SQLite connection to the Relevancy database, performing
45/// migrations as needed.
46pub struct RelevancyConnectionInitializer;
47
48impl ConnectionInitializer for RelevancyConnectionInitializer {
49    const NAME: &'static str = "relevancy db";
50    const END_VERSION: u32 = VERSION;
51
52    fn prepare(&self, conn: &Connection, _db_empty: bool) -> open_database::Result<()> {
53        let initial_pragmas = "
54            -- Use in-memory storage for TEMP tables.
55            PRAGMA temp_store = 2;
56            PRAGMA journal_mode = WAL;
57            PRAGMA foreign_keys = ON;
58        ";
59        conn.execute_batch(initial_pragmas)?;
60        Ok(())
61    }
62
63    fn init(&self, db: &Transaction<'_>) -> open_database::Result<()> {
64        Ok(db.execute_batch(SQL)?)
65    }
66
67    fn upgrade_from(&self, tx: &Transaction<'_>, version: u32) -> open_database::Result<()> {
68        match version {
69            // Upgrades 1-12 are missing because we started with version 13, because of a
70            // copy-and-paste error.
71            13 => {
72                tx.execute(
73                    "
74    CREATE TABLE user_interest(
75        kind INTEGER NOT NULL,
76        interest_code INTEGER NOT NULL,
77        count INTEGER NOT NULL,
78        PRIMARY KEY (kind, interest_code)
79    ) WITHOUT ROWID;
80                ",
81                    (),
82                )?;
83                Ok(())
84            }
85            14 => {
86                tx.execute(
87                    "CREATE TABLE multi_armed_bandit(
88        bandit TEXT NOT NULL,
89        arm TEXT NOT NULL,
90        alpha INTEGER NOT NULL,
91        beta INTEGER NOT NULL,
92        clicks INTEGER NOT NULL,
93        impressions INTEGER NOT NULL,
94        PRIMARY KEY (bandit, arm)
95    ) WITHOUT ROWID;",
96                    (),
97                )?;
98                Ok(())
99            }
100            _ => Err(open_database::Error::IncompatibleVersion(version)),
101        }
102    }
103}
104
105#[cfg(test)]
106mod test {
107    use super::*;
108    use sql_support::open_database::test_utils::MigratedDatabaseFile;
109
110    /// The first database schema we used
111    pub const V1_SCHEMA: &str = "
112    CREATE TABLE url_interest(
113        url_hash BLOB NOT NULL,
114        interest_code INTEGER NOT NULL,
115        PRIMARY KEY (url_hash, interest_code)
116    ) WITHOUT ROWID;
117
118        PRAGMA user_version=13;
119    ";
120
121    /// Test running all schema upgrades
122    ///
123    /// If an upgrade fails, then this test will fail with a panic.
124    #[test]
125    fn test_all_upgrades() {
126        let db_file = MigratedDatabaseFile::new(RelevancyConnectionInitializer, V1_SCHEMA);
127        db_file.run_all_upgrades();
128        db_file.assert_schema_matches_new_database();
129    }
130}