1use rusqlite::{Connection, Transaction};
7use sql_support::open_database::{self, ConnectionInitializer};
8
9pub const VERSION: u32 = 15;
17
18pub 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
44pub 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 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 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]
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}