Skip to content

init db and table, check health #55

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

init db and table, check health #55

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6b09b86
init db and table, check health
tankyleo Aug 15, 2025
14c39a9
fixup: Remove `String` allocation
tankyleo Aug 15, 2025
837b300
fixup: Add required `time` tokio feature
tankyleo Aug 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions rust/impls/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ chrono = "0.4.38"
tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4"] }
bb8-postgres = "0.7"
bytes = "1.4.0"
tokio = { version = "1.38.0", default-features = false, features = ["time"] }

[dev-dependencies]
tokio = { version = "1.38.0", default-features = false, features = ["rt-multi-thread", "macros"] }
Expand Down
110 changes: 106 additions & 4 deletions rust/impls/src/postgres_store.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,22 @@ const KEY_COLUMN: &str = "key";
const VALUE_COLUMN: &str = "value";
const VERSION_COLUMN: &str = "version";

const DB_CHECK_STMT: &str = "SELECT 1 FROM pg_database WHERE datname = $1";
const DB_INIT_CMD: &str = "CREATE DATABASE";
const TABLE_CHECK_STMT: &str = "SELECT 1 FROM vss_db WHERE false";
const TABLE_INIT_STMT: &str = "
CREATE TABLE IF NOT EXISTS vss_db (
user_token character varying(120) NOT NULL CHECK (user_token <> ''),
store_id character varying(120) NOT NULL CHECK (store_id <> ''),
key character varying(600) NOT NULL,
value bytea NULL,
version bigint NOT NULL,
created_at TIMESTAMP WITH TIME ZONE,
last_updated_at TIMESTAMP WITH TIME ZONE,
PRIMARY KEY (user_token, store_id, key)
);
";

/// The maximum number of key versions that can be returned in a single page.
///
/// This constant helps control memory and bandwidth usage for list operations,
Expand All @@ -46,17 +62,103 @@ pub struct PostgresBackendImpl {
pool: Pool<PostgresConnectionManager<NoTls>>,
}

async fn initialize_vss_database(postgres_endpoint: &str, db_name: &str) -> Result<(), Error> {
let postgres_dsn = format!("{}/{}", postgres_endpoint, "postgres");
let (client, connection) = tokio_postgres::connect(&postgres_dsn, NoTls).await
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, do we want to expose a way for the user to configure their connection, e.g. to enable TLS?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes certainly, though may be out of scope for this PR.

Copy link
Author

@tankyleo tankyleo Aug 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Filed issue #56

.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;
// Connection must be driven on separate task, will be dropped on client dropped
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does it need to be driven on separate task? Should we still keep track of this connection handle somewhere? Rather than just printing once and giving up, should we handle the error and/or add a loop to retry the connection?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does it need to be driven on a separate task?

From the docs:

A connection to a PostgreSQL database.

This is one half of what is returned when a new connection is established. It performs the actual IO with the server, and should generally be spawned off onto an executor to run in the background.

Connection implements Future, and only resolves when the connection is closed, either because a fatal error has occurred, or because its associated Client has dropped and all outstanding work has completed.

Is that satisfactory to you ?

Should we still keep track of this connection handle somewhere?

The connection future will only resolve when I drop the client, so not sure what I would do with the associated JoinHandle ?

Rather than just printing once and giving up, should we handle the error and/or add a loop to retry the connection?

Since this is part of the startup sequence, I was leaning towards having very little error handling, and failing fast and early. This lets the user know that something is not right.

Then once we are running, do all that we can to keep the show going.

tokio::spawn(async move {
if let Err(e) = connection.await {
eprintln!("connection error: {}", e);
}
});

// Check if the database already exists
let num_rows = client.execute(DB_CHECK_STMT, &[&db_name]).await
.map_err(|e| Error::new(ErrorKind::Other, format!("Query error: {}", e)))?;

if num_rows == 0 {
// Database does not exist, so create it
let stmt = format!("{} {}", DB_INIT_CMD, db_name);
client.execute(&stmt, &[]).await
.map_err(|e| Error::new(ErrorKind::Other, format!("Query error: {}", e)))?;
println!("Database '{}' created successfully", db_name);
} else {
println!("Database '{}' already exists skipping creation", db_name);
}

Ok(())
}


