Skip to content

Commit

Permalink
Add database creation code
Browse files Browse the repository at this point in the history
  • Loading branch information
mbr committed Apr 8, 2024
1 parent 0616f3d commit 58b804d
Showing 1 changed file with 43 additions and 0 deletions.
43 changes: 43 additions & 0 deletions src/postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ pub(crate) enum Error {
NoRolCreateDb,
#[error("user does not have permission `rolcreaterole`")]
NoRolCreateRole,
#[error("not a valid identifier: {0}")]
InvalidIdent(String),
#[error(transparent)]
PostgresError(#[from] tokio_postgres::Error),
}
Expand Down Expand Up @@ -80,4 +82,45 @@ impl PostgresConnection {

Ok(())
}

pub(crate) async fn create_instance(
&self,
role_name: &str,
role_password: &str,
db_name: &str,
) -> Result<(), Error> {
assert_valid_ident(role_name)?;
assert_valid_ident(db_name)?;

// Create role if it does not exist yet.
if self.client.query_opt("SELECT usename FROM pg_user WHERE usename = $1;", &[&role_name]).await?.is_none() {
// User does not exist, ensure we create it.
let create_role_sql = format!("CREATE ROLE {} NOCREATEDB NOCREATEROLE NOINHERIT LOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT 8 PASSWORD $1;", role_name);
self.client.execute(&create_role_sql, &[&role_name, &role_password]).await?;
};

// Same for database, only create if not existent yet.
if self.client.query_opt("SELECT datname FROM pg_database WHERE datname = $1", &[&db_name]).await?.is_none() {
let create_db_sql = format!("CREATE DATABASE {} WITH OWNER = {};", db_name, role_name);
self.client.execute(&create_db_sql, &[]).await?;
}

Ok(())
}
}

fn valid_ident(input: &str) -> bool {
let Some(first) = input.chars().next() else { return false; }

if !first.is_alphabetic() {
return false;
}

input.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '$')
}

fn assert_valid_ident(input: &str) -> Result<(), Error> {
if !valid_ident(input) {
Err(Error::InvalidIdent(input.to_owned()))
} else {Ok(())}
}

0 comments on commit 58b804d

Please sign in to comment.