dbquota/
queries.rs

1use std::{error, fmt};
2
3#[derive(Debug)]
4pub enum Error {
5    MySQL(mysql::Error),
6    RowError(mysql::FromRowError),
7    RowExpected,
8}
9
10impl fmt::Display for Error {
11    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
12        match *self {
13            Error::MySQL(ref err) => err.fmt(f),
14            Error::RowError(ref err) => err.fmt(f),
15            Error::RowExpected => write!(f, "row expected"),
16        }
17    }
18}
19
20impl error::Error for Error {}
21
22impl From<mysql::Error> for Error {
23    fn from(err: mysql::Error) -> Self {
24        Error::MySQL(err)
25    }
26}
27
28impl From<mysql::FromRowError> for Error {
29    fn from(err: mysql::FromRowError) -> Self {
30        Error::RowError(err)
31    }
32}
33
34pub struct Queries<'a> {
35    pool: &'a mysql::Pool,
36}
37
38pub fn new(pool: &mysql::Pool) -> Queries {
39    return Queries { pool: pool };
40}
41
42impl<'a> Queries<'a> {
43    pub fn update_db_size(&self) -> Result<(), Error> {
44        // create table
45        self.pool.prep_exec(
46            r#"CREATE TABLE IF NOT EXISTS quotas (
47        name VARCHAR(64) NOT NULL,
48        bytes BIGINT UNSIGNED NOT NULL,
49        mbytes FLOAT UNSIGNED NOT NULL,
50        quota BIGINT UNSIGNED,
51        enabled TINYINT(1) DEFAULT 0,
52        cdate timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
53        edate timestamp NULL DEFAULT NULL,
54        PRIMARY KEY(name))"#,
55            (),
56        )?;
57
58        // update table with db size
59        // optimize table <tbname> claims disk space
60        let mut tr = self.pool.start_transaction(true, None, None)?;
61        tr.prep_exec(
62            r#"INSERT INTO quotas (name, bytes, mbytes)
63            SELECT t.name, t.bytes, t.mbytes
64            FROM
65            (SELECT
66	        table_schema 'name',
67	        SUM(data_length + index_length) 'bytes',
68	        ROUND(SUM(data_length + index_length) / 1024 / 1024, 2) 'mbytes'
69            FROM information_schema.tables
70            GROUP BY table_schema) as t
71            ON DUPLICATE KEY UPDATE bytes=t.bytes, mbytes=t.mbytes"#,
72            (),
73        )?;
74        tr.commit()?;
75        Ok(())
76    }
77
78    pub fn enforce_quota(&self) -> Result<(), Error> {
79        let rows = self.pool.prep_exec(
80            "SELECT name FROM quotas WHERE bytes > quota AND enabled=1",
81            (),
82        )?;
83        for row in rows {
84            let row = row.map_err(Error::MySQL)?;
85            let dbname = mysql::from_row_opt::<String>(row)?;
86            // revoke insert, update on <dbname>.* FROM '<user>'@'%'
87            let users = self
88                .pool
89                .prep_exec("SELECT user, host FROM mysql.db WHERE Db=?", (&dbname,))?;
90            for user in users {
91                let user = user.map_err(Error::MySQL)?;
92                let (user, host) = mysql::from_row_opt::<(String, String)>(user)?;
93                let mut tr = self.pool.start_transaction(true, None, None)?;
94                tr.prep_exec(
95                    format!(
96                        "REVOKE INSERT, UPDATE, CREATE, ALTER ON `{}`.* FROM '{}'@'{}'",
97                        &dbname, user, host
98                    ),
99                    (),
100                )?;
101                tr.prep_exec(
102                    "UPDATE quotas SET enabled=2, edate=NOW() WHERE name=?",
103                    (&dbname,),
104                )?;
105                tr.commit()?;
106            }
107        }
108        Ok(())
109    }
110
111    pub fn revoke_quota(&self) -> Result<(), Error> {
112        let rows = self.pool.prep_exec(
113            "SELECT name FROM quotas WHERE bytes < quota AND enabled > 1",
114            (),
115        )?;
116        for row in rows {
117            let row = row.map_err(Error::MySQL)?;
118            let dbname = mysql::from_row_opt::<String>(row)?;
119            // revoke insert, update on <dbname>.* FROM '<user>'@'%'
120            let users = self
121                .pool
122                .prep_exec("SELECT user, host FROM mysql.db WHERE Db=?", (&dbname,))?;
123            for user in users {
124                let user = user.map_err(Error::MySQL)?;
125                let (user, host) = mysql::from_row_opt::<(String, String)>(user)?;
126                let mut tr = self.pool.start_transaction(true, None, None)?;
127                tr.prep_exec(
128                    format!(
129                        "GRANT ALL PRIVILEGES ON `{}`.* TO '{}'@'{}'",
130                        &dbname, user, host
131                    ),
132                    (),
133                )?;
134                tr.prep_exec(
135                    "UPDATE quotas SET enabled=1, edate=NOW() WHERE name=?",
136                    (&dbname,),
137                )?;
138                tr.commit()?;
139            }
140        }
141        Ok(())
142    }
143}
dbquota/queries.rs

dbquota/
queries.rs
