collab: Setup database for LLM service (#15882)

This PR puts the initial infrastructure for the LLM service's database
in place.

The LLM service will be using a separate Postgres database, with its own
set of migrations.

Currently we only connect to the database in development, as we don't
yet have the database setup for the staging/production environments.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-08-06 17:18:08 -04:00 committed by GitHub
parent a64906779b
commit 7f6d0919c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 627 additions and 74 deletions

View file

@ -0,0 +1,49 @@
use std::path::Path;
use std::time::Duration;
use anyhow::{anyhow, Result};
use collections::HashMap;
use sea_orm::ConnectOptions;
use sqlx::migrate::{Migrate, Migration, MigrationSource};
use sqlx::Connection;
/// Runs the database migrations for the specified database.
pub async fn run_database_migrations(
database_options: &ConnectOptions,
migrations_path: impl AsRef<Path>,
ignore_checksum_mismatch: bool,
) -> Result<Vec<(Migration, Duration)>> {
let migrations = MigrationSource::resolve(migrations_path.as_ref())
.await
.map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
let mut connection = sqlx::AnyConnection::connect(database_options.get_url()).await?;
connection.ensure_migrations_table().await?;
let applied_migrations: HashMap<_, _> = connection
.list_applied_migrations()
.await?
.into_iter()
.map(|migration| (migration.version, migration))
.collect();
let mut new_migrations = Vec::new();
for migration in migrations {
match applied_migrations.get(&migration.version) {
Some(applied_migration) => {
if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch {
Err(anyhow!(
"checksum mismatch for applied migration {}",
migration.description
))?;
}
}
None => {
let elapsed = connection.apply(&migration).await?;
new_migrations.push((migration, elapsed));
}
}
}
Ok(new_migrations)
}