misc: initial commit
This commit is contained in:
hdbg
8
src/banner.txt
Normal file
8
src/banner.txt
Normal file
@@ -0,0 +1,8 @@
|
||||
|
||||
|
||||
_ __ __ ___ __
|
||||
| '_ \ / _` \ \/ /
|
||||
| |_) | (_| |> <
|
||||
| .__/ \__, /_/\_\
|
||||
| | __/ |
|
||||
|_| |___/
|
||||
64
src/cli.rs
Normal file
64
src/cli.rs
Normal file
@@ -0,0 +1,64 @@
|
||||
use clap::{Args, Parser, Subcommand, builder::styling};
|
||||
|
||||
const STYLES: styling::Styles = styling::Styles::styled()
|
||||
.header(styling::AnsiColor::Green.on_default().bold())
|
||||
.usage(styling::AnsiColor::Green.on_default().bold())
|
||||
.literal(styling::AnsiColor::Blue.on_default().bold())
|
||||
.placeholder(styling::AnsiColor::Cyan.on_default());
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "pgx")]
|
||||
#[command(about = "Project-scoped PostgreSQL instance manager", long_about = None)]
|
||||
#[command(version)]
|
||||
#[command(styles = STYLES)]
|
||||
pub struct Cli {
|
||||
#[command(subcommand)]
|
||||
pub command: Commands,
|
||||
|
||||
#[arg(short, long, global = true)]
|
||||
pub verbose: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, clap::ValueEnum)]
|
||||
pub enum ConnectionFormat {
|
||||
/// Human-readable text format
|
||||
Text,
|
||||
/// JSON format
|
||||
Json,
|
||||
/// Environment variable format
|
||||
Env,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum ControlCommands {
|
||||
/// Start postgres instance
|
||||
Start,
|
||||
/// Stop postgres instance
|
||||
Stop,
|
||||
/// Restart postgres instance
|
||||
Restart,
|
||||
/// (WARNING!) Destroy postgres instance
|
||||
Destroy,
|
||||
/// Status of instance
|
||||
Status,
|
||||
/// View logs produced by postgres
|
||||
Logs { follow: bool },
|
||||
/// (Sensitive) get connection details
|
||||
Connection { format: ConnectionFormat },
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum Commands {
|
||||
/// Create a new project, or initialize instance for existing one
|
||||
Init,
|
||||
/// Create a new project, or initialize instance for existing one
|
||||
Sync,
|
||||
|
||||
/// Start the PostgreSQL container for the current project
|
||||
Instance {
|
||||
// Name of the instance you want to control. Defaults to current project
|
||||
name: Option<String>,
|
||||
#[command(subcommand)]
|
||||
cmd: ControlCommands,
|
||||
},
|
||||
}
|
||||
174
src/config.rs
Normal file
174
src/config.rs
Normal file
@@ -0,0 +1,174 @@
|
||||
use miette::miette;
|
||||
use miette::{Context, IntoDiagnostic, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_with::{DisplayFromStr, serde_as};
|
||||
use std::fmt::Display;
|
||||
use std::{
|
||||
path::{Path, PathBuf},
|
||||
str::FromStr,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub struct PostgresVersion {
|
||||
pub major: u32,
|
||||
pub minor: u32,
|
||||
}
|
||||
|
||||
impl FromStr for PostgresVersion {
|
||||
type Err = miette::Report;
|
||||
|
||||
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
|
||||
let Some((major_str, minor_str)) = s.split_once(".") else {
|
||||
return Err(miette!(
|
||||
help = "update hardcoded version",
|
||||
"expected two fragments in version"
|
||||
));
|
||||
};
|
||||
let major = major_str.parse().into_diagnostic()?;
|
||||
let minor = minor_str.parse().into_diagnostic()?;
|
||||
Ok(Self { major, minor })
|
||||
}
|
||||
}
|
||||
impl Display for PostgresVersion {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}.{}", self.major, self.minor)
|
||||
}
|
||||
}
|
||||
|
||||
const PROJECT_FILENAME: &'static str = "pgx.toml";
|
||||
|
||||
/// Configuration stored in pgx.toml
|
||||
#[serde_as]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PgxConfig {
|
||||
/// PostgreSQL version to use
|
||||
#[serde_as(as = "DisplayFromStr")]
|
||||
pub version: PostgresVersion,
|
||||
|
||||
/// Database password
|
||||
pub password: String,
|
||||
|
||||
/// Port to bind on host
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
impl PgxConfig {
|
||||
pub fn load(path: impl AsRef<Path>) -> Result<Self> {
|
||||
let path = path.as_ref();
|
||||
let content = std::fs::read_to_string(path)
|
||||
.into_diagnostic()
|
||||
.wrap_err_with(|| format!("Failed to read config file: {}", path.display()))?;
|
||||
|
||||
let config: PgxConfig = toml::from_str(&content)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to parse pgx.toml")?;
|
||||
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
pub fn save(&self, path: impl AsRef<Path>) -> Result<()> {
|
||||
let path = path.as_ref();
|
||||
let content = toml::to_string_pretty(self)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to serialize config")?;
|
||||
|
||||
std::fs::write(path, content)
|
||||
.into_diagnostic()
|
||||
.wrap_err_with(|| format!("Failed to write config file: {}", path.display()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Project {
|
||||
/// Project name (derived from directory name)
|
||||
pub name: String,
|
||||
|
||||
/// Path to the project directory containing pgx.toml
|
||||
pub path: PathBuf,
|
||||
|
||||
pub config: PgxConfig,
|
||||
}
|
||||
|
||||
impl Project {
|
||||
pub fn container_name(&self) -> String {
|
||||
let container_name = format!(
|
||||
"pgx-{}-{}",
|
||||
self.name,
|
||||
self.config.version.to_string().replace('.', "_")
|
||||
);
|
||||
container_name
|
||||
}
|
||||
|
||||
/// Load a project from the current directory
|
||||
pub fn load() -> Result<Option<Self>> {
|
||||
let project_path = get_project_path()?;
|
||||
let config_path = project_path.join(PROJECT_FILENAME);
|
||||
|
||||
if !config_path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let config = PgxConfig::load(&config_path)?;
|
||||
let name = Self::extract_project_name(&project_path)?;
|
||||
|
||||
Ok(Some(Project {
|
||||
name,
|
||||
path: project_path,
|
||||
config,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn new(config: PgxConfig) -> Result<Self> {
|
||||
let project_path = get_project_path()?;
|
||||
let name = Self::extract_project_name(&project_path)?;
|
||||
|
||||
let this = Self {
|
||||
name,
|
||||
path: project_path,
|
||||
config,
|
||||
};
|
||||
|
||||
this.save_config()?;
|
||||
|
||||
Ok(this)
|
||||
}
|
||||
|
||||
/// Extract project name from directory path
|
||||
fn extract_project_name(path: &Path) -> Result<String> {
|
||||
path.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| miette::miette!("Failed to extract project name from path"))
|
||||
}
|
||||
|
||||
/// Get the path to the pgx.toml file
|
||||
pub fn config_path(&self) -> PathBuf {
|
||||
self.path.join("pgx.toml")
|
||||
}
|
||||
|
||||
/// Save the current configuration
|
||||
pub fn save_config(&self) -> Result<()> {
|
||||
self.config.save(self.config_path())
|
||||
}
|
||||
}
|
||||
|
||||
fn get_project_path() -> Result<PathBuf, miette::Error> {
|
||||
let project_path = std::env::current_dir()
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to get current directory")?;
|
||||
Ok(project_path)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extract_project_name() {
|
||||
let path = PathBuf::from("/home/user/my-project");
|
||||
let name = Project::extract_project_name(&path).unwrap();
|
||||
assert_eq!(name, "my-project");
|
||||
}
|
||||
}
|
||||
632
src/controller.rs
Normal file
632
src/controller.rs
Normal file
@@ -0,0 +1,632 @@
|
||||
use miette::{bail, miette};
|
||||
use rand::{Rng, distr::Alphanumeric};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::{collections::HashMap, fmt::Write, pin::pin, str::FromStr};
|
||||
|
||||
use bollard::{
|
||||
Docker,
|
||||
errors::Error,
|
||||
query_parameters::{
|
||||
CreateContainerOptions, CreateImageOptions, InspectContainerOptions, ListImagesOptions,
|
||||
ListImagesOptionsBuilder, SearchImagesOptions, StartContainerOptions, StopContainerOptions,
|
||||
},
|
||||
secret::{ContainerConfig, ContainerCreateBody, CreateImageInfo},
|
||||
};
|
||||
use futures::{Stream, StreamExt, TryStreamExt};
|
||||
use indicatif::{MultiProgress, ProgressBar, ProgressState, ProgressStyle};
|
||||
use miette::{Context, IntoDiagnostic, Result, diagnostic};
|
||||
use tracing::info;
|
||||
|
||||
use crate::{
|
||||
config::{PgxConfig, PostgresVersion, Project},
|
||||
state::{InstanceState, StateManager},
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ContainerStatus {
|
||||
Running,
|
||||
Stopped,
|
||||
Paused,
|
||||
Restarting,
|
||||
Dead,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
const DOCKERHUB_POSTGRES: &str = "postgres";
|
||||
const DEFAULT_POSTGRES_PORT: u16 = 5432;
|
||||
const PORT_SEARCH_RANGE: u16 = 100;
|
||||
|
||||
fn format_image(ver: &PostgresVersion) -> String {
|
||||
format!("{DOCKERHUB_POSTGRES}:{}", ver.to_string())
|
||||
}
|
||||
|
||||
fn find_available_port() -> Result<u16> {
|
||||
use std::net::TcpListener;
|
||||
|
||||
for port in DEFAULT_POSTGRES_PORT..(DEFAULT_POSTGRES_PORT + PORT_SEARCH_RANGE) {
|
||||
if TcpListener::bind(("127.0.0.1", port)).is_ok() {
|
||||
return Ok(port);
|
||||
}
|
||||
}
|
||||
|
||||
miette::bail!(
|
||||
"No available ports found in range {}-{}",
|
||||
DEFAULT_POSTGRES_PORT,
|
||||
DEFAULT_POSTGRES_PORT + PORT_SEARCH_RANGE - 1
|
||||
)
|
||||
}
|
||||
|
||||
fn new_download_pb(multi: &MultiProgress, layer_id: &str) -> ProgressBar {
|
||||
let pb = multi.add(ProgressBar::new(0));
|
||||
pb.set_style(
|
||||
ProgressStyle::with_template(&format!(
|
||||
"{{spinner:.green}} [{{elapsed_precise}}] {{msg}} [{{wide_bar:.cyan/blue}}] {{bytes}}/{{total_bytes}} ({{eta}})"
|
||||
))
|
||||
.unwrap()
|
||||
.with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
|
||||
write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
|
||||
})
|
||||
.progress_chars("#>-"),
|
||||
);
|
||||
pb.set_message(format!("Layer {}", layer_id));
|
||||
pb
|
||||
}
|
||||
|
||||
// sadly type ... = impl ... is unstable
|
||||
pub async fn perform_download(
|
||||
multi: MultiProgress,
|
||||
chunks: impl Stream<Item = Result<CreateImageInfo, Error>>,
|
||||
) -> Result<()> {
|
||||
let mut chunks = pin!(chunks);
|
||||
let mut layer_progress: HashMap<String, ProgressBar> = HashMap::new();
|
||||
|
||||
while let Some(download_info) = chunks.try_next().await.into_diagnostic()? {
|
||||
download_check_for_error(&mut layer_progress, &download_info)?;
|
||||
|
||||
let layer_id = download_info.id.as_deref().unwrap_or("unknown");
|
||||
|
||||
// Get or create progress bar for this layer
|
||||
let pb = layer_progress
|
||||
.entry(layer_id.to_string())
|
||||
.or_insert_with(|| new_download_pb(&multi, layer_id));
|
||||
|
||||
download_drive_progress(pb, download_info);
|
||||
}
|
||||
|
||||
// Clean up any remaining progress bars
|
||||
for (_, pb) in layer_progress.drain() {
|
||||
pb.finish_and_clear();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn download_drive_progress(pb: &mut ProgressBar, download_info: CreateImageInfo) {
|
||||
match download_info.progress_detail {
|
||||
Some(info) => match (info.current, info.total) {
|
||||
(None, None) => {
|
||||
pb.inc(1);
|
||||
}
|
||||
(current, total) => {
|
||||
if let Some(total) = total {
|
||||
pb.set_length(total as u64);
|
||||
}
|
||||
if let Some(current) = current {
|
||||
pb.set_position(current as u64);
|
||||
}
|
||||
|
||||
if let (Some(current), Some(total)) = (current, total)
|
||||
&& (current == total)
|
||||
{
|
||||
pb.finish_with_message("Completed!");
|
||||
}
|
||||
}
|
||||
},
|
||||
None => {
|
||||
// No progress detail, just show activity
|
||||
pb.tick();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn download_check_for_error(
|
||||
layer_progress: &mut HashMap<String, ProgressBar>,
|
||||
download_info: &CreateImageInfo,
|
||||
) -> Result<()> {
|
||||
if let Some(error_detail) = &download_info.error_detail {
|
||||
for (_, pb) in layer_progress.drain() {
|
||||
pb.finish_and_clear();
|
||||
}
|
||||
|
||||
match (error_detail.code, &error_detail.message) {
|
||||
(None, Some(msg)) => miette::bail!("docker image download error: {}", msg),
|
||||
(Some(code), None) => miette::bail!("docker image download error: code {}", code),
|
||||
(Some(code), Some(msg)) => {
|
||||
miette::bail!(
|
||||
"docker image download error: code {}, message: {}",
|
||||
code,
|
||||
msg
|
||||
)
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub struct DockerController {
|
||||
daemon: Docker,
|
||||
}
|
||||
|
||||
impl DockerController {
|
||||
pub async fn new() -> Result<Self> {
|
||||
let docker = Docker::connect_with_local_defaults()
|
||||
.into_diagnostic()
|
||||
.wrap_err(
|
||||
"Failed to connect to Docker! pgx required Docker installed. Make sure it's running.",
|
||||
)?;
|
||||
|
||||
info!("docker.created");
|
||||
|
||||
docker
|
||||
.list_images(Some(ListImagesOptions::default()))
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Docker basic connectivity test refused")?;
|
||||
|
||||
Ok(Self { daemon: docker })
|
||||
}
|
||||
|
||||
pub async fn download_image(&self, image: String) -> Result<()> {
|
||||
let options = Some(CreateImageOptions {
|
||||
from_image: Some(image.clone()),
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
let download_progress = self.daemon.create_image(options, None, None);
|
||||
|
||||
let multi = MultiProgress::new();
|
||||
|
||||
println!("Downloading {image}");
|
||||
|
||||
perform_download(multi, download_progress).await?;
|
||||
|
||||
println!("Download complete!");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn ensure_version_downloaded(&self, ver: &PostgresVersion) -> Result<()> {
|
||||
let desired_image_tag = format_image(ver);
|
||||
|
||||
let images = self
|
||||
.daemon
|
||||
.list_images(Some(ListImagesOptions::default()))
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("failed to list installed docker images")?;
|
||||
|
||||
let is_downloaded = images
|
||||
.iter()
|
||||
.any(|img| img.repo_tags.contains(&desired_image_tag));
|
||||
|
||||
if !is_downloaded {
|
||||
self.download_image(desired_image_tag).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// TODO: make client to get available versions from dockerhub
|
||||
pub async fn available_versions(&self) -> Result<Vec<PostgresVersion>> {
|
||||
Ok(vec!["18.1", "17.7", "16.11", "15.15", "14.20"]
|
||||
.into_iter()
|
||||
.map(|v| PostgresVersion::from_str(v).unwrap())
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub async fn container_exists(&self, container_id: &str) -> Result<bool> {
|
||||
match self
|
||||
.daemon
|
||||
.inspect_container(container_id, None::<InspectContainerOptions>)
|
||||
.await
|
||||
{
|
||||
Ok(_) => Ok(true),
|
||||
Err(bollard::errors::Error::DockerResponseServerError {
|
||||
status_code: 404, ..
|
||||
}) => Ok(false),
|
||||
Err(e) => Err(e)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to inspect container"),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn is_container_running(&self, container_name: &str) -> Result<bool> {
|
||||
let container = self
|
||||
.daemon
|
||||
.inspect_container(container_name, None::<InspectContainerOptions>)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to inspect container")?;
|
||||
|
||||
Ok(container.state.and_then(|s| s.running).unwrap_or(false))
|
||||
}
|
||||
|
||||
pub async fn create_postgres_container(
|
||||
&self,
|
||||
container_name: &str,
|
||||
version: &PostgresVersion,
|
||||
password: &str,
|
||||
port: u16,
|
||||
) -> Result<String> {
|
||||
use bollard::models::{HostConfig, PortBinding};
|
||||
use std::collections::HashMap;
|
||||
|
||||
let image = format_image(version);
|
||||
|
||||
let env = vec![
|
||||
format!("POSTGRES_PASSWORD={}", password),
|
||||
format!("POSTGRES_USER={}", USERNAME),
|
||||
format!("POSTGRES_DB={}", DATABASE),
|
||||
];
|
||||
|
||||
let mut port_bindings = HashMap::new();
|
||||
port_bindings.insert(
|
||||
"5432/tcp".to_string(),
|
||||
Some(vec![PortBinding {
|
||||
host_ip: Some("127.0.0.1".to_string()),
|
||||
host_port: Some(port.to_string()),
|
||||
}]),
|
||||
);
|
||||
|
||||
let host_config = HostConfig {
|
||||
port_bindings: Some(port_bindings),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut labels = HashMap::new();
|
||||
labels.insert("pgx.postgres.version".to_string(), version.to_string());
|
||||
|
||||
let config = ContainerCreateBody {
|
||||
image: Some(image),
|
||||
env: Some(env),
|
||||
host_config: Some(host_config),
|
||||
labels: Some(labels),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let options = CreateContainerOptions {
|
||||
name: Some(container_name.to_owned()),
|
||||
platform: String::new(),
|
||||
};
|
||||
|
||||
let response = self
|
||||
.daemon
|
||||
.create_container(Some(options), config)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to create container")?;
|
||||
|
||||
Ok(response.id)
|
||||
}
|
||||
|
||||
pub async fn start_container(&self, container_id: &str) -> Result<()> {
|
||||
self.daemon
|
||||
.start_container(container_id, None::<StartContainerOptions>)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to start container")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn container_exists_by_id(&self, container_id: &str) -> Result<bool> {
|
||||
match self
|
||||
.daemon
|
||||
.inspect_container(container_id, None::<InspectContainerOptions>)
|
||||
.await
|
||||
{
|
||||
Ok(_) => Ok(true),
|
||||
Err(bollard::errors::Error::DockerResponseServerError {
|
||||
status_code: 404, ..
|
||||
}) => Ok(false),
|
||||
Err(e) => Err(e)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to inspect container by ID"),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn is_container_running_by_id(&self, container_id: &str) -> Result<bool> {
|
||||
let container = self
|
||||
.daemon
|
||||
.inspect_container(container_id, None::<InspectContainerOptions>)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to inspect container")?;
|
||||
|
||||
Ok(container.state.and_then(|s| s.running).unwrap_or(false))
|
||||
}
|
||||
|
||||
pub async fn start_container_by_id(&self, container_id: &str) -> Result<()> {
|
||||
self.start_container(container_id).await
|
||||
}
|
||||
|
||||
pub async fn stop_container(&self, container_id: &str, timeout: i32) -> Result<()> {
|
||||
self.daemon
|
||||
.stop_container(
|
||||
container_id,
|
||||
Some(StopContainerOptions {
|
||||
t: Some(timeout),
|
||||
signal: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to stop container")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_container_postgres_version(
|
||||
&self,
|
||||
container_id: &str,
|
||||
) -> Result<PostgresVersion> {
|
||||
let container = self
|
||||
.daemon
|
||||
.inspect_container(container_id, None::<InspectContainerOptions>)
|
||||
.await
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to inspect container")?;
|
||||
|
||||
let labels = container
|
||||
.config
|
||||
.and_then(|c| c.labels)
|
||||
.ok_or_else(|| miette!("Container has no labels"))?;
|
||||
|
||||
let version_str = labels
|
||||
.get("pgx.postgres.version")
|
||||
.ok_or_else(|| miette!("Container missing pgx.postgres.version label"))?;
|
||||
|
||||
PostgresVersion::from_str(version_str)
|
||||
.map_err(|_| miette!("Invalid version in label: {}", version_str))
|
||||
}
|
||||
}
|
||||
|
||||
const USERNAME: &str = "postgres";
|
||||
const DATABASE: &str = "postgres";
|
||||
|
||||
const PASSWORD_LENGTH: usize = 16;
|
||||
pub fn generate_password() -> String {
|
||||
let password = (&mut rand::rng())
|
||||
.sample_iter(Alphanumeric)
|
||||
.take(PASSWORD_LENGTH)
|
||||
.map(|b| b as char)
|
||||
.collect();
|
||||
password
|
||||
}
|
||||
|
||||
const MAX_RETRIES: u32 = 10;
|
||||
const VERIFY_DURATION_SECS: u64 = 5;
|
||||
|
||||
pub struct Controller {
|
||||
pub docker: DockerController,
|
||||
project: Option<Project>,
|
||||
state: StateManager,
|
||||
}
|
||||
|
||||
impl Controller {
|
||||
pub async fn new() -> Result<Self> {
|
||||
Ok(Self {
|
||||
docker: DockerController::new().await?,
|
||||
project: Project::load()?,
|
||||
state: StateManager::load()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn init_project(&self) -> Result<()> {
|
||||
if let Some(project) = &self.project {
|
||||
return self.reconcile(project).await;
|
||||
}
|
||||
|
||||
println!("Initializing new pgx project...");
|
||||
|
||||
let mut versions = self.docker.available_versions().await?;
|
||||
versions.sort();
|
||||
let latest_version = versions
|
||||
.last()
|
||||
.ok_or(miette!("expected to have at least one version"))?;
|
||||
|
||||
let config = PgxConfig {
|
||||
version: *latest_version,
|
||||
password: generate_password(),
|
||||
port: find_available_port()?,
|
||||
};
|
||||
let project = Project::new(config)?;
|
||||
|
||||
println!("Created pgx.toml in {}", project.path.display());
|
||||
println!(" Project: {}", project.name);
|
||||
println!(" PostgreSQL version: {}", project.config.version);
|
||||
println!(" Port: {}", project.config.port);
|
||||
println!(" Password: {}", "*".repeat(project.config.password.len()));
|
||||
|
||||
self.reconcile(&project).await?;
|
||||
|
||||
println!("\nProject initialized successfully!");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn reconcile(&self, project: &Project) -> Result<()> {
|
||||
self.docker
|
||||
.ensure_version_downloaded(&project.config.version)
|
||||
.await?;
|
||||
|
||||
self.ensure_container_running(project).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ensure_container_running(&self, project: &Project) -> Result<()> {
|
||||
let mut state = StateManager::load()?;
|
||||
let instance_state = state.get_mut(&project.name);
|
||||
|
||||
let container_id = match instance_state {
|
||||
Some(instance) => match self.ensure_container_exists(&instance).await? {
|
||||
Some(id) => id,
|
||||
None => self.update_project_container(project, &mut state).await?,
|
||||
},
|
||||
None => self.update_project_container(project, &mut state).await?,
|
||||
};
|
||||
|
||||
let container_version = self
|
||||
.docker
|
||||
.get_container_postgres_version(&container_id)
|
||||
.await?;
|
||||
|
||||
self.ensure_matches_project_version(project, &mut state, &container_id, container_version)
|
||||
.await?;
|
||||
|
||||
if self
|
||||
.docker
|
||||
.is_container_running_by_id(&container_id)
|
||||
.await?
|
||||
{
|
||||
println!("Container is already running");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("Starting container...");
|
||||
|
||||
for attempt in 1..=MAX_RETRIES {
|
||||
let result = self.try_starting_container(&container_id, attempt).await;
|
||||
|
||||
match result {
|
||||
Ok(_) => break,
|
||||
Err(err) => println!("Error: {:#?}", err),
|
||||
}
|
||||
|
||||
if attempt < MAX_RETRIES {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
|
||||
miette::bail!("Failed to start container after {} attempts", MAX_RETRIES)
|
||||
}
|
||||
|
||||
async fn try_starting_container(
|
||||
&self,
|
||||
container_id: &String,
|
||||
attempt: u32,
|
||||
) -> Result<(), miette::Error> {
|
||||
Ok(
|
||||
match self.docker.start_container_by_id(container_id).await {
|
||||
Ok(_) => {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(VERIFY_DURATION_SECS))
|
||||
.await;
|
||||
|
||||
if self.docker.is_container_running_by_id(container_id).await? {
|
||||
println!("Container started successfully and verified running");
|
||||
|
||||
return Ok(());
|
||||
} else {
|
||||
println!(
|
||||
"Container stopped unexpectedly after start (attempt {}/{})",
|
||||
attempt, MAX_RETRIES
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!(
|
||||
"Failed to start container (attempt {}/{}): {}",
|
||||
attempt, MAX_RETRIES, e
|
||||
);
|
||||
}
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
async fn update_project_container(
|
||||
&self,
|
||||
project: &Project,
|
||||
state: &mut StateManager,
|
||||
) -> Result<String, miette::Error> {
|
||||
println!("Creating container {}...", project.container_name());
|
||||
let id = self
|
||||
.docker
|
||||
.create_postgres_container(
|
||||
&project.container_name(),
|
||||
&project.config.version,
|
||||
&project.config.password,
|
||||
project.config.port,
|
||||
)
|
||||
.await?;
|
||||
println!("Container created successfully");
|
||||
state.set(
|
||||
project.name.clone(),
|
||||
crate::state::InstanceState::new(
|
||||
id.clone(),
|
||||
project.config.version.to_string(),
|
||||
DATABASE.to_string(),
|
||||
USERNAME.to_string(),
|
||||
project.config.port,
|
||||
),
|
||||
);
|
||||
state.save()?;
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
async fn ensure_container_exists(
|
||||
&self,
|
||||
instance: &InstanceState,
|
||||
) -> Result<Option<String>, miette::Error> {
|
||||
let mut container_id = None;
|
||||
let id = &instance.container_id;
|
||||
if self.docker.container_exists_by_id(id).await? {
|
||||
container_id = Some(id.clone());
|
||||
}
|
||||
Ok(container_id)
|
||||
}
|
||||
|
||||
async fn ensure_matches_project_version(
|
||||
&self,
|
||||
project: &Project,
|
||||
_state: &mut StateManager,
|
||||
_container_id: &String,
|
||||
container_version: PostgresVersion,
|
||||
) -> Result<(), miette::Error> {
|
||||
Ok(if container_version != project.config.version {
|
||||
let needs_upgrade = container_version < project.config.version;
|
||||
|
||||
if needs_upgrade {
|
||||
bail!("Upgrades are currently unsupported! :(");
|
||||
// println!(
|
||||
// "Upgrading PostgreSQL from {} to {}...",
|
||||
// container_version, project.config.version
|
||||
// );
|
||||
// self.docker.stop_container(container_id, 10).await?;
|
||||
// self.docker
|
||||
// .upgrade_container_image(
|
||||
// container_id,
|
||||
// container_name,
|
||||
// &project.config.version,
|
||||
// &project.config.password,
|
||||
// project.config.port,
|
||||
// )
|
||||
// .await?;
|
||||
|
||||
// if let Some(instance_state) = state.get_mut(&project.name) {
|
||||
// instance_state.postgres_version = project.config.version.to_string();
|
||||
// state.save()?;
|
||||
// }
|
||||
} else {
|
||||
miette::bail!(
|
||||
"Cannot downgrade PostgreSQL from {} to {}. Downgrades are not supported.",
|
||||
container_version,
|
||||
project.config.version
|
||||
);
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
39
src/main.rs
Normal file
39
src/main.rs
Normal file
@@ -0,0 +1,39 @@
|
||||
mod cli;
|
||||
mod config;
|
||||
mod state;
|
||||
|
||||
mod controller;
|
||||
|
||||
use clap::Parser;
|
||||
use cli::Cli;
|
||||
use miette::Result;
|
||||
use tracing::info;
|
||||
|
||||
use crate::controller::Controller;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
println!("{}", include_str!("./banner.txt"));
|
||||
let controller = Controller::new().await?;
|
||||
|
||||
let cli = Cli::parse();
|
||||
init_tracing(cli.verbose);
|
||||
|
||||
info!("pgx.start");
|
||||
|
||||
match cli.command {
|
||||
cli::Commands::Init => controller.init_project().await?,
|
||||
cli::Commands::Instance { name, cmd } => todo!(),
|
||||
cli::Commands::Sync => todo!(),
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn init_tracing(verbose: bool) {
|
||||
use tracing_subscriber::{fmt, prelude::*};
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(fmt::layer().with_target(false).with_level(true))
|
||||
.init();
|
||||
}
|
||||
203
src/state.rs
Normal file
203
src/state.rs
Normal file
@@ -0,0 +1,203 @@
|
||||
use miette::{Context, IntoDiagnostic, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// State information for a single PostgreSQL instance
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InstanceState {
|
||||
/// Docker container ID
|
||||
pub container_id: String,
|
||||
|
||||
/// PostgreSQL version running in the container
|
||||
pub postgres_version: String,
|
||||
|
||||
/// Database name
|
||||
pub database_name: String,
|
||||
|
||||
/// User name
|
||||
pub user_name: String,
|
||||
|
||||
/// Port the container is bound to
|
||||
pub port: u16,
|
||||
|
||||
/// Timestamp when the instance was created (Unix timestamp)
|
||||
pub created_at: u64,
|
||||
|
||||
/// Timestamp when the instance was last started (Unix timestamp)
|
||||
pub last_started_at: Option<u64>,
|
||||
}
|
||||
|
||||
/// Manages the global state file at ~/.pgx/state.json
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StateManager {
|
||||
/// Map of project name to instance state
|
||||
#[serde(default)]
|
||||
instances: HashMap<String, InstanceState>,
|
||||
}
|
||||
|
||||
/// Get the path to the state file (~/.pgx/state.json)
|
||||
|
||||
fn state_file_path() -> Result<PathBuf> {
|
||||
let home = std::env::var("HOME")
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to get HOME environment variable")?;
|
||||
|
||||
Ok(PathBuf::from(home).join(".pgx").join("state.json"))
|
||||
}
|
||||
|
||||
/// Get the path to the .pgx directory
|
||||
pub fn pgx_dir() -> Result<PathBuf> {
|
||||
let home = std::env::var("HOME")
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to get HOME environment variable")?;
|
||||
|
||||
Ok(PathBuf::from(home).join(".pgx"))
|
||||
}
|
||||
|
||||
impl StateManager {
|
||||
/// Load the state manager from disk, or create a new one if it doesn't exist
|
||||
pub fn load() -> Result<Self> {
|
||||
let state_path = state_file_path()?;
|
||||
|
||||
if !state_path.exists() {
|
||||
// Create the directory if it doesn't exist
|
||||
if let Some(parent) = state_path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to create .pgx directory")?;
|
||||
}
|
||||
|
||||
// Return empty state
|
||||
return Ok(StateManager {
|
||||
instances: HashMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(&state_path)
|
||||
.into_diagnostic()
|
||||
.wrap_err_with(|| format!("Failed to read state file: {}", state_path.display()))?;
|
||||
|
||||
let state: StateManager = serde_json::from_str(&content)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to parse state.json")?;
|
||||
|
||||
Ok(state)
|
||||
}
|
||||
|
||||
/// Save the state manager to disk
|
||||
pub fn save(&self) -> Result<()> {
|
||||
let state_path = state_file_path()?;
|
||||
|
||||
// Ensure directory exists
|
||||
if let Some(parent) = state_path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to create .pgx directory")?;
|
||||
}
|
||||
|
||||
let content = serde_json::to_string_pretty(self)
|
||||
.into_diagnostic()
|
||||
.wrap_err("Failed to serialize state")?;
|
||||
|
||||
std::fs::write(&state_path, content)
|
||||
.into_diagnostic()
|
||||
.wrap_err_with(|| format!("Failed to write state file: {}", state_path.display()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the state for a specific project
|
||||
pub fn get(&self, project_name: &str) -> Option<&InstanceState> {
|
||||
self.instances.get(project_name)
|
||||
}
|
||||
|
||||
/// Get mutable state for a specific project
|
||||
pub fn get_mut(&mut self, project_name: &str) -> Option<&mut InstanceState> {
|
||||
self.instances.get_mut(project_name)
|
||||
}
|
||||
|
||||
/// Set the state for a specific project
|
||||
pub fn set(&mut self, project_name: String, state: InstanceState) {
|
||||
self.instances.insert(project_name, state);
|
||||
}
|
||||
|
||||
/// Update the state for a specific project, creating it if it doesn't exist
|
||||
pub fn update<F>(&mut self, project_name: &str, updater: F) -> Result<()>
|
||||
where
|
||||
F: FnOnce(&mut InstanceState),
|
||||
{
|
||||
if let Some(state) = self.instances.get_mut(project_name) {
|
||||
updater(state);
|
||||
Ok(())
|
||||
} else {
|
||||
miette::bail!("No state found for project: {}", project_name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove the state for a specific project
|
||||
pub fn remove(&mut self, project_name: &str) -> Option<InstanceState> {
|
||||
self.instances.remove(project_name)
|
||||
}
|
||||
|
||||
/// Get all instances
|
||||
pub fn all_instances(&self) -> &HashMap<String, InstanceState> {
|
||||
&self.instances
|
||||
}
|
||||
}
|
||||
|
||||
impl InstanceState {
|
||||
pub fn new(
|
||||
container_id: String,
|
||||
postgres_version: String,
|
||||
database_name: String,
|
||||
user_name: String,
|
||||
port: u16,
|
||||
) -> Self {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
|
||||
InstanceState {
|
||||
container_id,
|
||||
postgres_version,
|
||||
database_name,
|
||||
user_name,
|
||||
port,
|
||||
created_at: now,
|
||||
last_started_at: Some(now),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_state_manager_operations() {
|
||||
let mut manager = StateManager {
|
||||
instances: HashMap::new(),
|
||||
};
|
||||
|
||||
let state = InstanceState::new(
|
||||
"container123".to_string(),
|
||||
"16".to_string(),
|
||||
"mydb".to_string(),
|
||||
"postgres".to_string(),
|
||||
5432,
|
||||
);
|
||||
|
||||
manager.set("my-project".to_string(), state);
|
||||
|
||||
assert!(manager.get("my-project").is_some());
|
||||
assert_eq!(
|
||||
manager.get("my-project").unwrap().container_id,
|
||||
"container123"
|
||||
);
|
||||
|
||||
manager.remove("my-project");
|
||||
assert!(manager.get("my-project").is_none());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user