Skip to content

Commit fe593a9

Browse files
SilasMarvinSilasMarvin
committed
Gated rust_bridge behind a feature flag
1 parent 5661eec commit fe593a9

File tree

12 files changed

+128
-75
lines changed

12 files changed

+128
-75
lines changed

‎pgml-sdks/pgml/Cargo.lock

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pgml-sdks/pgml/Cargo.toml

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@ name = "pgml"
1414
crate-type = ["lib", "cdylib"]
1515

1616
[dependencies]
17-
rust_bridge = {path = "../rust-bridge/rust-bridge", version = "0.1.0"}
17+
# rust_bridge = {path = "../rust-bridge/rust-bridge", version = "0.1.0", optional = true }
18+
rust_bridge = {git = "https://github.com/postgresml/postgresml", version = "0.1.0", optional = true }
1819
sqlx = { version = "0.7.3", features = [ "runtime-tokio-rustls", "postgres", "json", "time", "uuid"] }
1920
serde_json = "1.0.9"
2021
anyhow = "1.0.9"
@@ -50,5 +51,6 @@ serde_with = "3.8.1"
5051

5152
[features]
5253
default = []
53-
python = ["dep:pyo3", "dep:pyo3-asyncio"]
54-
javascript = ["dep:neon"]
54+
rust_bridge = ["dep:rust_bridge"]
55+
python = ["rust_bridge", "dep:pyo3", "dep:pyo3-asyncio"]
56+
javascript = ["rust_bridge", "dep:neon"]

‎pgml-sdks/pgml/src/builtins.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,26 @@
11
use anyhow::Context;
2-
use rust_bridge::{alias, alias_methods};
32
use sqlx::Row;
43
use tracing::instrument;
54

6-
/// Provides access to builtin database methods
7-
#[derive(alias, Debug, Clone)]
8-
pub struct Builtins {
9-
database_url: Option<String>,
10-
}
11-
125
use crate::{get_or_initialize_pool, query_runner::QueryRunner, types::Json};
136

7+
#[cfg(feature = "rust_bridge")]
8+
use rust_bridge::{alias, alias_methods};
9+
1410
#[cfg(feature = "python")]
1511
use crate::{query_runner::QueryRunnerPython, types::JsonPython};
1612

