draft: add Postgres extension for Ahnlich

ahnlich/types/Cargo.toml

@@ -16,6 +16,7 @@ prost.workspace = true
 prost-types = "0.13"
 serde.workspace = true
 ascii85 = "0.2.1"
+pgrx = { version = "=0.16.0", features = ["pg15"] }
 
 
 [build-dependencies]

ahnlich/types/build.rs

@@ -56,7 +56,7 @@ fn main() -> Result<()> {
         .open(out_dir.join("lib.rs"))
         .expect("Failed to create mod file");
 
-    // nonlinear algorthim, storekeyid, storevalue, metadatakey and value,
+    // nonlinear algorithm, storekeyid, storevalue, metadatakey and value
     tonic_build::configure()
         .build_client(true)
         .build_client(true)
@@ -65,23 +65,35 @@ fn main() -> Result<()> {
             "algorithm.nonlinear.NonLinearAlgorithm",
             "#[derive(serde::Serialize, serde::Deserialize)]",
         )
+        .type_attribute(
+            "keyval.StoreInput",
+            "#[derive(pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
+        .type_attribute(
+            "keyval.StoreInput.value",
+            "#[derive(PartialOrd, Ord, Hash, Eq, pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
         .type_attribute(
             "keyval.StoreValue",
-            "#[derive(serde::Serialize, serde::Deserialize)]",
+            "#[derive(pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
+        .type_attribute(
+            "keyval.StoreValue.value",
+            "#[derive(PartialOrd, Ord, Hash, Eq, pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
         )
         .type_attribute(
             "metadata.MetadataValue",
-            "#[derive(PartialOrd, Ord, Hash, Eq)]",
+            "#[derive(PartialOrd, Ord, Hash, Eq, pgrx::PostgresType)]",
+        )
+        .type_attribute(
+            "metadata.MetadataValue.value",
+            "#[derive(PartialOrd, Ord, Hash, Eq, serde::Serialize, serde::Deserialize)]",
         )
         .type_attribute("keyval.StoreName", "#[derive(Eq, Hash, Ord, PartialOrd)]")
         .type_attribute(
             "db.server.StoreInfo",
             "#[derive(Hash, Eq, Ord, PartialOrd)]",
         )
-        .type_attribute(
-            "metadata.MetadataValue.value",
-            "#[derive(PartialOrd, Ord, Hash, Eq)]",
-        )
         .type_attribute(
             "ai.models.AIModel",
             "#[derive(serde::Serialize, serde::Deserialize)]",
@@ -91,9 +103,21 @@ fn main() -> Result<()> {
             "#[derive(Eq, PartialOrd, Ord, Hash)]",
         )
         .type_attribute("client.ConnectedClient", "#[derive(PartialOrd, Ord, Eq)]")
+        .type_attribute(
+            "ai.server.SingleInputToEmbedding",
+            "#[derive(pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
+        .type_attribute(
+            "ai.query.ConvertStoreInputToEmbeddings",
+            "#[derive(pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
+        .type_attribute(
+            "ai.server.StoreInputToEmbeddingsList",
+            "#[derive(pgrx::PostgresType, serde::Serialize, serde::Deserialize)]",
+        )
         .compile_protos(&protofiles, &[proto_dir])
         .inspect_err(|err| println!("{err}"))
-        .expect("failed");
+        .expect("Failed");
 
     restructure_generated_code(&out_dir, &mut file);
 
@@ -162,5 +186,5 @@ fn restructure_generated_code(out_dir: &PathBuf, file: &mut std::fs::File) {
         .join("\n");
 
     file.write_all(buffer.as_bytes())
-        .expect("could not generate lib.rs");
+        .expect("Could not generate lib.rs");
 }

ahnlich/types/src/ai/query.rs

@@ -211,7 +211,9 @@ pub struct Set {
     )]
     pub execution_provider: ::core::option::Option<i32>,
 }
-#[derive(Clone, PartialEq, ::prost::Message)]
+#[derive(
+    pgrx::PostgresType, serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message,
+)]
 pub struct ConvertStoreInputToEmbeddings {
     /// Convert store inputs to embeddings
     ///

ahnlich/types/src/ai/server.rs

@@ -70,14 +70,18 @@ pub struct AiStoreInfo {
     #[prost(uint64, tag = "4")]
     pub embedding_size: u64,
 }
-#[derive(Clone, PartialEq, ::prost::Message)]
+#[derive(
+    pgrx::PostgresType, serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message,
+)]
 pub struct SingleInputToEmbedding {
     #[prost(message, optional, tag = "1")]
     pub input: ::core::option::Option<super::super::keyval::StoreInput>,
     #[prost(message, optional, tag = "2")]
     pub embedding: ::core::option::Option<super::super::keyval::StoreKey>,
 }
-#[derive(Clone, PartialEq, ::prost::Message)]
+#[derive(
+    pgrx::PostgresType, serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message,
+)]
 pub struct StoreInputToEmbeddingsList {
     #[prost(message, repeated, tag = "1")]
     pub values: ::prost::alloc::vec::Vec<SingleInputToEmbedding>,

ahnlich/types/src/keyval.rs

@@ -9,14 +9,27 @@ pub struct StoreKey {
     #[prost(float, repeated, tag = "1")]
     pub key: ::prost::alloc::vec::Vec<f32>,
 }
-#[derive(Clone, PartialEq, ::prost::Message)]
+#[derive(
+    pgrx::PostgresType, serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message,
+)]
 pub struct StoreInput {
     #[prost(oneof = "store_input::Value", tags = "2, 3")]
     pub value: ::core::option::Option<store_input::Value>,
 }
 /// Nested message and enum types in `StoreInput`.
 pub mod store_input {
-    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    #[derive(
+        PartialOrd,
+        Ord,
+        Hash,
+        Eq,
+        pgrx::PostgresType,
+        serde::Serialize,
+        serde::Deserialize,
+        Clone,
+        PartialEq,
+        ::prost::Oneof,
+    )]
     pub enum Value {
         #[prost(string, tag = "2")]
         RawString(::prost::alloc::string::String),
@@ -40,7 +53,9 @@ pub struct AiStoreEntry {
     #[prost(message, optional, tag = "2")]
     pub value: ::core::option::Option<StoreValue>,
 }
-#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message)]
+#[derive(
+    pgrx::PostgresType, serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message,
+)]
 pub struct StoreValue {
     #[prost(map = "string, message", tag = "1")]
     pub value:

ahnlich/types/src/metadata.rs

@@ -1,12 +1,22 @@
 // This file is @generated by prost-build.
-#[derive(PartialOrd, Ord, Hash, Eq, Clone, PartialEq, ::prost::Message)]
+#[derive(PartialOrd, Ord, Hash, Eq, pgrx::PostgresType, Clone, PartialEq, ::prost::Message)]
 pub struct MetadataValue {
     #[prost(oneof = "metadata_value::Value", tags = "2, 3")]
     pub value: ::core::option::Option<metadata_value::Value>,
 }
 /// Nested message and enum types in `MetadataValue`.
 pub mod metadata_value {
-    #[derive(PartialOrd, Ord, Hash, Eq, Clone, PartialEq, ::prost::Oneof)]
+    #[derive(
+        PartialOrd,
+        Ord,
+        Hash,
+        Eq,
+        serde::Serialize,
+        serde::Deserialize,
+        Clone,
+        PartialEq,
+        ::prost::Oneof,
+    )]
     pub enum Value {
         #[prost(string, tag = "2")]
         RawString(::prost::alloc::string::String),

extensions/rust/ahnlich-postgres-ext/.cargo/config.toml

@@ -0,0 +1,3 @@
+[target.'cfg(target_os="macos")']
+# Postgres symbols won't be available until runtime
+rustflags = ["-Clink-arg=-Wl,-undefined,dynamic_lookup"]

extensions/rust/ahnlich-postgres-ext/.gitignore

@@ -0,0 +1,9 @@
+.DS_Store
+.idea/
+/target
+*.iml
+**/*.rs.bk
+Cargo.lock
+pg_regress/results
+pg_regress/regression.diffs
+pg_regress/regression.out

extensions/rust/ahnlich-postgres-ext/Cargo.toml

@@ -0,0 +1,113 @@
+[package]
+name = "ahnlich_postgres_ext"
+version = "0.0.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib", "lib"]
+
+[[bin]]
+name = "pgrx_embed_ahnlich_postgres_ext"
+path = "./src/bin/pgrx_embed.rs"
+
+[features]
+default = ["pg15"]
+pg13 = ["pgrx/pg13", "pgrx-tests/pg13"]
+pg14 = ["pgrx/pg14", "pgrx-tests/pg14"]
+pg15 = ["pgrx/pg15", "pgrx-tests/pg15"]
+pg16 = ["pgrx/pg16", "pgrx-tests/pg16"]
+pg17 = ["pgrx/pg17", "pgrx-tests/pg17"]
+pg18 = ["pgrx/pg18", "pgrx-tests/pg18"]
+pg_test = []
+
+[dependencies]
+pgrx = "=0.16.0"
+ahnlich_client_rs = { path = "../../../ahnlich/client" }
+ahnlich_types = { path = "../../../ahnlich/types" }
+db = { path = "../../../ahnlich/db", version = "*" }
+ai = { path = "../../../ahnlich/ai", version = "*" }
+utils = { path = "../../../ahnlich/utils", version = "*" }
+once_cell = "1.19.0"
+serde_json = "1.0.116"
+tokio = { version = "1.37.0", features = [
+    "net",
+    "macros",
+    "io-util",
+    "rt-multi-thread",
+    "sync",
+    "signal",
+] }
+tonic = "0.12.3"
+# async-trait = "0.1"
+# serde = { version = "1.0.*", features = ["derive", "rc"] }
+# bincode = "1.3.3"
+# ndarray = { version = "0.16.1", features = ["serde", "rayon"] }
+# image = "0.25.5"
+# serde_json = "1.0.116"
+# itertools = "0.10.0"
+# clap = { version = "4.5.4", features = ["derive"] }
+# futures = "0.3.30"
+# pretty_assertions = "1.4.0"
+# tracing = "0.1"
+# thiserror = "1.0"
+# blake3 = "1.5.1"
+# tokio-util = { version = "0.7.11", features = ["rt"] }
+# rand = "0.8"
+# rayon = "1.10"
+# opentelemetry = { version = "0.23.0", features = ["trace"] }
+# tracing-opentelemetry = "0.24.0"
+# log = "0.4"
+# fallible_collections = "0.4.9"
+# dirs = "5.0.1"
+# strum = { version = "0.26", features = ["derive"] }
+# papaya = { version = "0.2.0", features = ["serde"] }
+# tonic = "0.12.3"
+# http = "1.2.0"
+# prost = "0.13"
+
+[dev-dependencies]
+pgrx-tests = "=0.16.0"
+
+[profile.dev]
+panic = "unwind"
+
+[profile.release]
+panic = "unwind"
+opt-level = 3
+lto = "fat"
+codegen-units = 1
+
+# [package]
+# name = "ahnlich_client_rs"
+# version = "0.1.0"
+# authors = ["Diretnan Domnan <diretnandomnan@gmail.com>"]
+# categories = ["database-implementations", "database", "web-programming"]
+# keywords = ["ahnlich", "in-memory", "ai"]
+# edition = "2024"
+# license-file = "LICENSE"
+# homepage = "https://github.com/deven96/ahnlich"
+# repository = "https://github.com/deven96/ahnlich"
+# description = "Client library for accessing in-memory ahnlich vector database and ai-proxy"
+
+# [badges]
+# maintenance = { status = "actively-developed" }
+# # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+# [dependencies]
+# ahnlich_types = { path = "../types", version = "*" }
+# thiserror.workspace = true
+# once_cell.workspace = true
+# bincode.workspace = true
+# async-trait.workspace = true
+# tokio.workspace = true
+# tonic.workspace = true
+# http.workspace = true
+# fallible_collections.workspace = true
+
+# [dev-dependencies]
+# db = { path = "../db", version = "*" }
+# ai = { path = "../ai", version = "*" }
+# pretty_assertions.workspace = true
+# ndarray.workspace = true
+# utils = { path = "../utils", version = "*" }
+# test-case = "3.3.1"

extensions/rust/ahnlich-postgres-ext/README.md

@@ -0,0 +1,33 @@
+# Ahnlich Postgres extension (beta)
+
+## Setup instructions
+
+1. Go to: <https://github.com/pgcentralfoundation/pgrx?tab=readme-ov-file#system-requirements> & <https://github.com/pgcentralfoundation/pgrx?tab=readme-ov-file#getting-started> and follow the instructions to install pgrx on your local machine
+
+## Run instructions
+
+```sh
+cargo pgrx run
+```
+
+When the `psql` shell opens, now you can run:
+
+```sql
+CREATE EXTENSION ahnlich_postgres_ext;
+```
+
+Next, run the Ahnlich client init function:
+
+```sql
+SELECT init_client('');
+```
+
+Then you can run any of the other functions defined in `lib.rs`:
+
+```sql
+SELECT ping();
+```
+
+```sql
+SELECT ping_with_args(2992);
+```

extensions/rust/ahnlich-postgres-ext/ahnlich_postgres_ext.control

@@ -0,0 +1,6 @@
+comment = 'ahnlich_postgres_ext:  Created by pgrx'
+default_version = '@CARGO_VERSION@'
+module_pathname = 'ahnlich_postgres_ext'
+relocatable = false
+superuser = true
+trusted = false

extensions/rust/ahnlich-postgres-ext/pg_regress/expected/setup.out

@@ -0,0 +1,3 @@
+-- this setup file is run immediately after the regression database is (re)created
+-- the file is optional but you likely want to create the extension
+CREATE EXTENSION ahnlich_postgres_ext;

extensions/rust/ahnlich-postgres-ext/pg_regress/sql/setup.sql

@@ -0,0 +1,3 @@
+-- this setup file is run immediately after the regression database is (re)created
+-- the file is optional but you likely want to create the extension
+CREATE EXTENSION ahnlich_postgres_ext;

extensions/rust/ahnlich-postgres-ext/src/bin/pgrx_embed.rs

@@ -0,0 +1 @@
+::pgrx::pgrx_embed!();