dbmigrate_lib/drivers/
postgres.rs

1use std::str::FromStr;
2
3use native_tls::TlsConnector;
4use postgres_client::{Client, Config};
5use postgres_native_tls::MakeTlsConnector;
6
7use super::Driver;
8use errors::{Result, ResultExt};
9
10/// The PostgreSQL driver
11//#[derive(Debug)]
12pub struct Postgres {
13    client: Client,
14}
15
16impl Postgres {
17    /// Create PostgreSQL driver
18    pub fn new(url: &str) -> Result<Postgres> {
19        let config = Config::from_str(url)?;
20        let connector = TlsConnector::new().unwrap();
21        let connector = MakeTlsConnector::new(connector);
22        let client = config.connect(connector)?;
23        Postgres::from_client(client)
24    }
25    /// Create PostgreSQL driver using an existing client
26    pub fn from_client(client: Client) -> Result<Postgres> {
27        let mut pg = Postgres { client };
28        pg.ensure_migration_table_exists();
29        Ok(pg)
30    }
31}
32
33impl Driver for Postgres {
34    fn ensure_migration_table_exists(&mut self) {
35        self.client
36            .simple_query(
37                "
38            CREATE TABLE IF NOT EXISTS __dbmigrate_table(id INTEGER, current INTEGER);
39            INSERT INTO __dbmigrate_table (id, current)
40            SELECT 1, 0
41            WHERE NOT EXISTS(SELECT * FROM __dbmigrate_table WHERE id = 1);
42        ",
43            )
44            .unwrap();
45    }
46
47    fn remove_migration_table(&mut self) {
48        self.client
49            .execute("DROP TABLE __dbmigrate_table;", &[])
50            .unwrap();
51    }
52
53    fn get_current_number(&mut self) -> i32 {
54        let stmt = self
55            .client
56            .prepare(
57                "
58            SELECT current FROM __dbmigrate_table WHERE id = 1;
59        ",
60            )
61            .unwrap();
62        let results = self.client.query(&stmt, &[]).unwrap();
63        results.first().unwrap().get("current")
64    }
65
66    fn set_current_number(&mut self, number: i32) {
67        let stmt = self
68            .client
69            .prepare("UPDATE __dbmigrate_table SET current = $1 WHERE id = 1;")
70            .unwrap();
71        self.client.execute(&stmt, &[&number]).unwrap();
72    }
73
74    fn migrate(&mut self, migration: String, number: i32) -> Result<()> {
75        self.client
76            .simple_query(&migration)
77            .chain_err(|| "Migration failed")?;
78        self.set_current_number(number);
79        Ok(())
80    }
81}