Content-Length: 30716 | pFad | http://github.com/postgresml/postgresml/pull/1201.diff

thub.com diff --git a/pgml-sdks/pgml/Cargo.lock b/pgml-sdks/pgml/Cargo.lock index c9bff6f4d..b377c4167 100644 --- a/pgml-sdks/pgml/Cargo.lock +++ b/pgml-sdks/pgml/Cargo.lock @@ -60,6 +60,54 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ab91ebe16eb252986481c5b62f6098f3b698a45e34b5b98200cf20dd2484a44" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87" + +[[package]] +name = "anstyle-parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "317b9a89c1868f5ea6ff1d9539a69f45dffc21ce321ac1fd1160dfa48c8e2140" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b" +dependencies = [ + "windows-sys 0.48.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0699d10d2f4d628a98ee7b57b289abbc98ff3bad977cb3152709d4bf2330628" +dependencies = [ + "anstyle", + "windows-sys 0.48.0", +] + [[package]] name = "anyhow" version = "1.0.71" @@ -110,6 +158,12 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "bitflags" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "327762f6e5a765692301e5bb513e0d9fef63be86bbc14528052b1cd3e6f03e07" + [[package]] name = "block-buffer" version = "0.10.4" @@ -164,6 +218,63 @@ dependencies = [ "winapi", ] +[[package]] +name = "clap" +version = "4.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fffed7514f420abec6d183b1d3acfd9099c79c3a10a06ade4f8203f1411272" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63361bae7eef3771745f02d8d892bec2fee5f6e34af316ba556e7f97a7069ff1" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.28", +] + +[[package]] +name = "clap_lex" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1" + +[[package]] +name = "colorchoice" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" + +[[package]] +name = "colored" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2674ec482fbc38012cf31e6c42ba0177b431a0cb6f15fe40efa5aab1bda516f6" +dependencies = [ + "is-terminal", + "lazy_static", + "windows-sys 0.48.0", +] + [[package]] name = "console" version = "0.15.7" @@ -269,6 +380,31 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossterm" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64e6c0fbe2c17357405f7c758c1ef960fce08bdfb2c03d88d2a18d7e09c4b67" +dependencies = [ + "bitflags 1.3.2", + "crossterm_winapi", + "libc", + "mio", + "parking_lot 0.12.1", + "signal-hook", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -279,6 +415,16 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctrlc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a011bbe2c35ce9c1f143b7af6f94f29a167beb4cd1d29e6740ce836f723120e" +dependencies = [ + "nix", + "windows-sys 0.48.0", +] + [[package]] name = "darling" version = "0.14.4" @@ -351,6 +497,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "dyn-clone" +version = "1.0.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "545b22097d44f8a9581187cdf93de7a71e4722bf51200cfaba810865b49a495d" + [[package]] name = "either" version = "1.8.1" @@ -811,6 +963,22 @@ dependencies = [ "syn 2.0.28", ] +[[package]] +name = "inquire" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33e7c1ddeb15c9abcbfef6029d8e29f69b52b6d6c891031b88ed91b5065803b" +dependencies = [ + "bitflags 1.3.2", + "crossterm", + "dyn-clone", + "lazy_static", + "newline-converter", + "thiserror", + "unicode-segmentation", + "unicode-width", +] + [[package]] name = "instant" version = "0.1.12" @@ -837,6 +1005,17 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28b29a3cd74f0f4598934efe3aeba42bae0eb4680554128851ebbecb02af14e6" +[[package]] +name = "is-terminal" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b" +dependencies = [ + "hermit-abi 0.3.2", + "rustix 0.38.3", + "windows-sys 0.48.0", +] + [[package]] name = "itertools" version = "0.10.5" @@ -895,6 +1074,12 @@ version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef53942eb7bf7ff43a617b3e2c1c4a5ecf5944a7c1bc12d7ee39bbb15e5c1519" +[[package]] +name = "linux-raw-sys" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "969488b55f8ac402214f3f5fd243ebb7206cf82de60d3172994707a4bcc2b829" + [[package]] name = "lock_api" version = "0.4.10" @@ -997,6 +1182,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "927a765cd3fc26206e66b296465fa9d3e5ab003e651c1b3c060e7956d96b19d2" dependencies = [ "libc", + "log", "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.48.0", ] @@ -1060,6 +1246,26 @@ dependencies = [ "smallvec", ] +[[package]] +name = "newline-converter" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f71d09d5c87634207f894c6b31b6a2b2c64ea3bdcf71bd5599fdbbe1600c00f" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "nix" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "598beaf3cc6fdd9a5dfb1630c2800c7acd31df7aaf0f565796fba2b53ca1af1b" +dependencies = [ + "bitflags 1.3.2", + "cfg-if", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -1117,7 +1323,7 @@ version = "0.10.55" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "345df152bc43501c5eb9e4654ff05f794effb78d4efe3d53abc158baddc0703d" dependencies = [ - "bitflags", + "bitflags 1.3.2", "cfg-if", "foreign-types", "libc", @@ -1233,13 +1439,18 @@ checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" [[package]] name = "pgml" -version = "0.9.6" +version = "0.10.0" dependencies = [ "anyhow", "async-trait", "chrono", + "clap", + "colored", + "ctrlc", "futures", "indicatif", + "inquire", + "is-terminal", "itertools", "lopdf", "md5", @@ -1451,7 +1662,7 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" dependencies = [ - "bitflags", + "bitflags 1.3.2", ] [[package]] @@ -1460,7 +1671,7 @@ version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "567664f262709473930a4bf9e51bf2ebf3348f2e748ccc50dea20646858f8f29" dependencies = [ - "bitflags", + "bitflags 1.3.2", ] [[package]] @@ -1574,11 +1785,24 @@ version = "0.37.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "84f3f8f960ed3b5a59055428714943298bf3fa2d4a1d53135084e0544829d995" dependencies = [ - "bitflags", + "bitflags 1.3.2", "errno", "io-lifetimes", "libc", - "linux-raw-sys", + "linux-raw-sys 0.3.8", + "windows-sys 0.48.0", +] + +[[package]] +name = "rustix" +version = "0.38.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac5ffa1efe7548069688cd7028f32591853cd7b5b756d41bcffd2353e4fc75b4" +dependencies = [ + "bitflags 2.4.1", + "errno", + "libc", + "linux-raw-sys 0.4.11", "windows-sys 0.48.0", ] @@ -1697,7 +1921,7 @@ version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fc758eb7bffce5b308734e9b0c1468893cae9ff70ebf13e7090be8dcbcc83a8" dependencies = [ - "bitflags", + "bitflags 1.3.2", "core-foundation", "core-foundation-sys", "libc", @@ -1803,6 +2027,36 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "signal-hook" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29ad2e15f37ec9a6cc544097b78a1ec90001e9f71b81338ca39f430adaca99af" +dependencies = [ + "libc", + "mio", + "signal-hook", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + [[package]] name = "slab" version = "0.4.8" @@ -1864,7 +2118,7 @@ dependencies = [ "ahash 0.7.6", "atoi", "base64 0.13.1", - "bitflags", + "bitflags 1.3.2", "byteorder", "bytes", "crc", @@ -2012,7 +2266,7 @@ dependencies = [ "cfg-if", "fastrand", "redox_syscall 0.3.5", - "rustix", + "rustix 0.37.26", "windows-sys 0.48.0", ] @@ -2324,6 +2578,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "utf8parse" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" + [[package]] name = "uuid" version = "1.3.4" diff --git a/pgml-sdks/pgml/Cargo.toml b/pgml-sdks/pgml/Cargo.toml index 5db2bbefd..ac7dc23ed 100644 --- a/pgml-sdks/pgml/Cargo.toml +++ b/pgml-sdks/pgml/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pgml" -version = "0.9.6" +version = "0.10.0" edition = "2021" authors = ["PosgresML "] homepage = "https://postgresml.org/" @@ -38,6 +38,11 @@ serde = "1.0.181" futures = "0.3.28" walkdir = "2.4.0" lopdf = { version = "0.31.0", features = ["nom_parser"] } +clap = { version = "4", features = ["derive"]} +is-terminal = "0.4" +colored = "2" +ctrlc = "3" +inquire = "0.6" [features] default = [] diff --git a/pgml-sdks/pgml/javascript-cli/index.js b/pgml-sdks/pgml/javascript-cli/index.js new file mode 100755 index 000000000..165c4f768 --- /dev/null +++ b/pgml-sdks/pgml/javascript-cli/index.js @@ -0,0 +1,3 @@ +#!/usr/bin/env node +const pgml = require("pgml"); +pgml.cli().then(() => {}); diff --git a/pgml-sdks/pgml/javascript-cli/package-lock.json b/pgml-sdks/pgml/javascript-cli/package-lock.json new file mode 100644 index 000000000..650c18fc4 --- /dev/null +++ b/pgml-sdks/pgml/javascript-cli/package-lock.json @@ -0,0 +1,32 @@ +{ + "name": "pgml-cli", + "version": "0.10.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "pgml-cli", + "version": "0.10.0", + "license": "MIT", + "dependencies": { + "pgml": "file:../javascript" + }, + "bin": { + "pgml": "index.js" + } + }, + "../javascript": { + "name": "pgml", + "version": "0.9.6", + "license": "MIT", + "devDependencies": { + "@types/node": "^20.3.1", + "cargo-cp-artifact": "^0.1" + } + }, + "node_modules/pgml": { + "resolved": "../javascript", + "link": true + } + } +} diff --git a/pgml-sdks/pgml/javascript-cli/package.json b/pgml-sdks/pgml/javascript-cli/package.json new file mode 100644 index 000000000..791b52e6a --- /dev/null +++ b/pgml-sdks/pgml/javascript-cli/package.json @@ -0,0 +1,27 @@ +{ + "name": "pgml-cli", + "version": "0.10.0", + "description": "CLI for PostgresML, the GPU-powered AI application database.", + "keywords": [ + "postgres", + "machine learning", + "vector databases", + "embeddings" + ], + "bin": { + "pgml": "index.js" + }, + "author": { + "name": "PostgresML", + "email": "team@postgresml.org", + "url": "https://postgresml.org" + }, + "repository": { + "type": "git", + "url": "https://github.com/postgresml/postgresml" + }, + "license": "MIT", + "dependencies": { + "pgml": "0.10.0" + } +} diff --git a/pgml-sdks/pgml/javascript/package-lock.json b/pgml-sdks/pgml/javascript/package-lock.json index 08aedb865..9ab5f611e 100644 --- a/pgml-sdks/pgml/javascript/package-lock.json +++ b/pgml-sdks/pgml/javascript/package-lock.json @@ -1,12 +1,12 @@ { "name": "pgml", - "version": "0.8.1", + "version": "0.9.6", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "pgml", - "version": "0.8.1", + "version": "0.9.6", "license": "MIT", "devDependencies": { "@types/node": "^20.3.1", diff --git a/pgml-sdks/pgml/pyproject.toml b/pgml-sdks/pgml/pyproject.toml index df80ecb74..0d1843cbb 100644 --- a/pgml-sdks/pgml/pyproject.toml +++ b/pgml-sdks/pgml/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "maturin" [project] name = "pgml" requires-python = ">=3.7" -version = "0.9.6" +version = "0.10.0" description = "Python SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases." authors = [ {name = "PostgresML", email = "team@postgresml.org"}, diff --git a/pgml-sdks/pgml/python/pgml/__main__.py b/pgml-sdks/pgml/python/pgml/__main__.py new file mode 100644 index 000000000..9bbdae756 --- /dev/null +++ b/pgml-sdks/pgml/python/pgml/__main__.py @@ -0,0 +1,7 @@ +import asyncio +from pgml import cli + +async def main(): + await cli() + +asyncio.run(main()) diff --git a/pgml-sdks/pgml/src/builtins.rs b/pgml-sdks/pgml/src/builtins.rs index 188948c72..db023b951 100644 --- a/pgml-sdks/pgml/src/builtins.rs +++ b/pgml-sdks/pgml/src/builtins.rs @@ -101,7 +101,7 @@ mod tests { let query = "SELECT * from pgml.collections"; let results = builtins.query(query).fetch_all().await?; assert!(results.as_array().is_some()); - Ok(()) + Ok(()) } #[sqlx::test] diff --git a/pgml-sdks/pgml/src/cli.rs b/pgml-sdks/pgml/src/cli.rs new file mode 100644 index 000000000..359006c9f --- /dev/null +++ b/pgml-sdks/pgml/src/cli.rs @@ -0,0 +1,359 @@ +use clap::{Parser, Subcommand}; +use colored::Colorize; +use inquire::Text; +use is_terminal::IsTerminal; +use itertools::Itertools; +#[cfg(feature = "python")] +use pyo3::exceptions::PyRuntimeError; +#[cfg(feature = "python")] +use pyo3::prelude::*; +use sqlx::{Acquire, Executor}; +use std::io::Write; + +//github.com/ PostgresML CLI +#[cfg(feature = "python")] +#[derive(Parser, Debug, Clone)] +#[command(author, version, about, long_about = None, name = "pgml", bin_name = "pgml")] +struct Python { + //github.com/ We're running this as `python -m`, this argument is ignored + #[arg(short)] + module: Option, + + #[command(subcommand)] + subcommand: Subcommands, +} + +//github.com/ PostgresML CLI +#[cfg(feature = "javascript")] +#[derive(Parser, Debug, Clone)] +#[command(author, version, about, long_about = None, name = "pgml", bin_name = "pgml")] +struct Javascript { + //github.com/ Ignore this argument, we're running as `node`. + #[arg(name = "pgmlcli")] + pgmlcli: Option, + + #[command(subcommand)] + subcommand: Subcommands, +} + +#[derive(Subcommand, Debug, Clone)] +enum Subcommands { + //github.com/ Connect your PostgresML database to another PostgreSQL database. + Connect { + //github.com/ Name for this connection. Allows to configure multiple connections + //github.com/ from PostgresML to any number of databases. + #[arg(long)] + name: Option, + + //github.com/ Host name or IP address of your database. + //github.com/ The database must be reachable from our cloud via a private link + //github.com/ or the Internet. + #[arg(long)] + host: Option, + + //github.com/ The port on which the database server is running. + #[arg(long)] + port: Option, + + //github.com/ A user that has read permissions to your schemas and tables. + #[arg(long)] + user: Option, + + //github.com/ The password for the user. + #[arg(long)] + password: Option, + + //github.com/ The name of the Postgres database. + #[arg(long)] + database_name: Option, + + //github.com/ If you're using another schema that's not public, + //github.com/ you can specify it here. + #[arg(long)] + schema: Option, + + //github.com/ Don't do anything, just print the commands. + #[arg(long, default_value = "false")] + dry_run: bool, + + //github.com/ Drop the connection before creating it. + #[arg(long, default_value = "false")] + drop: bool, + + //github.com/ DATABASE_URL for your PostgresML database. + #[arg(long)] + database_url: Option, + }, +} + +enum Level { + Happy, + Sad, + #[allow(dead_code)] + Concerned, +} + +#[cfg(feature = "python")] +#[pyfunction] +pub fn cli(py: pyo3::Python) -> pyo3::PyResult<&pyo3::PyAny> { + ctrlc::set_handler(move || { + println!(""); + std::process::exit(1); + }) + .expect("failed to set ctrl-c handler"); + + pyo3_asyncio::tokio::future_into_py(py, async move { + match cli_internal().await { + Ok(_) => Ok(()), + Err(err) => Err(PyRuntimeError::new_err(format!("{}", err))), + } + }) +} + +#[cfg(feature = "javascript")] +pub fn cli( + mut cx: neon::context::FunctionContext, +) -> neon::result::JsResult { + ctrlc::set_handler(move || { + println!(""); + std::process::exit(1); + }) + .expect("failed to set ctrl-c handler"); + + use neon::prelude::*; + use rust_bridge::javascript::IntoJsResult; + let channel = cx.channel(); + let (deferred, promise) = cx.promise(); + deferred + .try_settle_with(&channel, move |mut cx| { + let runtime = crate::get_or_set_runtime(); + let x = runtime.block_on(cli_internal()); + let x = match x { + Ok(x) => x, + Err(e) => { + // Node has its own ctrl-c handler, so we need to handle it here. + if e.to_string() + .contains("Operation was interrupted by the user") + { + std::process::exit(1); + } else { + panic!("{e}"); + } + } + }; + x.into_js_result(&mut cx) + }) + .expect("Error sending js"); + Ok(promise) +} + +async fn cli_internal() -> anyhow::Result<()> { + #[cfg(feature = "python")] + let subcommand = { + let args = Python::parse(); + args.subcommand + }; + + #[cfg(feature = "javascript")] + let subcommand = { + let args = Javascript::parse(); + args.subcommand + }; + + match subcommand { + Subcommands::Connect { + name, + host, + port, + user, + password, + database_name, + dry_run, + schema, + drop, + database_url, + } => { + connect( + name, + host, + port, + user, + password, + database_name, + schema, + dry_run, + drop, + database_url, + ) + .await?; + } + }; + + Ok(()) +} + +async fn execute_sql(sql: &str) -> anyhow::Result<()> { + let pool = crate::get_or_initialize_pool(&None).await?; + let mut connection = pool.acquire().await?; + let mut transaction = connection.begin().await?; + + for query in sql.split(";") { + transaction.execute(query).await?; + } + + transaction.commit().await?; + + Ok(()) +} + +async fn connect( + name: Option, + host: Option, + port: Option, + user: Option, + password: Option, + database_name: Option, + schema: Option, + dry_run: bool, + drop: bool, + database_url: Option, +) -> anyhow::Result<()> { + println!(""); + println!("The connector will configure a Postgres Foreign Data Wrapper connection"); + println!("from PostgresML to your Postgres database of choice. If we're missing any details,"); + println!("we'll ask for them now."); + println!(""); + + if std::env::var("DATABASE_URL").is_err() && database_url.is_none() { + println!("Required DATABASE_URL environment variable is not set."); + println!("We need it to connect to your PostgresML database."); + println!(""); + let database_url = user_input!(None::, "DATABASE_URL"); + std::env::set_var("DATABASE_URL", database_url); + println!(""); + } else if let Some(database_url) = database_url { + std::env::set_var("DATABASE_URL", database_url); + } + + let name = user_input!(name, "Connection name", Some("production")); + let host = user_input!(host, "PostgreSQL host"); + let port = user_input!(port, "PostgreSQL port", Some("5432")); + let user = user_input!(user, "PostgreSQL user", Some("postgres")); + let password = user_input!(password, "Password"); + let database_name = user_input!(database_name, "PostgreSQL database", Some("postgres")); + let schema = user_input!(schema, "PostgreSQL schema", Some("public")); + + let sql = include_str!("sql/fdw.sql") + .replace("{host}", &host) + .replace("{port}", &port) + .replace("{user}", &user) + .replace("{password}", &password) + .replace("{database_name}", &database_name) + .replace("{db_name}", &name) + .replace("{schema}", &schema); + let drop_sql = include_str!("sql/fdw_drop.sql") + .replace("{db_name}", &name) + .replace("{schema}", &schema); + + if dry_run { + println!(""); + if drop { + println!("{}", syntax_highlight(&drop_sql)); + } + println!("{}", syntax_highlight(&sql)); + println!(""); + } else { + println!(""); + print!("Everything looks good, creating connection..."); + std::io::stdout().flush().unwrap(); + + if drop { + match execute_sql(&drop_sql).await { + Ok(_) => (), + Err(err) => { + println!("{}", colorize("error", Level::Sad)); + println!("{}", err); + std::process::exit(1); + } + }; + } + + match execute_sql(&sql).await { + Ok(_) => { + println!("{}", colorize("done", Level::Happy)); + println!(""); + println!("You can now use your PostgreSQL tables inside your PostgresML database."); + println!("If you connect with psql, you can view your tables by updating your search_path:"); + println!(""); + println!( + "{}", + syntax_highlight(&format!("SET search_path TO {}_public, public;", name)) + ); + println!(""); + } + Err(err) => { + println!("{}", colorize("error", Level::Sad)); + println!("{}", err); + } + }; + } + + Ok(()) +} + +fn syntax_highlight(text: &str) -> String { + if !std::io::stdout().is_terminal() { + return text.to_owned(); + } + + text.split(" ") + .into_iter() + .map(|word| { + let uppercase = word.chars().all(|c| c.is_ascii_uppercase()); + + if uppercase { + word.cyan().to_string() + } else { + word.to_owned() + } + }) + .join(" ") +} + +fn colorize(text: &str, level: Level) -> String { + if !std::io::stdout().is_terminal() { + return text.to_owned(); + } + + match level { + Level::Happy => text.green().to_string(), + Level::Sad => text.red().to_string(), + Level::Concerned => text.yellow().to_string(), + } +} + +macro_rules! user_input { + ($var:expr, $prompt:expr, $default:expr) => {{ + if $var.is_none() { + let prompt = format!("{}:", $prompt); + let prompt = if let Some(default) = $default { + Text::new(&prompt).with_default(default).prompt()? + } else { + Text::new(&prompt).prompt()? + }; + prompt.to_string() + } else { + $var.unwrap() + } + }}; + + ($var:expr, $prompt:expr) => {{ + user_input!($var, $prompt, None) + }}; + + ($var:expr) => {{ + user_input!($var, strginfy!($var)) + }}; +} + +use user_input; diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs index cd0eaaeef..e8056dd9f 100644 --- a/pgml-sdks/pgml/src/lib.rs +++ b/pgml-sdks/pgml/src/lib.rs @@ -4,7 +4,7 @@ //! //! With this SDK, you can seamlessly manage various database tables related to documents, text chunks, text splitters, LLM (Language Model) models, and embeddings. By leveraging the SDK's capabilities, you can efficiently index LLM embeddings using PgVector for fast and accurate queries. -use sqlx::PgPool; +use sqlx::{postgres::PgPoolOptions, PgPool}; use std::collections::HashMap; use std::env; use std::sync::RwLock; @@ -13,6 +13,7 @@ use tracing::Level; use tracing_subscriber::FmtSubscriber; mod builtins; +mod cli; mod collection; mod filter_builder; mod languages; @@ -60,7 +61,15 @@ async fn get_or_initialize_pool(database_url: &Option) -> anyhow::Result if let Some(pool) = pools.get(url) { Ok(pool.clone()) } else { - let pool = PgPool::connect_lazy(url)?; + let timeout = std::env::var("PGML_CHECKOUT_TIMEOUT") + .unwrap_or_else(|_| "5000".to_string()) + .parse::() + .expect("Error parsing PGML_CHECKOUT_TIMEOUT, expected an integer"); + + let pool = PgPoolOptions::new() + .acquire_timeout(std::time::Duration::from_millis(timeout)) + .connect_lazy(&url)?; + pools.insert(url.to_string(), pool.clone()); Ok(pool) } @@ -146,6 +155,7 @@ fn migrate(py: pyo3::Python) -> pyo3::PyResult<&pyo3::PyAny> { fn pgml(_py: pyo3::Python, m: &pyo3::types::PyModule) -> pyo3::PyResult<()> { m.add_function(pyo3::wrap_pyfunction!(init_logger, m)?)?; m.add_function(pyo3::wrap_pyfunction!(migrate, m)?)?; + m.add_function(pyo3::wrap_pyfunction!(cli::cli, m)?)?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -192,6 +202,7 @@ fn migrate( fn main(mut cx: neon::context::ModuleContext) -> neon::result::NeonResult<()> { cx.export_function("init_logger", init_logger)?; cx.export_function("migrate", migrate)?; + cx.export_function("cli", cli::cli)?; cx.export_function("newCollection", collection::CollectionJavascript::new)?; cx.export_function("newModel", model::ModelJavascript::new)?; cx.export_function("newSplitter", splitter::SplitterJavascript::new)?; diff --git a/pgml-sdks/pgml/src/sql/fdw.sql b/pgml-sdks/pgml/src/sql/fdw.sql new file mode 100644 index 000000000..75ae64f00 --- /dev/null +++ b/pgml-sdks/pgml/src/sql/fdw.sql @@ -0,0 +1,22 @@ + + CREATE SERVER "{db_name}" + FOREIGN DATA WRAPPER postgres_fdw + OPTIONS ( + host '{host}', + port '{port}', + dbname '{database_name}' + ); + + CREATE USER MAPPING + FOR CURRENT_USER + SERVER "{db_name}" + OPTIONS ( + user '{user}', + password '{password}' + ); + + CREATE SCHEMA "{db_name}_{schema}"; + + IMPORT FOREIGN SCHEMA "{schema}" + FROM SERVER "{db_name}" + INTO "{db_name}_{schema}"; diff --git a/pgml-sdks/pgml/src/sql/fdw_drop.sql b/pgml-sdks/pgml/src/sql/fdw_drop.sql new file mode 100644 index 000000000..8735489ee --- /dev/null +++ b/pgml-sdks/pgml/src/sql/fdw_drop.sql @@ -0,0 +1,8 @@ + + DROP SCHEMA IF EXISTS "{db_name}_{schema}" CASCADE; + + DROP USER MAPPING IF EXISTS + FOR CURRENT_USER + SERVER "{db_name}"; + + DROP SERVER IF EXISTS "{db_name}" CASCADE; diff --git a/pgml-sdks/pgml/src/utils.rs b/pgml-sdks/pgml/src/utils.rs index c8a2f8039..a8c040bc9 100644 --- a/pgml-sdks/pgml/src/utils.rs +++ b/pgml-sdks/pgml/src/utils.rs @@ -49,7 +49,9 @@ pub fn get_file_contents(path: &Path) -> anyhow::Result { "pdf" => { let doc = Document::load(path) .with_context(|| format!("Error reading PDF file: {}", path.display()))?; - doc.get_pages().into_keys().map(|page_number| { + doc.get_pages() + .into_keys() + .map(|page_number| { doc.extract_text(&[page_number]).with_context(|| { format!("Error extracting content from PDF file: {}", path.display()) })

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier! Saves Data!