17-
#[alias_methods(new, query, transform, embed, embed_batch)]
13+
/// Provides access to builtin database methods
14+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
15+
#[derive(Debug, Clone)]
16+
pub struct Builtins {
17+
database_url: Option<String>,
18+
}
19+
20+
#[cfg_attr(
21+
feature = "rust_bridge",
22+
alias_methods(new, query, transform, embed, embed_batch)
23+
)]
1824
impl Builtins {
1925
pub fn new(database_url: Option<String>) -> Self {
2026
Self { database_url }

‎pgml-sdks/pgml/src/collection.rs

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ use anyhow::Context;
22
use indicatif::MultiProgress;
33
use itertools::Itertools;
44
use regex::Regex;
5-
use rust_bridge::{alias, alias_methods};
65
use sea_query::Alias;
76
use sea_query::{Expr, NullOrdering, Order, PostgresQueryBuilder, Query};
87
use sea_query_binder::SqlxBinder;
@@ -35,6 +34,9 @@ use crate::{
3534
utils,
3635
};
3736

37+
#[cfg(feature = "rust_bridge")]
38+
use rust_bridge::{alias, alias_methods};
39+
3840
#[cfg(feature = "python")]
3941
use crate::{
4042
pipeline::PipelinePython,
@@ -43,7 +45,7 @@ use crate::{
4345
};
4446

4547
/// A RAGStream Struct
46-
#[derive(alias)]
48+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
4749
#[allow(dead_code)]
4850
pub struct RAGStream {
4951
general_json_async_iterator: Option<GeneralJsonAsyncIterator>,
@@ -57,7 +59,7 @@ impl Clone for RAGStream {
5759
}
5860
}
5961

60-
#[alias_methods(stream, sources)]
62+
#[cfg_attr(feature = "rust_bridge", alias_methods(stream, sources))]
6163
impl RAGStream {
6264
pub fn stream(&mut self) -> anyhow::Result<GeneralJsonAsyncIterator> {
6365
self.general_json_async_iterator
@@ -137,7 +139,8 @@ pub(crate) struct CollectionDatabaseData {
137139
}
138140

139141
/// A collection of documents
140-
#[derive(alias, Debug, Clone)]
142+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
143+
#[derive(Debug, Clone)]
141144
pub struct Collection {
142145
pub(crate) name: String,
143146
pub(crate) database_url: Option<String>,
@@ -146,29 +149,32 @@ pub struct Collection {
146149
pub(crate) database_data: Option<CollectionDatabaseData>,
147150
}
148151

149-
#[alias_methods(
150-
new,
151-
upsert_documents,
152-
get_documents,
153-
delete_documents,
154-
get_pipelines,
155-
get_pipeline,
156-
add_pipeline,
157-
remove_pipeline,
158-
enable_pipeline,
159-
disable_pipeline,
160-
search,
161-
add_search_event,
162-
vector_search,
163-
query,
164-
rag,
165-
rag_stream,
166-
exists,
167-
archive,
168-
upsert_directory,
169-
upsert_file,
170-
generate_er_diagram,
171-
get_pipeline_status
152+
#[cfg_attr(
153+
feature = "rust_bridge",
154+
alias_methods(
155+
new,
156+
upsert_documents,
157+
get_documents,
158+
delete_documents,
159+
get_pipelines,
160+
get_pipeline,
161+
add_pipeline,
162+
remove_pipeline,
163+
enable_pipeline,
164+
disable_pipeline,
165+
search,
166+
add_search_event,
167+
vector_search,
168+
query,
169+
rag,
170+
rag_stream,
171+
exists,
172+
archive,
173+
upsert_directory,
174+
upsert_file,
175+
generate_er_diagram,
176+
get_pipeline_status
177+
)
172178
)]
173179
impl Collection {
174180
/// Creates a new [Collection]

‎pgml-sdks/pgml/src/model.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use rust_bridge::{alias, alias_methods};
21
use sqlx::{Pool, Postgres};
32
use tracing::instrument;
43

@@ -11,6 +10,9 @@ use crate::{
1110
#[cfg(feature = "python")]
1211
use crate::types::JsonPython;
1312

13+
#[cfg(feature = "rust_bridge")]
14+
use rust_bridge::{alias, alias_methods};
15+
1416
/// A few notes on the following enums:
1517
/// - Sqlx does provide type derivation for enums, but it's not very good
1618
/// - Queries using these enums require a number of additional queries to get their oids and
@@ -52,7 +54,8 @@ pub(crate) struct ModelDatabaseData {
5254
}
5355

5456
/// A model used for embedding, inference, etc...
55-
#[derive(alias, Debug, Clone)]
57+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
58+
#[derive(Debug, Clone)]
5659
pub struct Model {
5760
pub(crate) name: String,
5861
pub(crate) runtime: ModelRuntime,
@@ -66,7 +69,7 @@ impl Default for Model {
6669
}
6770
}
6871

69-
#[alias_methods(new, transform)]
72+
#[cfg_attr(feature = "rust_bridge", alias_methods(new, transform))]
7073
impl Model {
7174
/// Creates a new [Model]
7275
pub fn new(name: Option<String>, source: Option<String>, parameters: Option<Json>) -> Self {

‎pgml-sdks/pgml/src/open_source_ai.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use anyhow::Context;
22
use futures::{Stream, StreamExt};
3-
use rust_bridge::{alias, alias_methods};
43
use std::time::{SystemTime, UNIX_EPOCH};
54
use uuid::Uuid;
65

@@ -10,11 +9,15 @@ use crate::{
109
TransformerPipeline,
1110
};
1211

12+
#[cfg(feature = "rust_bridge")]
13+
use rust_bridge::{alias, alias_methods};
14+
1315
#[cfg(feature = "python")]
1416
use crate::types::{GeneralJsonAsyncIteratorPython, GeneralJsonIteratorPython, JsonPython};
1517

1618
/// A drop in replacement for OpenAI
17-
#[derive(alias, Debug, Clone)]
19+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
20+
#[derive(Debug, Clone)]
1821
pub struct OpenSourceAI {
1922
database_url: Option<String>,
2023
}
@@ -160,12 +163,15 @@ impl Iterator for AsyncToSyncJsonIterator {
160163
}
161164
}
162165

163-
#[alias_methods(
164-
new,
165-
chat_completions_create,
166-
chat_completions_create_async,
167-
chat_completions_create_stream,
168-
chat_completions_create_stream_async
166+
#[cfg_attr(
167+
feature = "rust_bridge",
168+
alias_methods(
169+
new,
170+
chat_completions_create,
171+
chat_completions_create_async,
172+
chat_completions_create_stream,
173+
chat_completions_create_stream_async
174+
)
169175
)]
170176
impl OpenSourceAI {
171177
/// Creates a new [OpenSourceAI]

‎pgml-sdks/pgml/src/pipeline.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use anyhow::Context;
2-
use rust_bridge::{alias, alias_methods};
32
use serde::Deserialize;
43
use serde_json::json;
54
use sqlx::{Executor, PgConnection, Pool, Postgres, Transaction};
@@ -16,6 +15,9 @@ use crate::{
1615
types::{DateTime, Json, TryToNumeric},
1716
};
1817

18+
#[cfg(feature = "rust_bridge")]
19+
use rust_bridge::{alias, alias_methods};
20+
1921
#[cfg(feature = "python")]
2022
use crate::types::JsonPython;
2123

@@ -176,7 +178,8 @@ pub struct PipelineDatabaseData {
176178
}
177179

178180
/// A pipeline that describes transformations to documents
179-
#[derive(alias, Debug, Clone)]
181+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
182+
#[derive(Debug, Clone)]
180183
pub struct Pipeline {
181184
pub(crate) name: String,
182185
pub(crate) schema: Option<Json>,
@@ -202,7 +205,7 @@ fn json_to_schema(schema: &Json) -> anyhow::Result<ParsedSchema> {
202205
})
203206
}
204207

205-
#[alias_methods(new)]
208+
#[cfg_attr(feature = "rust_bridge", alias_methods(new))]
206209
impl Pipeline {
207210
/// Creates a [Pipeline]
208211
///

‎pgml-sdks/pgml/src/query_builder.rs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,29 @@
33
// No new things should be added here, instead add new items to collection.vector_search
44

55
use anyhow::Context;
6-
use rust_bridge::{alias, alias_methods};
76
use serde_json::json;
87
use tracing::instrument;
98

109
use crate::{pipeline::Pipeline, types::Json, Collection};
1110

11+
#[cfg(feature = "rust_bridge")]
12+
use rust_bridge::{alias, alias_methods};
13+
1214
#[cfg(feature = "python")]
1315
use crate::{pipeline::PipelinePython, types::JsonPython};
1416

15-
#[derive(alias, Clone, Debug)]
17+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
18+
#[derive(Clone, Debug)]
1619
pub struct QueryBuilder {
1720
collection: Collection,
1821
query: Json,
1922
pipeline: Option<Pipeline>,
2023
}
2124

22-
#[alias_methods(limit, filter, vector_recall, to_full_string, fetch_all)]
25+
#[cfg_attr(
26+
feature = "rust_bridge",
27+
alias_methods(limit, filter, vector_recall, to_full_string, fetch_all(skip = "C"))
28+
)]
2329
impl QueryBuilder {
2430
pub fn new(collection: Collection) -> Self {
2531
let query = json!({

‎pgml-sdks/pgml/src/query_runner.rs

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use rust_bridge::{alias, alias_methods};
21
use sqlx::postgres::PgArguments;
32
use sqlx::query::Query;
43
use sqlx::{Postgres, Row};
@@ -8,6 +7,9 @@ use crate::{get_or_initialize_pool, types::Json};
87
#[cfg(feature = "python")]
98
use crate::types::JsonPython;
109

10+
#[cfg(feature = "rust_bridge")]
11+
use rust_bridge::{alias, alias_methods};
12+
1113
#[derive(Clone, Debug)]
1214
enum BindValue {
1315
String(String),
@@ -17,21 +19,25 @@ enum BindValue {
1719
Json(Json),
1820
}
1921

20-
#[derive(alias, Clone, Debug)]
22+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
23+
#[derive(Clone, Debug)]
2124
pub struct QueryRunner {
2225
query: String,
2326
bind_values: Vec<BindValue>,
2427
database_url: Option<String>,
2528
}
2629

27-
#[alias_methods(
28-
fetch_all,
29-
execute,
30-
bind_string,
31-
bind_int,
32-
bind_float,
33-
bind_bool,
34-
bind_json
30+
#[cfg_attr(
31+
feature = "rust_bridge",
32+
alias_methods(
33+
fetch_all,
34+
execute,
35+
bind_string,
36+
bind_int,
37+
bind_float,
38+
bind_bool,
39+
bind_json
40+
)
3541
)]
3642
impl QueryRunner {
3743
pub fn new(query: &str, database_url: Option<String>) -> Self {

0 commit comments

Comments
 (0)