mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-20 17:02:53 -04:00
Introduce elasticsearch database + tools
This commit is contained in:
Brandon Thomas
@@ -75,8 +75,3 @@ TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME=hifigan_superres_twilight_33000.pt
|
||||
|
||||
# Development key. Permissible to check in.
|
||||
NEWRELIC_API_KEY=c402be1c4c4216ec83d0cfc0851520cfFFFFNRAL
|
||||
|
||||
STORYTELLER_ROOT = "/home/tensor/code/storyteller" # your path here!
|
||||
STORYTELLER_FRONTEND = "storyteller-frontend"
|
||||
STORYTELLER_ML = "storyteller-ml"
|
||||
STORYTELLER_RUST = "storyteller-rust"
|
||||
+6
-4
@@ -2,18 +2,20 @@
|
||||
*secret*sh
|
||||
.env
|
||||
.env-aichatbot-secrets
|
||||
.env-elasticsearch-secrets
|
||||
.env-secrets
|
||||
_develop/tools/.secrets
|
||||
_develop/tools/db_secrets.toml
|
||||
crates/service/web/storyteller_web/config/storyteller-web.development-secrets.env
|
||||
db.secrets
|
||||
db_secrets.toml
|
||||
env-secrets
|
||||
production_database_secret_connect.sh
|
||||
production_database_secret_query.sh
|
||||
production_database_voice_clone_report_download.sh
|
||||
secrets.toml
|
||||
secrets.yaml
|
||||
_develop/tools/.secrets
|
||||
_develop/tools/db_secrets.toml
|
||||
twitch_secrets.toml
|
||||
crates/service/web/storyteller_web/config/storyteller-web.development-secrets.env
|
||||
|
||||
# New env secrets (per-app)
|
||||
*-secrets.env
|
||||
@@ -28,7 +30,7 @@ crates/lib/sqlite_queries/target/
|
||||
*.swp
|
||||
*~
|
||||
|
||||
# IntelliJ/CLion noise
|
||||
# IntelliJ/CLion/RustRover noise
|
||||
.idea/
|
||||
|
||||
# Python noise
|
||||
|
||||
Generated
+154
-4
@@ -766,6 +766,19 @@ dependencies = [
|
||||
"futures-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-compression"
|
||||
version = "0.3.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "942c7cd7ae39e91bde4820d74132e9862e62c2f386c3aa90ccf55949f5bad63a"
|
||||
dependencies = [
|
||||
"flate2",
|
||||
"futures-core",
|
||||
"memchr",
|
||||
"pin-project-lite 0.2.9",
|
||||
"tokio 1.25.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stripe"
|
||||
version = "0.15.0"
|
||||
@@ -910,6 +923,12 @@ version = "0.2.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a4521f3e3d031370679b3b140beb36dfe4801b09ac77e30c61941f97df3ef28b"
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b41b7ea54a0c9d92199de89e20e58d49f02f8e699814ef3fdf266f6f748d15c7"
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.12.3"
|
||||
@@ -1659,14 +1678,38 @@ dependencies = [
|
||||
"syn 2.0.26",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling"
|
||||
version = "0.13.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a01d95850c592940db9b8194bc39f4bc0e89dee5c4265e4b1807c34a9aba453c"
|
||||
dependencies = [
|
||||
"darling_core 0.13.4",
|
||||
"darling_macro 0.13.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling"
|
||||
version = "0.14.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0dd3cd20dc6b5a876612a6e5accfe7f3dd883db6d07acfbf14c128f61550dfa"
|
||||
dependencies = [
|
||||
"darling_core",
|
||||
"darling_macro",
|
||||
"darling_core 0.14.2",
|
||||
"darling_macro 0.14.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling_core"
|
||||
version = "0.13.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "859d65a907b6852c9361e3185c862aae7fafd2887876799fa55f5f99dc40d610"
|
||||
dependencies = [
|
||||
"fnv",
|
||||
"ident_case",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"strsim 0.10.0",
|
||||
"syn 1.0.107",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1683,13 +1726,24 @@ dependencies = [
|
||||
"syn 1.0.107",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling_macro"
|
||||
version = "0.13.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9c972679f83bdf9c42bd905396b6c3588a843a17f0f16dfcfa3e2c5d57441835"
|
||||
dependencies = [
|
||||
"darling_core 0.13.4",
|
||||
"quote",
|
||||
"syn 1.0.107",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling_macro"
|
||||
version = "0.14.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7618812407e9402654622dd402b0a89dff9ba93badd6540781526117b92aab7e"
|
||||
dependencies = [
|
||||
"darling_core",
|
||||
"darling_core 0.14.2",
|
||||
"quote",
|
||||
"syn 1.0.107",
|
||||
]
|
||||
@@ -1786,7 +1840,7 @@ version = "0.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c11bdc11a0c47bc7d37d582b5285da6849c96681023680b906673c5707af7b0f"
|
||||
dependencies = [
|
||||
"darling",
|
||||
"darling 0.14.2",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.107",
|
||||
@@ -1830,9 +1884,11 @@ name = "dev-database-seed"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"buckets",
|
||||
"clap",
|
||||
"cloud_storage",
|
||||
"config",
|
||||
"easyenv",
|
||||
"elasticsearch",
|
||||
"enums",
|
||||
"errors",
|
||||
"filesys",
|
||||
@@ -2088,6 +2144,12 @@ dependencies = [
|
||||
"toml",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dyn-clone"
|
||||
version = "1.0.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23d2f3407d9a573d666de4b5bdf10569d73ca9478087346697dcbae6244bfbcd"
|
||||
|
||||
[[package]]
|
||||
name = "easyenv"
|
||||
version = "0.1.0"
|
||||
@@ -2113,6 +2175,69 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "elasticsearch"
|
||||
version = "8.5.0-alpha.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "40d9bd57d914cc66ce878f098f63ed7b5d5b64c30644a5adb950b008f874a6c6"
|
||||
dependencies = [
|
||||
"base64 0.11.0",
|
||||
"bytes 1.3.0",
|
||||
"dyn-clone",
|
||||
"lazy_static",
|
||||
"percent-encoding",
|
||||
"reqwest",
|
||||
"rustc_version",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"url",
|
||||
"void",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "elasticsearch-cli"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"buckets",
|
||||
"clap",
|
||||
"cloud_storage",
|
||||
"config",
|
||||
"easyenv",
|
||||
"elasticsearch",
|
||||
"elasticsearch_schema",
|
||||
"enums",
|
||||
"errors",
|
||||
"filesys",
|
||||
"hashing",
|
||||
"idempotency",
|
||||
"log",
|
||||
"mimetypes",
|
||||
"mysql_queries",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"sqlx",
|
||||
"storyteller_root",
|
||||
"strum",
|
||||
"tokens",
|
||||
"tokio 1.25.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "elasticsearch_schema"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"elasticsearch",
|
||||
"enums",
|
||||
"errors",
|
||||
"log",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokens",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "encoding_rs"
|
||||
version = "0.8.28"
|
||||
@@ -4642,6 +4767,7 @@ version = "0.11.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "21eed90ec8570952d53b772ecf8f206aa1ec9a3d76b2521c56c42973f2d91ee9"
|
||||
dependencies = [
|
||||
"async-compression",
|
||||
"base64 0.21.0",
|
||||
"bytes 1.3.0",
|
||||
"encoding_rs",
|
||||
@@ -5113,6 +5239,28 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_with"
|
||||
version = "1.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "678b5a069e50bf00ecd22d0cd8ddf7c236f68581b03db652061ed5eb13a312ff"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_with_macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_with_macros"
|
||||
version = "1.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e182d6ec6f05393cc0e5ed1bf81ad6db3a8feedf8ee515ecdd369809bcce8082"
|
||||
dependencies = [
|
||||
"darling 0.13.4",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.107",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serial_test"
|
||||
version = "2.0.0"
|
||||
@@ -5554,6 +5702,8 @@ dependencies = [
|
||||
"diesel",
|
||||
"dotenv",
|
||||
"easyenv",
|
||||
"elasticsearch",
|
||||
"elasticsearch_schema",
|
||||
"enums",
|
||||
"errors",
|
||||
"futures 0.1.31",
|
||||
|
||||
@@ -13,6 +13,7 @@ members = [
|
||||
"crates/cli/database_cleanup",
|
||||
"crates/cli/dev_database_seed",
|
||||
"crates/cli/dev_upload_media_file",
|
||||
"crates/cli/elasticsearch-cli",
|
||||
"crates/cli/kubectl_pod_cleanup",
|
||||
"crates/cli/scrape_test",
|
||||
|
||||
@@ -32,6 +33,7 @@ members = [
|
||||
"crates/lib/deprecated/container_common", # TODO(bt, 2023-02-13): Pull apart and rehome pieces
|
||||
"crates/lib/deprecated/http_server_common", # TODO(bt, 2023-03-20): Move useful pieces to actix_helpers
|
||||
"crates/lib/easyenv", # NB: Vendored version of https://github.com/echelon/easyenv.rs that we'll improve.
|
||||
"crates/lib/elasticsearch_schema",
|
||||
"crates/lib/errors",
|
||||
"crates/lib/files/archives",
|
||||
"crates/lib/files/filesys",
|
||||
|
||||
@@ -143,13 +143,76 @@ file (located at `/etc/hosts`) to include the following configuration lines:
|
||||
127.0.0.1 devproxy.storyteller.ai
|
||||
```
|
||||
|
||||
### (9a) Install Redis
|
||||
### (9) Install Redis
|
||||
|
||||
```
|
||||
sudo apt install redis
|
||||
```
|
||||
|
||||
### (9b) Run one or more applications:
|
||||
### (10) Install Elasticsearch
|
||||
|
||||
Mac: (TODO - haven't installed yet)
|
||||
|
||||
Ubuntu 22.04:
|
||||
|
||||
Install:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://artifacts.elastic.co/GPG-KEY-elasticsearch | sudo gpg --dearmor -o /usr/share/keyrings/elastic.gpg
|
||||
|
||||
echo "deb [signed-by=/usr/share/keyrings/elastic.gpg] https://artifacts.elastic.co/packages/7.x/apt stable main" | sudo tee -a /etc/apt/sources.list.d/elastic-7.x.list
|
||||
|
||||
sudo apt update
|
||||
sudo apt install elasticsearch
|
||||
```
|
||||
|
||||
Configure:
|
||||
|
||||
```text
|
||||
# Edit /etc/elasticsearch/elasticsearch.yml
|
||||
# Set the following:
|
||||
network.host: localhost
|
||||
```
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
sudo systemctl start elasticsearch
|
||||
sudo systemctl enable elasticsearch
|
||||
```
|
||||
|
||||
Verify:
|
||||
|
||||
```bash
|
||||
curl -X GET 'http://localhost:9200'
|
||||
```
|
||||
|
||||
Install Kibana (prototyping UI):
|
||||
|
||||
```bash
|
||||
# Install
|
||||
sudo apt-get install kibana
|
||||
|
||||
# Run
|
||||
sudo /bin/systemctl daemon-reload
|
||||
sudo /bin/systemctl enable kibana.service
|
||||
sudo systemctl start kibana.service
|
||||
|
||||
# Access the Kibana Status Page:
|
||||
# http://localhost:5601/status
|
||||
|
||||
# Access the query console UI:
|
||||
# Click Sidebar > Management > Dev Tools > Console
|
||||
# http://localhost:5601/
|
||||
```
|
||||
|
||||
Extra reading:
|
||||
|
||||
- [Guide 1](https://www.digitalocean.com/community/tutorials/how-to-install-and-configure-elasticsearch-on-ubuntu-22-04)
|
||||
- [Guide 2](https://www.elastic.co/guide/en/elasticsearch/reference/current/deb.html)
|
||||
- [Kibana](https://www.elastic.co/guide/en/kibana/current/deb.html)
|
||||
|
||||
### (11) Run one or more applications:
|
||||
DO NOT CHECK THESE TWO FILES INTO THE REPO
|
||||
Ask Brandon for the .env-secrets and .storteller-web.development-secrets.env and place it in the root of the folder.
|
||||
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"settings" : {
|
||||
"index.number_of_shards" : 3,
|
||||
"index.number_of_replicas" : 2
|
||||
},
|
||||
"mappings" : {
|
||||
"properties" : {
|
||||
"token" : { "type" : "keyword" },
|
||||
|
||||
"title" : { "type" : "search_as_you_type" },
|
||||
|
||||
"ietf_language_tag" : { "type" : "keyword" },
|
||||
"ietf_primary_language_subtag" : { "type" : "keyword" },
|
||||
|
||||
"creator_username" : { "type" : "keyword" },
|
||||
"creator_display_name" : { "type" : "keyword" },
|
||||
"creator_user_token" : { "type" : "keyword" },
|
||||
"creator_set_visibility" : { "type" : "keyword" },
|
||||
|
||||
"created_at" : { "type" : "date" },
|
||||
"updated_at" : { "type" : "date" }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -30,6 +30,7 @@ tokens = { path = "../../_types/tokens" }
|
||||
|
||||
# External
|
||||
log = "0.4.14"
|
||||
clap = "2.34.0" # TODO: This is really out of date, but has conflicts when updated
|
||||
|
||||
# Runtime
|
||||
# NB: We switched to Tokio from `async-std` for broader ecosystem support.
|
||||
@@ -41,5 +42,7 @@ tokio = { version = "1.6.0", features = ["macros"] }
|
||||
# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828
|
||||
sqlx = { version = "0.6.2", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] }
|
||||
|
||||
elasticsearch = "8.5.0-alpha.1"
|
||||
|
||||
[dev-dependencies]
|
||||
# None yet
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
use std::path::PathBuf;
|
||||
use clap::{App, Arg};
|
||||
use errors::{anyhow, AnyhowResult};
|
||||
|
||||
pub struct CliArgs {
|
||||
pub seed_cloud_bucket: bool,
|
||||
pub seed_elasticsearch: bool,
|
||||
}
|
||||
|
||||
pub fn parse_cli_args() -> AnyhowResult<CliArgs> {
|
||||
let matches = App::new("dev-database-seed")
|
||||
.arg(Arg::with_name("seed_cloud_bucket")
|
||||
.long("bucket")
|
||||
.help("Seed the cloud bucket with files")
|
||||
.takes_value(false)
|
||||
.required(false))
|
||||
.arg(Arg::with_name("seed_elasticsearch")
|
||||
.long("elasticsearch")
|
||||
.help("Seed the local elasticsearch")
|
||||
.takes_value(false)
|
||||
.required(false))
|
||||
.get_matches();
|
||||
|
||||
Ok(CliArgs {
|
||||
seed_cloud_bucket: matches.is_present("seed_cloud_bucket"),
|
||||
seed_elasticsearch: matches.is_present("seed_elasticsearch"),
|
||||
})
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::time::Duration;
|
||||
use elasticsearch::Elasticsearch;
|
||||
|
||||
use log::info;
|
||||
use sqlx::mysql::MySqlPoolOptions;
|
||||
@@ -6,11 +7,13 @@ use sqlx::mysql::MySqlPoolOptions;
|
||||
use cloud_storage::bucket_client::BucketClient;
|
||||
use config::shared_constants::{DEFAULT_MYSQL_CONNECTION_STRING, DEFAULT_RUST_LOG};
|
||||
use errors::AnyhowResult;
|
||||
use crate::cli_args::parse_cli_args;
|
||||
|
||||
use crate::seeding::users::seed_user_accounts;
|
||||
use crate::seeding::voice_conversion::seed_voice_conversion;
|
||||
use crate::seeding::zero_shot_tts::seed_zero_shot_tts;
|
||||
|
||||
pub mod cli_args;
|
||||
pub mod seeding;
|
||||
|
||||
#[tokio::main]
|
||||
@@ -20,7 +23,7 @@ pub async fn main() -> AnyhowResult<()> {
|
||||
info!("Database seed CLI script.");
|
||||
|
||||
// NB: Read secrets (eg. ACCESS_KEY)
|
||||
easyenv::from_filename("env-secrets")?;
|
||||
easyenv::from_filename(".env-secrets")?;
|
||||
|
||||
let db_connection_string =
|
||||
easyenv::get_env_string_or_default(
|
||||
@@ -32,10 +35,22 @@ pub async fn main() -> AnyhowResult<()> {
|
||||
.connect(&db_connection_string)
|
||||
.await?;
|
||||
|
||||
let public_bucket_client = get_bucket_client()?;
|
||||
let args = parse_cli_args()?;
|
||||
|
||||
let mut maybe_public_bucket_client = None;
|
||||
|
||||
if args.seed_cloud_bucket {
|
||||
maybe_public_bucket_client = Some(get_bucket_client()?);
|
||||
}
|
||||
|
||||
let mut maybe_elasticsearch = None;
|
||||
|
||||
if args.seed_elasticsearch {
|
||||
maybe_elasticsearch = Some(get_elasticsearch_client());
|
||||
}
|
||||
|
||||
seed_user_accounts(&pool).await?;
|
||||
seed_zero_shot_tts(&pool, &public_bucket_client).await?;
|
||||
seed_zero_shot_tts(&pool, maybe_public_bucket_client.as_ref()).await?;
|
||||
seed_voice_conversion(&pool).await?;
|
||||
|
||||
info!("Done!");
|
||||
@@ -64,3 +79,9 @@ fn get_bucket_client() -> AnyhowResult<BucketClient> {
|
||||
|
||||
Ok(public_bucket_client)
|
||||
}
|
||||
|
||||
fn get_elasticsearch_client() -> AnyhowResult<Elasticsearch> {
|
||||
// TODO(bt,2023-10-26): Allow connecting to instances by URL instead of the default dev URL.
|
||||
let client = Elasticsearch::default();
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
@@ -24,7 +24,7 @@ use tokens::tokens::zs_voices::ZsVoiceToken;
|
||||
|
||||
use crate::seeding::users::HANASHI_USERNAME;
|
||||
|
||||
pub async fn seed_zero_shot_tts(mysql_pool: &Pool<MySql>, bucket_client: &BucketClient) -> AnyhowResult<()> {
|
||||
pub async fn seed_zero_shot_tts(mysql_pool: &Pool<MySql>, maybe_bucket_client: Option<&BucketClient>) -> AnyhowResult<()> {
|
||||
info!("Seeding zero shot TTS...");
|
||||
|
||||
let user_token = match get_user_token_by_username(HANASHI_USERNAME, mysql_pool).await? {
|
||||
@@ -44,7 +44,7 @@ pub async fn seed_zero_shot_tts(mysql_pool: &Pool<MySql>, bucket_client: &Bucket
|
||||
];
|
||||
|
||||
for (voice_name, bucket_hash, wav_file, user_token) in records {
|
||||
create_voice_records(voice_name, bucket_hash, wav_file, user_token, mysql_pool, bucket_client).await?;
|
||||
create_voice_records(voice_name, bucket_hash, wav_file, user_token, mysql_pool, maybe_bucket_client).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -56,7 +56,7 @@ async fn create_voice_records(
|
||||
wav_file: &str,
|
||||
creator_user_token: &UserToken,
|
||||
mysql_pool: &Pool<MySql>,
|
||||
bucket_client: &BucketClient,
|
||||
maybe_bucket_client: Option<&BucketClient>,
|
||||
) -> AnyhowResult<(ZsVoiceToken, ZsVoiceDatasetToken)> {
|
||||
info!("Creating voice records for voice {} ...", voice_name);
|
||||
|
||||
@@ -72,28 +72,13 @@ async fn create_voice_records(
|
||||
mysql_pool,
|
||||
}).await?;
|
||||
|
||||
info!("Uploading wav file {} ...", wav_file);
|
||||
|
||||
let public_upload_path = MediaFileBucketPath::generate_new(
|
||||
Some("dev_zs_"),
|
||||
Some(".bin"));
|
||||
|
||||
info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str());
|
||||
|
||||
let mut file_path = get_storyteller_rust_root();
|
||||
file_path.push(format!("assets/seed/{}", wav_file));
|
||||
|
||||
info!("Reading seed file: {:?}", file_path);
|
||||
|
||||
let bytes = file_read_bytes(file_path)?;
|
||||
let mimetype = get_mimetype_for_bytes(&bytes).unwrap_or("audio/wav");
|
||||
|
||||
let _r = bucket_client.upload_file_with_content_type(
|
||||
public_upload_path.get_full_object_path_str(),
|
||||
bytes.as_ref(),
|
||||
mimetype)
|
||||
.await?;
|
||||
let public_upload_path;
|
||||
|
||||
if let Some(bucket_client) = maybe_bucket_client {
|
||||
public_upload_path = seed_file_to_bucket(wav_file, bucket_client).await?;
|
||||
} else {
|
||||
public_upload_path = MediaFileBucketPath::from_object_hash("fake", None, None);
|
||||
}
|
||||
|
||||
info!("Creating dataset sample record...");
|
||||
|
||||
@@ -132,3 +117,29 @@ async fn create_voice_records(
|
||||
|
||||
Ok((voice_token, dataset_token))
|
||||
}
|
||||
|
||||
async fn seed_file_to_bucket(wav_file: &str, bucket_client: &BucketClient) -> AnyhowResult<MediaFileBucketPath> {
|
||||
info!("Uploading wav file {} ...", wav_file);
|
||||
|
||||
let public_upload_path = MediaFileBucketPath::generate_new(
|
||||
Some("dev_zs_"),
|
||||
Some(".bin"));
|
||||
|
||||
info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str());
|
||||
|
||||
let mut file_path = get_storyteller_rust_root();
|
||||
file_path.push(format!("assets/seed/{}", wav_file));
|
||||
|
||||
info!("Reading seed file: {:?}", file_path);
|
||||
|
||||
let bytes = file_read_bytes(file_path)?;
|
||||
let mimetype = get_mimetype_for_bytes(&bytes).unwrap_or("audio/wav");
|
||||
|
||||
let _r = bucket_client.upload_file_with_content_type(
|
||||
public_upload_path.get_full_object_path_str(),
|
||||
bytes.as_ref(),
|
||||
mimetype)
|
||||
.await?;
|
||||
|
||||
Ok(public_upload_path)
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ pub struct CliArgs {
|
||||
}
|
||||
|
||||
pub fn parse_cli_args() -> AnyhowResult<CliArgs> {
|
||||
let matches = App::new("scrape_test")
|
||||
let matches = App::new("dev_upload_media")
|
||||
.arg(Arg::with_name("file")
|
||||
.short("f")
|
||||
.long("file")
|
||||
@@ -15,18 +15,11 @@ pub fn parse_cli_args() -> AnyhowResult<CliArgs> {
|
||||
.help("File to upload")
|
||||
.takes_value(true)
|
||||
.required(true))
|
||||
//.arg(Arg::with_name("print_html")
|
||||
// .long("html")
|
||||
// .help("Print the HTML scraped")
|
||||
// .takes_value(false)
|
||||
// .required(false))
|
||||
.get_matches();
|
||||
|
||||
let file_path = matches.value_of("file")
|
||||
.map(|val| val.trim().to_string())
|
||||
.ok_or(anyhow!("no url supplied"))?;
|
||||
|
||||
//let print_html = matches.is_present("print_html");
|
||||
.ok_or(anyhow!("no file path supplied"))?;
|
||||
|
||||
Ok(CliArgs {
|
||||
file_path: PathBuf::from(file_path),
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
[package]
|
||||
name = "elasticsearch-cli"
|
||||
edition = "2018"
|
||||
version = "0.0.1"
|
||||
authors = [
|
||||
"Brandon Thomas <bt@brand.io>",
|
||||
"Brandon Thomas <echelon@gmail.com>",
|
||||
]
|
||||
publish = false
|
||||
|
||||
[[bin]]
|
||||
name = "elasticsearch-cli"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
# Internal
|
||||
buckets = { path = "../../_types/buckets" }
|
||||
cloud_storage = { path = "../../lib/api_clients/cloud_storage" }
|
||||
config = { path = "../../lib/deprecated/config" }
|
||||
easyenv = { path = "../../lib/easyenv" }
|
||||
elasticsearch_schema = { path = "../../lib/elasticsearch_schema" }
|
||||
enums = { path = "../../_types/enums" }
|
||||
errors = { path = "../../lib/errors" }
|
||||
filesys = { path = "../../lib/files/filesys" }
|
||||
hashing = { path = "../../lib/files/hashing" }
|
||||
idempotency = { path = "../../lib/idempotency" }
|
||||
mimetypes = { path = "../../lib/files/mimetypes" }
|
||||
mysql_queries = { path = "../../lib/mysql_queries" }
|
||||
storyteller_root = { path = "../../lib/storyteller_root" }
|
||||
tokens = { path = "../../_types/tokens" }
|
||||
|
||||
# External
|
||||
log = "0.4.14"
|
||||
clap = "2.34.0" # TODO: This is really out of date, but has conflicts when updated
|
||||
strum = { version = "0.24", features = ["derive"] }
|
||||
|
||||
# Runtime
|
||||
# NB: We switched to Tokio from `async-std` for broader ecosystem support.
|
||||
tokio = { version = "1.6.0", features = ["macros"] }
|
||||
|
||||
# MySql Database
|
||||
# "offline" supports typesafe builds without a live database
|
||||
# "rustls" doesn't appear to work with IP addresses, per:
|
||||
# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828
|
||||
sqlx = { version = "0.6.2", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] }
|
||||
|
||||
elasticsearch = "8.5.0-alpha.1"
|
||||
|
||||
# Serialization
|
||||
serde = "1.0.125"
|
||||
serde_derive = "1.0.125"
|
||||
serde_json = "1.0.64"
|
||||
|
||||
[dev-dependencies]
|
||||
# None yet
|
||||
@@ -0,0 +1,84 @@
|
||||
use std::str::FromStr;
|
||||
|
||||
use clap::{App, Arg, ArgMatches};
|
||||
use strum::{EnumCount, EnumString, IntoEnumIterator};
|
||||
use strum::EnumIter;
|
||||
|
||||
use errors::{anyhow, AnyhowResult, bail};
|
||||
|
||||
pub struct CliArgs {
|
||||
pub mysql_environment: Environment,
|
||||
pub elasticsearch_environment: Environment,
|
||||
pub action: Action,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum Environment {
|
||||
Development,
|
||||
Production,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, EnumIter, EnumCount, EnumString, strum::Display)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum Action {
|
||||
ReindexTts,
|
||||
SearchTts,
|
||||
}
|
||||
|
||||
pub fn parse_cli_args() -> AnyhowResult<CliArgs> {
|
||||
let matches = App::new("dev-database-seed")
|
||||
.arg(Arg::with_name("action")
|
||||
.long("action")
|
||||
.help("action to take")
|
||||
.takes_value(true)
|
||||
.required(true))
|
||||
.arg(Arg::with_name("mysql")
|
||||
.long("mysql")
|
||||
.help("Production or development")
|
||||
.takes_value(true)
|
||||
.required(false))
|
||||
.arg(Arg::with_name("elasticsearch")
|
||||
.long("elasticsearch")
|
||||
.help("Production or development")
|
||||
.takes_value(true)
|
||||
.required(false))
|
||||
.get_matches();
|
||||
|
||||
Ok(CliArgs {
|
||||
mysql_environment: to_environment(&matches, "mysql")?,
|
||||
elasticsearch_environment: to_environment(&matches, "elasticsearch")?,
|
||||
action: get_action(&matches)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn to_environment(matches: &ArgMatches, key: &str) -> AnyhowResult<Environment> {
|
||||
let value = matches.value_of(key)
|
||||
.map(|s| s.to_lowercase());
|
||||
|
||||
Ok(match value.as_deref() {
|
||||
None => Environment::Development,
|
||||
Some("dev") | Some("development") => Environment::Development,
|
||||
Some("prod") | Some("production") => Environment::Production,
|
||||
_ => bail!("invalid environment: {:?}", value),
|
||||
})
|
||||
}
|
||||
|
||||
fn get_action(matches: &ArgMatches) -> AnyhowResult<Action> {
|
||||
let action = matches.value_of("action")
|
||||
.map(|s| action_from_str(s))
|
||||
.transpose()?
|
||||
.ok_or(anyhow!("no action provided"))?;
|
||||
|
||||
Ok(action)
|
||||
}
|
||||
|
||||
fn action_from_str(value: &str) -> AnyhowResult<Action> {
|
||||
let action = Action::from_str(value)
|
||||
.map_err(|err| {
|
||||
let choices = Action::iter()
|
||||
.map(|e| e.to_string())
|
||||
.collect::<Vec<_>>();
|
||||
anyhow!("parse error: {:?}, provided: \"{}\" choices: {:?}", err, value, choices)
|
||||
})?;
|
||||
Ok(action)
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
//! elasticsearch-cli
|
||||
//!
|
||||
//! This is not a dev seed tool, but rather an operational tool meant to interface with both
|
||||
//! development and production.
|
||||
//!
|
||||
//! The intent is to be able to quickly populate documents and indices.
|
||||
//!
|
||||
use elasticsearch::Elasticsearch;
|
||||
use elasticsearch::http::transport::Transport;
|
||||
use log::info;
|
||||
use sqlx::{MySql, Pool};
|
||||
use sqlx::mysql::MySqlPoolOptions;
|
||||
|
||||
use config::shared_constants::DEFAULT_RUST_LOG;
|
||||
use elasticsearch_schema::searches::search_tts_models::search_tts_models;
|
||||
use errors::AnyhowResult;
|
||||
|
||||
use crate::cli_args::{Action, Environment, parse_cli_args};
|
||||
use crate::plans::create_all_tts_documents::create_all_tts_documents;
|
||||
|
||||
pub mod cli_args;
|
||||
pub mod plans;
|
||||
|
||||
#[tokio::main]
|
||||
pub async fn main() -> AnyhowResult<()> {
|
||||
println!("elasticsearch-cli: operational tooling for Elasticsearch");
|
||||
|
||||
easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG));
|
||||
|
||||
// NB: This secrets file differs from the rest because we might actually want to cross
|
||||
// development/production boundaries for seeding, or scope to production for rebuilding indices.
|
||||
// We don't want to pull in secrets from other sources. (Hopefully this isn't getting out of
|
||||
// hand at this point.)
|
||||
easyenv::from_filename(".env-elasticsearch-secrets")?;
|
||||
|
||||
let args = parse_cli_args()?;
|
||||
|
||||
let mysql= get_mysql(args.mysql_environment).await?;
|
||||
let elasticsearch = get_elasticsearch_client(args.elasticsearch_environment)?;
|
||||
|
||||
match args.action {
|
||||
Action::ReindexTts => {
|
||||
info!("Reindexing TTS...");
|
||||
create_all_tts_documents(&mysql, &elasticsearch).await?;
|
||||
}
|
||||
Action::SearchTts => {
|
||||
info!("Searching TTS...");
|
||||
search_tts_models(&elasticsearch, "zel", Some("en")).await;
|
||||
}
|
||||
}
|
||||
|
||||
info!("Done!");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_mysql(environment: Environment) -> AnyhowResult<Pool<MySql>> {
|
||||
info!("Connecting to {:?} MySQL...", environment);
|
||||
|
||||
let connection_string_env = match environment {
|
||||
Environment::Development => "MYSQL_DEVELOPMENT_URL",
|
||||
Environment::Production => "MYSQL_PRODUCTION_URL",
|
||||
};
|
||||
|
||||
let pool = MySqlPoolOptions::new()
|
||||
.max_connections(easyenv::get_env_num("MYSQL_MAX_CONNECTIONS", 3)?)
|
||||
.connect(&easyenv::get_env_string_required(connection_string_env)?)
|
||||
.await?;
|
||||
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
fn get_elasticsearch_client(environment: Environment) -> AnyhowResult<Elasticsearch> {
|
||||
info!("Connecting to {:?} Elasticsearch...", environment);
|
||||
|
||||
let connection_string_env = match environment {
|
||||
Environment::Development => "ELASTICSEARCH_DEVELOPMENT_URL",
|
||||
Environment::Production => "ELASTICSEARCH_PRODUCTION_URL",
|
||||
};
|
||||
|
||||
let transport = Transport::single_node(&easyenv::get_env_string_required(connection_string_env)?)?;
|
||||
|
||||
// TODO(bt,2023-10-26): Allow connecting to instances by URL instead of the default dev URL.
|
||||
let client = Elasticsearch::new(transport);
|
||||
Ok(client)
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
use std::fs::read_to_string;
|
||||
|
||||
use elasticsearch::{BulkOperation, BulkParts, Elasticsearch};
|
||||
use log::{error, info};
|
||||
use serde_json::Value;
|
||||
use sqlx::{MySql, Pool};
|
||||
|
||||
use elasticsearch_schema::documents::tts_model_document::{TTS_MODEL_INDEX, TtsModelDocument};
|
||||
use elasticsearch_schema::traits::document::Document;
|
||||
use elasticsearch_schema::utils::create_index_if_not_exists::{create_index_if_not_exists, CreateIndexArgs};
|
||||
use errors::AnyhowResult;
|
||||
use mysql_queries::queries::tts::tts_models::list_tts_models::{list_tts_models, TtsModelRecordForList};
|
||||
use storyteller_root::get_storyteller_rust_root;
|
||||
use tokens::tokens::tts_models::TtsModelToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
pub async fn create_all_tts_documents(mysql: &Pool<MySql>, elasticsearch: &Elasticsearch) -> AnyhowResult<()> {
|
||||
info!("Create all TTS documents.");
|
||||
|
||||
const SCOPE_USERNAME : Option<&str> = None;
|
||||
const REQUIRE_MOD_APPROVED : bool = false;
|
||||
|
||||
// TODO(bt,2023-10-26): Paginate this query (!!)
|
||||
let all_tts_models = list_tts_models(
|
||||
mysql,
|
||||
SCOPE_USERNAME,
|
||||
REQUIRE_MOD_APPROVED
|
||||
).await?;
|
||||
|
||||
create_tts_model_index(&elasticsearch, true).await?;
|
||||
|
||||
for record in all_tts_models {
|
||||
create_document_from_record(elasticsearch, &record).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_document_from_record(elasticsearch: &Elasticsearch, record: &TtsModelRecordForList) -> AnyhowResult<()> {
|
||||
info!("Create record for {:?} - {:?}", record.model_token, record.title);
|
||||
|
||||
let document = TtsModelDocument {
|
||||
token: TtsModelToken::new_from_str(&record.model_token),
|
||||
title: record.title.clone(),
|
||||
ietf_language_tag: record.ietf_language_tag.clone(),
|
||||
ietf_primary_language_subtag: record.ietf_primary_language_subtag.clone(),
|
||||
creator_username: record.creator_username.clone(),
|
||||
creator_display_name: record.creator_display_name.clone(),
|
||||
creator_user_token: UserToken::new_from_str(&record.creator_user_token),
|
||||
creator_set_visibility: record.creator_set_visibility,
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
};
|
||||
|
||||
let op : BulkOperation<_> = BulkOperation::index(&document)
|
||||
.id(document.get_document_id())
|
||||
.into();
|
||||
|
||||
let response = elasticsearch
|
||||
.bulk(BulkParts::Index(TTS_MODEL_INDEX))
|
||||
.body(vec![op])
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let json: Value = response.json().await?;
|
||||
|
||||
let had_errors = json["errors"].as_bool().unwrap_or(false);
|
||||
if had_errors {
|
||||
let failed: Vec<&Value> = json["items"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|v| !v["error"].is_null())
|
||||
.collect();
|
||||
|
||||
// TODO: retry failures
|
||||
error!("Errors during indexing. Failures: {}", failed.len());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// NB: Adapted from elasticsearch crate examples source
|
||||
async fn create_tts_model_index(client: &Elasticsearch, delete: bool) -> AnyhowResult<()> {
|
||||
|
||||
info!("Creating TTS model index...");
|
||||
|
||||
let index_path = get_storyteller_rust_root()
|
||||
.join("_elasticsearch/index_definitions/tts_models_v1.json");
|
||||
|
||||
info!("Reading index file: {:?}", index_path);
|
||||
|
||||
let index_definition = read_to_string(index_path)?;
|
||||
|
||||
create_index_if_not_exists(CreateIndexArgs {
|
||||
client,
|
||||
index_name: TTS_MODEL_INDEX,
|
||||
index_definition: &index_definition,
|
||||
delete_existing: true,
|
||||
}).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod create_all_tts_documents;
|
||||
@@ -0,0 +1,31 @@
|
||||
[package]
|
||||
name = "elasticsearch_schema"
|
||||
edition = "2018"
|
||||
version = "0.0.1"
|
||||
authors = [
|
||||
"Brandon Thomas <bt@brand.io>",
|
||||
"Brandon Thomas <echelon@gmail.com>",
|
||||
]
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
name = "elasticsearch_schema"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
# Internal
|
||||
enums = { path = "../../_types/enums" }
|
||||
errors = { path = "../errors" }
|
||||
tokens = { path = "../../_types/tokens" }
|
||||
|
||||
# External
|
||||
elasticsearch = "8.5.0-alpha.1"
|
||||
log = "0.4.14"
|
||||
|
||||
# Serialization
|
||||
chrono = { version = "0.4.19", features = ["serde"] }
|
||||
serde = { version = "1.0.152", features = ["derive"] }
|
||||
serde_json = "1.0.64"
|
||||
|
||||
[dev-dependencies]
|
||||
# None yet
|
||||
@@ -0,0 +1 @@
|
||||
pub mod tts_model_document;
|
||||
@@ -0,0 +1,43 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use enums::common::visibility::Visibility;
|
||||
use tokens::tokens::tts_models::TtsModelToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
use crate::traits::document::Document;
|
||||
|
||||
/// The current name for the TTS model index.
|
||||
/// We may need to perform migrations in the future, so this may grow to keeping
|
||||
/// tabs of multiple indices in the future.
|
||||
pub const TTS_MODEL_INDEX : &str = "tts_model_v1";
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug)]
|
||||
pub struct TtsModelDocument {
|
||||
pub token: TtsModelToken,
|
||||
pub title: String,
|
||||
|
||||
pub ietf_language_tag: String,
|
||||
pub ietf_primary_language_subtag: String,
|
||||
|
||||
pub creator_username: String,
|
||||
pub creator_display_name: String,
|
||||
pub creator_user_token: UserToken,
|
||||
pub creator_set_visibility: Visibility,
|
||||
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl Document for TtsModelDocument {
|
||||
fn get_document_id(&self) -> String {
|
||||
self.token.to_string()
|
||||
}
|
||||
|
||||
fn get_document_path(&self) -> PathBuf {
|
||||
let document_id = self.get_document_id();
|
||||
PathBuf::from(format!("/{TTS_MODEL_INDEX}/_doc/{document_id}"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
//! elasticsearch_schema
|
||||
//!
|
||||
//! The purpose of this library is to keep a list of schemas for our ElasticSearch documents
|
||||
//! as well as index definitions.
|
||||
//!
|
||||
|
||||
// Never allow these
|
||||
#![forbid(private_in_public)]
|
||||
#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things
|
||||
|
||||
// Okay to toggle
|
||||
//#![forbid(unreachable_patterns)]
|
||||
//#![forbid(unused_imports)]
|
||||
//#![forbid(unused_mut)]
|
||||
//#![forbid(unused_variables)]
|
||||
|
||||
// Always allow
|
||||
#![allow(dead_code)]
|
||||
#![allow(non_snake_case)]
|
||||
|
||||
pub mod documents;
|
||||
pub mod searches;
|
||||
pub mod traits;
|
||||
pub mod utils;
|
||||
@@ -0,0 +1 @@
|
||||
pub mod search_tts_models;
|
||||
@@ -0,0 +1,150 @@
|
||||
use elasticsearch::{Elasticsearch, SearchParts};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use errors::AnyhowResult;
|
||||
|
||||
use crate::documents::tts_model_document::TtsModelDocument;
|
||||
|
||||
pub async fn search_tts_models(
|
||||
client: &Elasticsearch,
|
||||
search_term: &str,
|
||||
maybe_language_subtag: Option<&str>,
|
||||
) -> AnyhowResult<Vec<TtsModelDocument>> {
|
||||
|
||||
let search_json = match maybe_language_subtag {
|
||||
None => query_tts_models(search_term),
|
||||
Some(language) => query_tts_models_with_required_language(search_term, language)
|
||||
};
|
||||
|
||||
let search_response = client
|
||||
.search(SearchParts::None)
|
||||
.body(search_json)
|
||||
.allow_no_indices(true)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let _status_code = search_response.status_code();
|
||||
|
||||
let mut response_json = search_response.json::<Value>().await?;
|
||||
|
||||
let mut hits = response_json.get_mut("hits")
|
||||
.map(|hits| hits.take());
|
||||
|
||||
let mut hits = hits.map(|mut hits| {
|
||||
hits.get_mut("hits")
|
||||
.map(|hits| hits.take())
|
||||
}).flatten();
|
||||
|
||||
let mut documents = Vec::new();
|
||||
|
||||
match hits {
|
||||
Some(Value::Array(inner_hits)) => {
|
||||
for mut hit in inner_hits {
|
||||
let maybe_object = hit.get_mut("_source")
|
||||
.map(|source| source.take());
|
||||
if let Some(value) = maybe_object {
|
||||
let document = serde_json::from_value::<TtsModelDocument>(value)?;
|
||||
documents.push(document);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {},
|
||||
}
|
||||
|
||||
Ok(documents)
|
||||
}
|
||||
|
||||
fn query_tts_models(search_term: &str) -> Value {
|
||||
json!({
|
||||
"query": {
|
||||
"bool": {
|
||||
"must": [
|
||||
{
|
||||
"bool": {
|
||||
"should": [
|
||||
{
|
||||
"fuzzy": {
|
||||
"title": {
|
||||
"value": search_term,
|
||||
"fuzziness": 2
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"match": {
|
||||
"title": {
|
||||
"query": search_term,
|
||||
"boost": 1
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"multi_match": {
|
||||
"query": search_term,
|
||||
"type": "bool_prefix",
|
||||
"fields": [
|
||||
"title",
|
||||
"title._2gram",
|
||||
"title._3gram"
|
||||
],
|
||||
"boost": 50
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn query_tts_models_with_required_language(search_term: &str, language_tag: &str) -> Value {
|
||||
json!({
|
||||
"query": {
|
||||
"bool": {
|
||||
"must": [
|
||||
{
|
||||
"match": {
|
||||
"ietf_primary_language_subtag": language_tag
|
||||
}
|
||||
},
|
||||
{
|
||||
"bool": {
|
||||
"should": [
|
||||
{
|
||||
"fuzzy": {
|
||||
"title": {
|
||||
"value": search_term,
|
||||
"fuzziness": 2
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"match": {
|
||||
"title": {
|
||||
"query": search_term,
|
||||
"boost": 1
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"multi_match": {
|
||||
"query": search_term,
|
||||
"type": "bool_prefix",
|
||||
"fields": [
|
||||
"title",
|
||||
"title._2gram",
|
||||
"title._3gram"
|
||||
],
|
||||
"boost": 50
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
pub trait Document {
|
||||
/// Get the document unique identifier
|
||||
fn get_document_id(&self) -> String;
|
||||
|
||||
/// Return the document path
|
||||
fn get_document_path(&self) -> PathBuf;
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod document;
|
||||
@@ -0,0 +1,61 @@
|
||||
use elasticsearch::Elasticsearch;
|
||||
use elasticsearch::http::StatusCode;
|
||||
use elasticsearch::indices::{IndicesCreateParts, IndicesDeleteParts, IndicesExistsParts};
|
||||
use log::{error, info};
|
||||
use serde_json::Value;
|
||||
|
||||
use errors::{AnyhowResult, bail};
|
||||
|
||||
pub struct CreateIndexArgs <'a> {
|
||||
pub client: &'a Elasticsearch,
|
||||
pub index_name: &'a str,
|
||||
pub index_definition: &'a str,
|
||||
pub delete_existing: bool,
|
||||
}
|
||||
|
||||
// NB: Adapted from elasticsearch crate examples source
|
||||
pub async fn create_index_if_not_exists<'a>(args: CreateIndexArgs<'a>) -> AnyhowResult<()> {
|
||||
info!("Checking if Elasticsearch index exists: {:?}", args.index_name);
|
||||
|
||||
let exists = args.client
|
||||
.indices()
|
||||
.exists(IndicesExistsParts::Index(&[args.index_name]))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if exists.status_code().is_success() && args.delete_existing {
|
||||
info!("Deleting Elasticsearch index: {:?}", args.index_name);
|
||||
|
||||
let delete = args.client
|
||||
.indices()
|
||||
.delete(IndicesDeleteParts::Index(&[args.index_name]))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !delete.status_code().is_success() {
|
||||
error!("Problem deleting index: {}", delete.text().await?);
|
||||
}
|
||||
}
|
||||
|
||||
if exists.status_code() == StatusCode::NOT_FOUND || args.delete_existing {
|
||||
// NB(bt): This is weirdly tricky.
|
||||
// Passing a value to `.body()` with serde_json's json!() literal macro works.
|
||||
// However, Elasticsearch server fails if the argument to `.body()` is a string.
|
||||
// Furthermore, Elasticsearch server fails if serde_json::from_str() is called
|
||||
// without a bound type! We need to bind it to the `Value` type explicitly.
|
||||
let json_value : Value = serde_json::from_str(args.index_definition)?;
|
||||
|
||||
let response = args.client
|
||||
.indices()
|
||||
.create(IndicesCreateParts::Index(args.index_name))
|
||||
.body(json_value)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status_code().is_success() {
|
||||
bail!("Error while creating index");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod create_index_if_not_exists;
|
||||
@@ -25,6 +25,7 @@ container_common = { path = "../../../lib/deprecated/container_common" }
|
||||
cookies = { path = "../../../lib/cookies" }
|
||||
datetimes = { path = "../../../lib/datetimes" }
|
||||
easyenv = { path = "../../../lib/easyenv" }
|
||||
elasticsearch_schema = { path = "../../../lib/elasticsearch_schema" }
|
||||
enums = { path = "../../../_types/enums" }
|
||||
errors = { path = "../../../lib/errors" }
|
||||
google_drive_common = { path = "../../../lib/api_clients/google_drive_common" }
|
||||
@@ -109,6 +110,9 @@ sqlx = { version = "0.6.2", features = [ "mysql", "runtime-actix-native-tls", "o
|
||||
limitation = "0.1.1" # Redis rate limiter
|
||||
r2d2_redis = "0.14.0"
|
||||
|
||||
# Elasticsearch
|
||||
elasticsearch = "8.5.0-alpha.1"
|
||||
|
||||
# In-memory caching
|
||||
lru_time_cache = "0.11.11"
|
||||
|
||||
|
||||
@@ -19,6 +19,9 @@ STRIPE_SECRET_WEBHOOK_SIGNING_KEY=whsec_40b60bd75768052532b63dbeefa7a1b163e182c3
|
||||
NUM_WORKERS=16
|
||||
MYSQL_MAX_CONNECTIONS=50
|
||||
|
||||
# Elasticsearch (development)
|
||||
ELASTICSEARCH_URL=http://localhost:9200
|
||||
|
||||
# REDIS_0 is used for many things, including:
|
||||
# - generic inference job keepalive / polling / TTL
|
||||
# - (many others I haven't documented yet)
|
||||
|
||||
@@ -12,4 +12,5 @@ pub mod get_tts_result;
|
||||
pub mod get_tts_upload_model_job_status;
|
||||
pub mod list_tts_models;
|
||||
pub mod list_user_tts_inference_results;
|
||||
pub mod list_user_tts_models;
|
||||
pub mod list_user_tts_models;
|
||||
pub mod search_tts_models_handler;
|
||||
+155
@@ -0,0 +1,155 @@
|
||||
// NB: Incrementally getting rid of build warnings...
|
||||
#![forbid(unused_imports)]
|
||||
#![forbid(unused_mut)]
|
||||
#![forbid(unused_variables)]
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use actix_web::{HttpRequest, HttpResponse, web};
|
||||
use actix_web::error::ResponseError;
|
||||
use actix_web::http::StatusCode;
|
||||
use chrono::{DateTime, Utc};
|
||||
use log::error;
|
||||
|
||||
use elasticsearch_schema::searches::search_tts_models::search_tts_models;
|
||||
use enums::common::visibility::Visibility;
|
||||
use tokens::tokens::tts_models::TtsModelToken;
|
||||
use tokens::tokens::users::UserToken;
|
||||
|
||||
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
|
||||
use crate::server_state::ServerState;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SearchTtsModelsRequest {
|
||||
pub search_term: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
pub struct TtsModel {
|
||||
pub model_token: TtsModelToken,
|
||||
//pub tts_model_type: String,
|
||||
pub creator_user_token: UserToken,
|
||||
pub creator_username: String,
|
||||
pub creator_display_name: String,
|
||||
//pub creator_gravatar_hash: String,
|
||||
pub title: String,
|
||||
pub ietf_language_tag: String,
|
||||
pub ietf_primary_language_subtag: String,
|
||||
//pub is_front_page_featured: bool,
|
||||
//pub is_twitch_featured: bool,
|
||||
//pub maybe_suggested_unique_bot_command: Option<String>,
|
||||
|
||||
pub creator_set_visibility: Visibility,
|
||||
|
||||
//pub user_ratings: UserRatingsStats,
|
||||
|
||||
///// Category assignments
|
||||
///// From non-deleted, mod-approved categories only
|
||||
//pub category_tokens: HashSet<String>,
|
||||
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
//#[derive(Serialize, Clone)]
|
||||
//pub struct UserRatingsStats {
|
||||
// pub positive_count: u32,
|
||||
// pub negative_count: u32,
|
||||
// /// Total count does not take into account "neutral" ratings.
|
||||
// pub total_count: u32,
|
||||
//}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct SearchTtsModelsSuccessResponse {
|
||||
pub success: bool,
|
||||
pub models: Vec<TtsModel>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SearchTtsModelsError {
|
||||
ServerError,
|
||||
}
|
||||
|
||||
impl ResponseError for SearchTtsModelsError {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
match *self {
|
||||
SearchTtsModelsError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR,
|
||||
}
|
||||
}
|
||||
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
let error_reason = match self {
|
||||
SearchTtsModelsError::ServerError => "server error".to_string(),
|
||||
};
|
||||
|
||||
to_simple_json_error(&error_reason, self.status_code())
|
||||
}
|
||||
}
|
||||
|
||||
// NB: Not using derive_more::Display since Clion doesn't understand it.
|
||||
impl fmt::Display for SearchTtsModelsError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{:?}", self)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn search_tts_models_handler(
|
||||
_http_request: HttpRequest,
|
||||
request: web::Json<SearchTtsModelsRequest>,
|
||||
server_state: web::Data<Arc<ServerState>>) -> Result<HttpResponse, SearchTtsModelsError>
|
||||
{
|
||||
let results = search_tts_models(
|
||||
&server_state.elasticsearch,
|
||||
&request.search_term,
|
||||
None)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Searching error: {:?}", err);
|
||||
SearchTtsModelsError::ServerError
|
||||
})?;
|
||||
|
||||
let results = results.into_iter()
|
||||
.map(|result| TtsModel {
|
||||
model_token: result.token,
|
||||
creator_user_token: result.creator_user_token,
|
||||
creator_username: result.creator_username,
|
||||
creator_display_name: result.creator_display_name,
|
||||
title: result.title,
|
||||
ietf_language_tag: result.ietf_language_tag,
|
||||
ietf_primary_language_subtag: result.ietf_primary_language_subtag,
|
||||
creator_set_visibility: result.creator_set_visibility,
|
||||
created_at: result.created_at,
|
||||
updated_at: result.updated_at,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// TODO(bt,2023-10-27): For some reason Elasticsearch returns duplicates. Maybe we populated the
|
||||
// DB twice? Need to filter them out, or React chokes and gets stuck on duplicates. (Effectively
|
||||
// freezing them into the UI, despite component updates)
|
||||
|
||||
let mut added_tokens = HashSet::new();
|
||||
let mut new_results = Vec::with_capacity(results.len());
|
||||
|
||||
for result in results.into_iter() {
|
||||
if added_tokens.contains(&result.model_token) {
|
||||
continue;
|
||||
}
|
||||
added_tokens.insert(result.model_token.clone());
|
||||
new_results.push(result);
|
||||
}
|
||||
|
||||
let response = SearchTtsModelsSuccessResponse {
|
||||
success: true,
|
||||
models: new_results,
|
||||
};
|
||||
|
||||
let body = serde_json::to_string(&response)
|
||||
.map_err(|_e| SearchTtsModelsError::ServerError)?;
|
||||
|
||||
Ok(HttpResponse::Ok()
|
||||
.content_type("application/json")
|
||||
.body(body))
|
||||
}
|
||||
|
||||
@@ -24,6 +24,8 @@ use actix_cors::Cors;
|
||||
use actix_web::{App, HttpResponse, HttpServer, middleware, web};
|
||||
use actix_web::middleware::{DefaultHeaders, Logger};
|
||||
use anyhow::anyhow;
|
||||
use elasticsearch::Elasticsearch;
|
||||
use elasticsearch::http::transport::Transport;
|
||||
use futures::Future;
|
||||
use limitation::Limiter;
|
||||
use log::{error, info};
|
||||
@@ -230,6 +232,10 @@ async fn main() -> AnyhowResult<()> {
|
||||
RedisRateLimiter::new(limiter, "model_upload", limiter_enabled)
|
||||
};
|
||||
|
||||
info!("Connecting to elasticsearch...");
|
||||
|
||||
let elasticsearch = get_elasticsearch_client()?;
|
||||
|
||||
info!("Reading env vars and setting up utils...");
|
||||
|
||||
let bind_address = easyenv::get_env_string_or_default("BIND_ADDRESS", DEFAULT_BIND_ADDRESS);
|
||||
@@ -450,6 +456,7 @@ async fn main() -> AnyhowResult<()> {
|
||||
third_party_url_redirector,
|
||||
health_check_status,
|
||||
mysql_pool: pool,
|
||||
elasticsearch,
|
||||
redis_pool,
|
||||
redis_ttl_cache,
|
||||
redis_rate_limiters: RedisRateLimiters {
|
||||
@@ -671,3 +678,9 @@ pub async fn serve(server_state: ServerState) -> AnyhowResult<()>
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_elasticsearch_client() -> AnyhowResult<Elasticsearch> {
|
||||
let transport = Transport::single_node(&easyenv::get_env_string_required("ELASTICSEARCH_URL")?)?;
|
||||
let client = Elasticsearch::new(transport);
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
@@ -88,6 +88,7 @@ use crate::http_server::endpoints::tts::get_tts_upload_model_job_status::get_tts
|
||||
use crate::http_server::endpoints::tts::list_tts_models::list_tts_models_handler;
|
||||
use crate::http_server::endpoints::tts::list_user_tts_inference_results::list_user_tts_inference_results_handler;
|
||||
use crate::http_server::endpoints::tts::list_user_tts_models::list_user_tts_models_handler;
|
||||
use crate::http_server::endpoints::tts::search_tts_models_handler::search_tts_models_handler;
|
||||
use crate::http_server::endpoints::twitch::event_rules::create_event_rule::create_twitch_event_rule_handler;
|
||||
use crate::http_server::endpoints::twitch::event_rules::delete_event_rule::delete_twitch_event_rule_handler;
|
||||
use crate::http_server::endpoints::twitch::event_rules::edit_event_rule::edit_twitch_event_rule_handler;
|
||||
@@ -417,6 +418,11 @@ fn add_tts_routes<T, B> (app: App<T>) -> App<T>
|
||||
.route(web::get().to(list_tts_models_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(
|
||||
web::resource("/search")
|
||||
.route(web::post().to(search_tts_models_handler))
|
||||
.route(web::head().to(|| HttpResponse::Ok()))
|
||||
)
|
||||
.service(
|
||||
web::resource("/model/{token}")
|
||||
.route(web::get().to(get_tts_model_handler))
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use elasticsearch::Elasticsearch;
|
||||
use r2d2_redis::{r2d2, RedisConnectionManager};
|
||||
use sqlx::MySqlPool;
|
||||
|
||||
@@ -54,6 +55,8 @@ pub struct ServerState {
|
||||
|
||||
pub mysql_pool: MySqlPool,
|
||||
|
||||
pub elasticsearch: Elasticsearch,
|
||||
|
||||
pub redis_pool: r2d2::Pool<RedisConnectionManager>,
|
||||
pub redis_ttl_cache: RedisTtlCache,
|
||||
|
||||
|
||||
Reference in New Issue
Block a user