relevancy/
schema.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

use rusqlite::{Connection, Transaction};
use sql_support::open_database::{self, ConnectionInitializer};

/// The current database schema version.
///
/// For any changes to the schema [`SQL`], please make sure to:
///
///  1. Bump this version.
///  2. Add a migration from the old version to the new version in
///     [`RelevancyConnectionInitializer::upgrade_from`].
pub const VERSION: u32 = 15;

/// The current database schema.
pub const SQL: &str = "
    CREATE TABLE url_interest(
        url_hash BLOB NOT NULL,
        interest_code INTEGER NOT NULL,
        PRIMARY KEY (url_hash, interest_code)
    ) WITHOUT ROWID;

    -- Stores user interest vectors.  The `kind` field stores the raw code from the `InterestVectorKind` enum.
    CREATE TABLE user_interest(
        kind INTEGER NOT NULL,
        interest_code INTEGER NOT NULL,
        count INTEGER NOT NULL,
        PRIMARY KEY (kind, interest_code)
    ) WITHOUT ROWID;
    CREATE TABLE multi_armed_bandit(
        bandit TEXT NOT NULL,
        arm TEXT NOT NULL,
        alpha INTEGER NOT NULL,
        beta INTEGER NOT NULL,
        clicks INTEGER NOT NULL,
        impressions INTEGER NOT NULL,
        PRIMARY KEY (bandit, arm)
    ) WITHOUT ROWID;
";

/// Initializes an SQLite connection to the Relevancy database, performing
/// migrations as needed.
pub struct RelevancyConnectionInitializer;

impl ConnectionInitializer for RelevancyConnectionInitializer {
    const NAME: &'static str = "relevancy db";
    const END_VERSION: u32 = VERSION;

    fn prepare(&self, conn: &Connection, _db_empty: bool) -> open_database::Result<()> {
        let initial_pragmas = "
            -- Use in-memory storage for TEMP tables.
            PRAGMA temp_store = 2;
            PRAGMA journal_mode = WAL;
            PRAGMA foreign_keys = ON;
        ";
        conn.execute_batch(initial_pragmas)?;
        Ok(())
    }

    fn init(&self, db: &Transaction<'_>) -> open_database::Result<()> {
        Ok(db.execute_batch(SQL)?)
    }

    fn upgrade_from(&self, tx: &Transaction<'_>, version: u32) -> open_database::Result<()> {
        match version {
            // Upgrades 1-12 are missing because we started with version 13, because of a
            // copy-and-paste error.
            13 => {
                tx.execute(
                    "
    CREATE TABLE user_interest(
        kind INTEGER NOT NULL,
        interest_code INTEGER NOT NULL,
        count INTEGER NOT NULL,
        PRIMARY KEY (kind, interest_code)
    ) WITHOUT ROWID;
                ",
                    (),
                )?;
                Ok(())
            }
            14 => {
                tx.execute(
                    "CREATE TABLE multi_armed_bandit(
        bandit TEXT NOT NULL,
        arm TEXT NOT NULL,
        alpha INTEGER NOT NULL,
        beta INTEGER NOT NULL,
        clicks INTEGER NOT NULL,
        impressions INTEGER NOT NULL,
        PRIMARY KEY (bandit, arm)
    ) WITHOUT ROWID;",
                    (),
                )?;
                Ok(())
            }
            _ => Err(open_database::Error::IncompatibleVersion(version)),
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use sql_support::open_database::test_utils::MigratedDatabaseFile;

    /// The first database schema we used
    pub const V1_SCHEMA: &str = "
    CREATE TABLE url_interest(
        url_hash BLOB NOT NULL,
        interest_code INTEGER NOT NULL,
        PRIMARY KEY (url_hash, interest_code)
    ) WITHOUT ROWID;

        PRAGMA user_version=13;
    ";

    /// Test running all schema upgrades
    ///
    /// If an upgrade fails, then this test will fail with a panic.
    #[test]
    fn test_all_upgrades() {
        let db_file = MigratedDatabaseFile::new(RelevancyConnectionInitializer, V1_SCHEMA);
        db_file.run_all_upgrades();
        db_file.assert_schema_matches_new_database();
    }
}