Sanitize the database name in diesel_cli for CREATE / DROP database

operations.

This allows the usage of database names like foo-bar.

I chose not to use my query_helper functions and instead escape the
names myself in the tests directory to avoid adding a lib.rs to diesel_cli.

Fixes diesel-rs#350

Author: Eijebong

CHANGELOG.md

@@ -30,6 +30,8 @@ for Rust libraries in [RFC #1105](https://github.com/rust-lang/rfcs/blob/master/
 
 * Diesel CLI: Support the `--migration-dir` argument and the `MIGRATION_DIRECTORY` on every migrations related command.
 
+* Diesel CLI: Escape the database name for CREATE/DROP DATABASE operations. This allows you to have hyphens in your database name.
+
 ## [0.12.1] - 2017-05-07
 
 ### Changed

diesel_cli/src/database.rs

@@ -8,6 +8,9 @@ use diesel::sqlite::SqliteConnection;
 use diesel::mysql::MysqlConnection;
 use diesel::types::Bool;
 use diesel::*;
+#[cfg(any(feature="postgres", feature="mysql"))]
+use super::query_helper;
+
 
 use database_error::{DatabaseError, DatabaseResult};
 
@@ -131,7 +134,7 @@ fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
                 let (database, postgres_url) = change_database_of_url(database_url, "postgres");
                 println!("Creating database: {}", database);
                 let conn = try!(PgConnection::establish(&postgres_url));
-                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
+                query_helper::create_database(&database).execute(&conn)?;
             }
         },
         #[cfg(feature="sqlite")]
@@ -147,7 +150,7 @@ fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
                 let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
                 println!("Creating database: {}", database);
                 let conn = try!(MysqlConnection::establish(&mysql_url));
-                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
+                query_helper::create_database(&database).execute(&conn)?;
             }
         },
     }
@@ -179,7 +182,7 @@ fn drop_database(database_url: &str) -> DatabaseResult<()> {
             let conn = try!(PgConnection::establish(&postgres_url));
             if try!(pg_database_exists(&conn, &database)) {
                 println!("Dropping database: {}", database);
-                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
+                query_helper::drop_database(&database).if_exists().execute(&conn)?;
             }
         },
         #[cfg(feature="sqlite")]
@@ -198,7 +201,7 @@ fn drop_database(database_url: &str) -> DatabaseResult<()> {
             let conn = try!(MysqlConnection::establish(&mysql_url));
             if try!(mysql_database_exists(&conn, &database)) {
                 println!("Dropping database: {}", database);
-                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
+                query_helper::drop_database(&database).if_exists().execute(&conn)?;
             }
         },
     }

diesel_cli/src/main.rs

@@ -26,6 +26,8 @@ mod database_error;
 mod database;
 mod cli;
 mod pretty_printing;
+#[cfg(any(feature="postgres", feature="mysql"))]
+mod query_helper;
 
 use chrono::*;
 use clap::{ArgMatches,Shell};

diesel_cli/src/query_helper.rs

@@ -0,0 +1,66 @@
+use diesel::backend::Backend;
+use diesel::query_builder::*;
+use diesel::result::QueryResult;
+
+#[derive(Debug, Clone)]
+pub struct DropDatabaseStatement {
+    db_name: String,
+    if_exists: bool,
+}
+
+impl DropDatabaseStatement {
+    pub fn new(db_name: &str) -> Self {
+        DropDatabaseStatement {
+            db_name: db_name.to_owned(),
+            if_exists: false,
+        }
+    }
+
+    pub fn if_exists(self) -> Self {
+        DropDatabaseStatement { if_exists: true, ..self }
+    }
+}
+
+impl<DB: Backend> QueryFragment<DB> for DropDatabaseStatement {
+    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
+        out.push_sql("DROP DATABASE ");
+        if self.if_exists {
+            out.push_sql("IF EXISTS ");
+        }
+        out.push_identifier(&self.db_name)?;
+        Ok(())
+    }
+}
+
+impl_query_id!(noop: DropDatabaseStatement);
+
+#[derive(Debug, Clone)]
+pub struct CreateDatabaseStatement {
+    db_name: String,
+}
+
+impl CreateDatabaseStatement {
+    pub fn new(db_name: &str) -> Self {
+        CreateDatabaseStatement {
+            db_name: db_name.to_owned(),
+        }
+    }
+}
+
+impl<DB: Backend> QueryFragment<DB> for CreateDatabaseStatement {
+    fn walk_ast(&self, mut out: AstPass<DB>) -> QueryResult<()> {
+        out.push_sql("CREATE DATABASE ");
+        out.push_identifier(&self.db_name)?;
+        Ok(())
+    }
+}
+
+impl_query_id!(noop: CreateDatabaseStatement);
+
+pub fn drop_database(db_name: &str) -> DropDatabaseStatement {
+    DropDatabaseStatement::new(db_name)
+}
+
+pub fn create_database(db_name: &str) -> CreateDatabaseStatement {
+    CreateDatabaseStatement::new(db_name)
+}

diesel_cli/tests/database_reset.rs

@@ -131,3 +131,21 @@ fn reset_works_with_migration_dir_by_env() {
     assert!(db.table_exists("users"));
     assert!(db.table_exists("__diesel_schema_migrations"));
 }
+
+#[test]
+fn reset_sanitize_database_name() {
+    let p = project("name-with-dashes")
+        .folder("migrations")
+        .build();
+    let _db = database(&p.database_url()).create();
+
+    let result = p.command("database")
+        .arg("reset")
+        .run();
+
+    assert!(result.is_success(), "Result was unsuccessful {:?}", result.stdout());
+    assert!(result.stdout().contains("Dropping database:"),
+        "Unexpected stdout {}", result.stdout());
+    assert!(result.stdout().contains("Creating database:"),
+        "Unexpected stdout {}", result.stdout());
+}

diesel_cli/tests/support/mysql_database.rs

@@ -17,7 +17,7 @@ impl Database {
     pub fn create(self) -> Self {
         let (database, mysql_url) = self.split_url();
         let conn = MysqlConnection::establish(&mysql_url).unwrap();
-        conn.execute(&format!("CREATE DATABASE {}", database)).unwrap();
+        conn.execute(&format!("CREATE DATABASE `{}`", database)).unwrap();
         self
     }
 
@@ -56,6 +56,6 @@ impl Drop for Database {
     fn drop(&mut self) {
         let (database, mysql_url) = self.split_url();
         let conn = try_drop!(MysqlConnection::establish(&mysql_url), "Couldn't connect to database");
-        try_drop!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)), "Couldn't drop database");
+        try_drop!(conn.execute(&format!("DROP DATABASE IF EXISTS `{}`", database)), "Couldn't drop database");
     }
 }

diesel_cli/tests/support/postgres_database.rs

@@ -17,7 +17,7 @@ impl Database {
     pub fn create(self) -> Self {
         let (database, postgres_url) = self.split_url();
         let conn = PgConnection::establish(&postgres_url).unwrap();
-        conn.execute(&format!("CREATE DATABASE {}", database)).unwrap();
+        conn.execute(&format!(r#"CREATE DATABASE "{}""#, database)).unwrap();
         self
     }
 
@@ -56,7 +56,7 @@ impl Drop for Database {
         let (database, postgres_url) = self.split_url();
         let conn = try_drop!(PgConnection::establish(&postgres_url), "Couldn't connect to database");
         conn.silence_notices(|| {
-            try_drop!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)), "Couldn't drop database");
+            try_drop!(conn.execute(&format!(r#"DROP DATABASE IF EXISTS "{}""#, database)), "Couldn't drop database");
         });
     }
 }