Merge pull request #110 from theseus-rs/initial-postgresql-extensions

feat!: initial postgresql_extensions crate

Author: brianheineman

Cargo.lock

@@ -3,12 +3,14 @@ default-members = [
     "postgresql_archive",
     "postgresql_commands",
     "postgresql_embedded",
+    "postgresql_extensions",
 ]
 members = [
     "examples/*",
     "postgresql_archive",
     "postgresql_commands",
     "postgresql_embedded",
+    "postgresql_extensions",
 ]
 resolver = "2"
 
@@ -31,10 +33,10 @@ hex = "0.4.3"
 home = "0.5.9"
 http = "1.1.0"
 human_bytes = { version = "0.4.3", default-features = false }
-lazy_static = "1.5.0"
+indoc = "2.0.5"
 md-5 = "0.10.6"
 num-format = "0.4.4"
-quick-xml = "0.35.0"
+quick-xml = "0.36.1"
 rand = "0.8.5"
 regex = "1.10.5"
 reqwest = { version = "0.12.5", default-features = false }
@@ -43,21 +45,22 @@ reqwest-retry = "0.6.0"
 reqwest-tracing = "0.5.2"
 semver = "1.0.23"
 serde = "1.0.204"
-serde_json = "1.0.118"
+serde_json = "1.0.121"
 sha1 = "0.10.6"
 sha2 = "0.10.8"
 sha3 = "0.10.8"
-sqlx = { version = "0.7.4", default-features = false, features = ["postgres"] }
+sqlx = { version = "0.8.0", default-features = false, features = ["postgres"] }
 tar = "0.4.41"
 target-triple = "0.1.3"
 test-log = "0.2.16"
 tempfile = "3.10.1"
-thiserror = "1.0.62"
-tokio = "1.38.1"
+thiserror = "1.0.63"
+tokio = "1.39.2"
 tracing = "0.1.40"
+tracing-subscriber = "0.3.18"
 url = "2.5.2"
 xz2 = "0.1.7"
-zip = "2.1.3"
+zip = "2.1.6"
 
 [workspace.metadata.release]
 shared-version = true

Cargo.toml

@@ -7,5 +7,5 @@ version.workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
-postgres = { version = "0.19.7" }
+postgres = { version = "0.19.8" }
 postgresql_embedded = { path = "../../postgresql_embedded", features = ["blocking"] }

examples/postgres_embedded/Cargo.toml

@@ -8,5 +8,5 @@ version.workspace = true
 [dependencies]
 anyhow = { workspace = true }
 postgresql_embedded = { path = "../../postgresql_embedded" }
-sqlx = { version = "0.7.4", default-features = true, features = ["postgres", "runtime-tokio"] }
+sqlx = { workspace = true, features = ["runtime-tokio"] }
 tokio = { workspace = true, features = ["full"] }

examples/sqlx_embedded/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+edition.workspace = true
+name = "vector_extension"
+publish = false
+license.workspace = true
+version.workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+indoc = { workspace = true }
+postgresql_extensions = { path = "../../postgresql_extensions" }
+postgresql_embedded = { path = "../../postgresql_embedded" }
+sqlx = { workspace = true, features = ["runtime-tokio"] }
+tracing = { workspace = true }
+tracing-subscriber = { workspace = true }
+tokio = { workspace = true, features = ["full"] }

examples/vector_extension/Cargo.toml

@@ -0,0 +1,129 @@
+#![forbid(unsafe_code)]
+#![deny(clippy::pedantic)]
+
+use anyhow::Result;
+use indoc::indoc;
+use postgresql_embedded::{PostgreSQL, Settings, VersionReq};
+use sqlx::PgPool;
+use tracing::info;
+
+/// Example of how to install and configure the vector extension.
+///
+/// See: <https://github.com/tensorchord/pgvecto.rs/?tab=readme-ov-file#quick-start>
+#[tokio::main]
+async fn main() -> Result<()> {
+    tracing_subscriber::fmt().compact().init();
+
+    info!("Installing PostgreSQL");
+    let settings = Settings {
+        version: VersionReq::parse("=16.3.0")?,
+        ..Default::default()
+    };
+    let mut postgresql = PostgreSQL::new(settings);
+    postgresql.setup().await?;
+
+    info!("Installing the vector extension from TensorChord");
+    postgresql_extensions::install(
+        postgresql.settings(),
+        "tensor-chord",
+        "pgvecto.rs",
+        &VersionReq::parse("=0.3.0")?,
+    )
+    .await?;
+
+    info!("Starting PostgreSQL");
+    postgresql.start().await?;
+
+    let database_name = "vector-example";
+    info!("Creating database {database_name}");
+    postgresql.create_database(database_name).await?;
+
+    info!("Configuring extension");
+    let settings = postgresql.settings();
+    let database_url = settings.url(database_name);
+    let pool = PgPool::connect(database_url.as_str()).await?;
+    configure_extension(&pool).await?;
+    pool.close().await;
+
+    info!("Restarting database");
+    postgresql.stop().await?;
+    postgresql.start().await?;
+
+    info!("Enabling extension");
+    let pool = PgPool::connect(database_url.as_str()).await?;
+    enable_extension(&pool).await?;
+
+    info!("Creating table");
+    create_table(&pool).await?;
+
+    info!("Creating data");
+    create_data(&pool).await?;
+
+    info!("Stopping database");
+    postgresql.stop().await?;
+    Ok(())
+}
+
+async fn configure_extension(pool: &PgPool) -> Result<()> {
+    sqlx::query("ALTER SYSTEM SET shared_preload_libraries = \"vectors.so\"")
+        .execute(pool)
+        .await?;
+    sqlx::query("ALTER SYSTEM SET search_path = \"$user\", public, vectors")
+        .execute(pool)
+        .await?;
+    Ok(())
+}
+
+async fn enable_extension(pool: &PgPool) -> Result<()> {
+    sqlx::query("DROP EXTENSION IF EXISTS vectors")
+        .execute(pool)
+        .await?;
+    sqlx::query("CREATE EXTENSION vectors")
+        .execute(pool)
+        .await?;
+    Ok(())
+}
+
+async fn create_table(pool: &PgPool) -> Result<()> {
+    sqlx::query(indoc! {"
+        CREATE TABLE IF NOT EXISTS items (
+            id bigserial PRIMARY KEY,
+            embedding vector(3) NOT NULL
+        )
+    "})
+    .execute(pool)
+    .await?;
+    Ok(())
+}
+
+async fn create_data(pool: &PgPool) -> Result<()> {
+    sqlx::query(indoc! {"
+        INSERT INTO items (embedding)
+        VALUES
+            ('[1,2,3]'),
+            ('[4,5,6]')
+    "})
+    .execute(pool)
+    .await?;
+    sqlx::query(indoc! {"
+        INSERT INTO items (embedding)
+        VALUES
+            (ARRAY[1, 2, 3]::real[]),
+            (ARRAY[4, 5, 6]::real[]
+        )
+    "})
+    .execute(pool)
+    .await?;
+    Ok(())
+}
+
+#[cfg(target_os = "linux")]
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_main() -> Result<()> {
+        main()
+    }
+}

examples/vector_extension/src/main.rs

@@ -16,7 +16,6 @@ flate2 = { workspace = true, optional = true }
 hex = { workspace = true }
 http = { workspace = true }
 human_bytes = { workspace = true, default-features = false }
-lazy_static = { workspace = true }
 md-5 = { workspace = true, optional = true }
 num-format = { workspace = true }
 quick-xml = { workspace = true, features = ["serialize"], optional = true }