impl PostgresBackendImpl {
/// Constructs a [`PostgresBackendImpl`] using `dsn` for PostgreSQL connection information.
pub async fn new(dsn: &str) -> Result<Self, Error> {
let manager = PostgresConnectionManager::new_from_stringlike(dsn, NoTls).map_err(|e| {
pub async fn new(postgres_endpoint: &str, db_name: &str, init_db: bool) -> Result<Self, Error> {
if init_db {
tokio::time::timeout(
tokio::time::Duration::from_secs(3),
initialize_vss_database(postgres_endpoint, db_name),
)
.await
.map_err(|e| Error::new(ErrorKind::Other, format!("Is the postgres endpoint online? {}", e)))?
.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;
}
let vss_dsn = format!("{}/{}", postgres_endpoint, db_name);
let manager = PostgresConnectionManager::new_from_stringlike(vss_dsn, NoTls).map_err(|e| {
Error::new(ErrorKind::Other, format!("Connection manager error: {}", e))
})?;
let pool = Pool::builder()
.build(manager)
.await
.map_err(|e| Error::new(ErrorKind::Other, format!("Pool build error: {}", e)))?;
Ok(PostgresBackendImpl { pool })
let ret = PostgresBackendImpl { pool };
let touch_table = async {
if init_db {
ret.initialize_vss_table().await?;
}
ret.check_health().await
};
tokio::time::timeout(
tokio::time::Duration::from_secs(3),
touch_table,
)
.await
.map_err(|e| Error::new(ErrorKind::Other, format!("Does the database exist? If not use --init-db {}", e)))?
.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;

Ok(ret)
}

async fn initialize_vss_table(&self) -> Result<(), Error> {
let conn = self
.pool
.get()
.await
.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;
let num_rows = conn
.execute(TABLE_INIT_STMT, &[])
.await
.map_err(|e| {
Error::new(ErrorKind::Other, format!("Database operation failed. {}", e))
})?;
assert_eq!(num_rows, 0);
Ok(())
}

async fn check_health(&self) -> Result<(), Error> {
let conn = self
.pool
.get()
.await
.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;
let num_rows = conn
.execute(TABLE_CHECK_STMT, &[])
.await
.map_err(|e| {
Error::new(ErrorKind::Other, format!("Does the table exist? If not use --init-db {}", e))
})?;
assert_eq!(num_rows, 0);
Ok(())
}

fn build_vss_record(&self, user_token: String, store_id: String, kv: KeyValue) -> VssDbRecord {
Expand Down Expand Up @@ -413,7 +515,7 @@ mod tests {
define_kv_store_tests!(
PostgresKvStoreTest,
PostgresBackendImpl,
PostgresBackendImpl::new("postgresql://postgres:postgres@localhost:5432/postgres")
PostgresBackendImpl::new("postgresql://postgres:postgres@localhost:5432", "postgres", false)
.await
.unwrap()
);
Expand Down
10 changes: 7 additions & 3 deletions rust/server/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,11 @@ pub(crate) mod vss_service;

fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() != 2 {
eprintln!("Usage: {} <config-file-path>", args[0]);
if args.len() < 2 {
eprintln!("Usage: {} <config-file-path> [--init-db]", args[0]);
std::process::exit(1);
}
let init_db = args.iter().any(|arg| arg == "--init-db");

let config = match util::config::load_config(&args[1]) {
Ok(cfg) => cfg,
Expand Down Expand Up @@ -67,13 +68,16 @@ fn main() {
},
};
let authorizer = Arc::new(NoopAuthorizer {});
let postgresql_config = config.postgresql_config.expect("PostgreSQLConfig must be defined in config file.");
let store = Arc::new(
PostgresBackendImpl::new(&config.postgresql_config.expect("PostgreSQLConfig must be defined in config file.").to_connection_string())
PostgresBackendImpl::new(&postgresql_config.to_postgresql_endpoint(), &postgresql_config.database, init_db)
.await
.unwrap(),
);
println!("Loaded postgres!");
let rest_svc_listener =
TcpListener::bind(&addr).await.expect("Failed to bind listening port");
println!("Bound to {}", addr);
loop {
tokio::select! {
res = rest_svc_listener.accept() => {
Expand Down
6 changes: 3 additions & 3 deletions rust/server/src/util/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ pub(crate) struct PostgreSQLConfig {
}

impl PostgreSQLConfig {
pub(crate) fn to_connection_string(&self) -> String {
pub(crate) fn to_postgresql_endpoint(&self) -> String {
let username_env = std::env::var("VSS_POSTGRESQL_USERNAME");
let username = username_env.as_ref()
.ok()
Expand All @@ -35,8 +35,8 @@ impl PostgreSQLConfig {
.expect("PostgreSQL database password must be provided in config or env var VSS_POSTGRESQL_PASSWORD must be set.");

format!(
"postgresql://{}:{}@{}:{}/{}",
username, password, self.host, self.port, self.database
"postgresql://{}:{}@{}:{}",
username, password, self.host, self.port
)
}
}
Expand Down
2 changes: 1 addition & 1 deletion rust/server/vss-server-config.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,4 @@ username = "postgres" # Optional in TOML, can be overridden by env var `VSS_POS
password = "postgres" # Optional in TOML, can be overridden by env var `VSS_POSTGRESQL_PASSWORD`
host = "localhost"
port = 5432
database = "postgres"
database = "vss_db"
Loading