Skip to content

Commit 58b804d

Browse files
committed
Add database creation code
1 parent 0616f3d commit 58b804d

File tree

1 file changed

+43
-0
lines changed

1 file changed

+43
-0
lines changed

src/postgres.rs

+43
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ pub(crate) enum Error {
88
NoRolCreateDb,
99
#[error("user does not have permission `rolcreaterole`")]
1010
NoRolCreateRole,
11+
#[error("not a valid identifier: {0}")]
12+
InvalidIdent(String),
1113
#[error(transparent)]
1214
PostgresError(#[from] tokio_postgres::Error),
1315
}
@@ -80,4 +82,45 @@ impl PostgresConnection {
8082

8183
Ok(())
8284
}
85+
86+
pub(crate) async fn create_instance(
87+
&self,
88+
role_name: &str,
89+
role_password: &str,
90+
db_name: &str,
91+
) -> Result<(), Error> {
92+
assert_valid_ident(role_name)?;
93+
assert_valid_ident(db_name)?;
94+
95+
// Create role if it does not exist yet.
96+
if self.client.query_opt("SELECT usename FROM pg_user WHERE usename = $1;", &[&role_name]).await?.is_none() {
97+
// User does not exist, ensure we create it.
98+
let create_role_sql = format!("CREATE ROLE {} NOCREATEDB NOCREATEROLE NOINHERIT LOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT 8 PASSWORD $1;", role_name);
99+
self.client.execute(&create_role_sql, &[&role_name, &role_password]).await?;
100+
};
101+
102+
// Same for database, only create if not existent yet.
103+
if self.client.query_opt("SELECT datname FROM pg_database WHERE datname = $1", &[&db_name]).await?.is_none() {
104+
let create_db_sql = format!("CREATE DATABASE {} WITH OWNER = {};", db_name, role_name);
105+
self.client.execute(&create_db_sql, &[]).await?;
106+
}
107+
108+
Ok(())
109+
}
110+
}
111+
112+
fn valid_ident(input: &str) -> bool {
113+
let Some(first) = input.chars().next() else { return false; }
114+
115+
if !first.is_alphabetic() {
116+
return false;
117+
}
118+
119+
input.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '$')
120+
}
121+
122+
fn assert_valid_ident(input: &str) -> Result<(), Error> {
123+
if !valid_ident(input) {
124+
Err(Error::InvalidIdent(input.to_owned()))
125+
} else {Ok(())}
83126
}

0 commit comments

Comments
 (0)