Merge pull request #133 from allenap/escape-identifier

Escape database name in `createdb` and `dropdb`

Author: allenap

Cargo.toml

@@ -1,6 +1,11 @@
 [package]
 authors = ["Gavin Panella <gavinpanella@gmail.com>"]
-categories = ["command-line-utilities", "database", "development-tools", "development-tools::testing"]
+categories = [
+    "command-line-utilities",
+    "database",
+    "development-tools",
+    "development-tools::testing",
+]
 description = "Easily create and manage PostgreSQL clusters on demand for testing and development."
 edition = "2021"
 keywords = ["database", "postgres", "postgresql"]
@@ -11,7 +16,7 @@ repository = "https://github.com/allenap/rust-postgresfixture"
 version = "0.3.2"
 
 [badges]
-travis-ci = {repository = "allenap/rust-postgresfixture", branch = "master"}
+travis-ci = { repository = "allenap/rust-postgresfixture", branch = "master" }
 
 [lib]
 name = "postgresfixture"
@@ -23,16 +28,17 @@ name = "postgresfixture"
 path = "src/main.rs"
 
 [dependencies]
-clap = {version = "^3.1.0", features = ["derive", "env"]}
+clap = { version = "^3.1.0", features = ["derive", "env"] }
 color-eyre = "^0.6.1"
-ctrlc = {version = "^3.2.1", features = ["termination"]}
+ctrlc = { version = "^3.2.1", features = ["termination"] }
 either = "^1.6.1"
 nix = "^0.23.0"
 postgres = "^0.19.2"
+postgres-protocol = "^0.6.4"
 rand = "^0.8.5"
 regex = "^1.5.4"
 shell-quote = "^0.3.0"
-uuid = {version = "^0.8.2", features = ["v5"]}
+uuid = { version = "^0.8.2", features = ["v5"] }
 
 [dev-dependencies]
 tempdir = "^0.3.7"

src/cluster.rs

@@ -345,15 +345,21 @@ impl Cluster {
 
     /// Create the named database.
     pub fn createdb(&self, database: &str) -> Result<bool, ClusterError> {
-        let statement = format!("CREATE DATABASE {}", &database);
+        let statement = format!(
+            "CREATE DATABASE {}",
+            postgres_protocol::escape::escape_identifier(database)
+        );
         self.connect("template1")?
             .execute(statement.as_str(), &[])?;
         Ok(true)
     }
 
     /// Drop the named database.
     pub fn dropdb(&self, database: &str) -> Result<bool, ClusterError> {
-        let statement = format!("DROP DATABASE {}", &database);
+        let statement = format!(
+            "DROP DATABASE {}",
+            postgres_protocol::escape::escape_identifier(database)
+        );
         self.connect("template1")?
             .execute(statement.as_str(), &[])?;
         Ok(true)
@@ -594,4 +600,31 @@ mod tests {
             cluster.destroy().unwrap();
         }
     }
+
+    #[test]
+    fn cluster_databases_with_non_plain_names_can_be_created_and_dropped() {
+        // PostgreSQL identifiers containing hyphens, for example, or where we
+        // want to preserve capitalisation, are possible.
+        for runtime in Runtime::find_on_path() {
+            println!("{:?}", runtime);
+            let data_dir = tempdir::TempDir::new("data").unwrap();
+            let cluster = Cluster::new(&data_dir, runtime);
+            cluster.start().unwrap();
+            cluster.createdb("foo-bar").unwrap();
+            cluster.createdb("Foo-BAR").unwrap();
+
+            let expected: HashSet<String> =
+                ["foo-bar", "Foo-BAR", "postgres", "template0", "template1"]
+                    .iter()
+                    .cloned()
+                    .map(|s| s.to_string())
+                    .collect();
+            let observed: HashSet<String> = cluster.databases().unwrap().iter().cloned().collect();
+            assert_eq!(expected, observed);
+
+            cluster.dropdb("foo-bar").unwrap();
+            cluster.dropdb("Foo-BAR").unwrap();
+            cluster.destroy().unwrap();
+        }
+    }
 }