fix: correct vector example error

Author: brianheineman

Cargo.lock

@@ -33,6 +33,7 @@ hex = "0.4.3"
 home = "0.5.9"
 http = "1.1.0"
 human_bytes = { version = "0.4.3", default-features = false }
+indoc = "2.0.5"
 md-5 = "0.10.6"
 num-format = "0.4.4"
 quick-xml = "0.36.1"
@@ -56,6 +57,7 @@ tempfile = "3.10.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.6"

Cargo.toml

@@ -7,7 +7,10 @@ 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

@@ -2,46 +2,118 @@
 #![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?;
 
-    println!("Installing the vector extension from TensorChord");
+    info!("Installing the vector extension from TensorChord");
     postgresql_extensions::install(
-        &settings,
+        postgresql.settings(),
         "tensor-chord",
         "pgvecto.rs",
         &VersionReq::parse("=0.3.0")?,
     )
     .await?;
 
-    println!("Installing PostgreSQL");
-    let mut postgresql = PostgreSQL::new(settings);
-    postgresql.setup().await?;
+    info!("Starting PostgreSQL");
     postgresql.start().await?;
 
     let database_name = "vector-example";
-    println!("Creating database {database_name}");
+    info!("Creating database {database_name}");
     postgresql.create_database(database_name).await?;
 
-    println!("Connecting to database {database_name}");
+    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(())
+}
 
-    println!("Configuring the vector extension");
+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)
+        .execute(pool)
         .await?;
+    Ok(())
+}
 
-    println!("Stopping database");
-    postgresql.stop().await?;
+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(())
 }
 

examples/vector_extension/src/main.rs

@@ -1,4 +1,5 @@
-#[cfg(all(target_os = "linux", feature = "steampipe"))]
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+#[cfg(feature = "steampipe")]
 #[tokio::test]
 async fn test_lifecycle() -> anyhow::Result<()> {
     let installation_dir = tempfile::tempdir()?.path().to_path_buf();