optimize db api

This commit is contained in:
yggverse 2026-01-10 01:41:20 +02:00
parent f48e256fad
commit ee083dfc45
11 changed files with 215 additions and 192 deletions

View file

@ -15,6 +15,6 @@ chrono = "0.4.42"
clap = { version = "4.5.54", features = ["derive"] }
lancor = "0.1.1"
log = "0.4.29"
mysql = { package = "rssto-mysql", version = "0.1.0", features = ["transactional"], path = "../mysql" }
mysql = { package = "rssto-mysql", version = "0.1.0", features = ["transaction"], path = "../mysql" }
tokio = { version = "1.0", features = ["full"] }
tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }

View file

@ -2,7 +2,7 @@ mod argument;
use anyhow::Result;
use argument::Argument;
use mysql::Transactional;
use mysql::Database;
#[tokio::main]
async fn main() -> Result<()> {
@ -32,12 +32,17 @@ async fn main() -> Result<()> {
"{}://{}:{}",
arg.llm_scheme, arg.llm_host, arg.llm_port
))?;
let db = Database::pool(
&arg.mysql_host,
arg.mysql_port,
&arg.mysql_username,
&arg.mysql_password,
&arg.mysql_database,
)?;
// find existing ID by model name or create a new one
// * this feature should be moved to a separate CLI tool
let provider_id = {
let mut db = tx(&arg)?;
match db.provider_id_by_name(&arg.llm_model)? {
let mut conn = db.connection()?;
match conn.provider_id_by_name(&arg.llm_model)? {
Some(provider_id) => {
debug!(
"Use existing DB provider #{} matches model name `{}`",
@ -46,12 +51,11 @@ async fn main() -> Result<()> {
provider_id
}
None => {
let provider_id = db.insert_provider(&arg.llm_model)?;
let provider_id = conn.insert_provider(&arg.llm_model)?;
info!(
"Provider `{}` not found in database, created new one with ID `{provider_id}`",
&arg.llm_model
);
db.commit()?;
provider_id
}
}
@ -60,42 +64,38 @@ async fn main() -> Result<()> {
info!("Daemon started");
loop {
debug!("New queue begin...");
{
// disconnect from the database immediately when exiting this scope,
// in case the `update` queue is enabled and pending for a while.
let mut db = tx(&arg)?;
for source in db.contents_queue_for_provider_id(provider_id)? {
debug!(
"Begin generating `content_id` #{} using `provider_id` #{provider_id}.",
source.content_id
);
let mut tx = db.transaction()?;
for source in tx.contents_queue_for_provider_id(provider_id)? {
debug!(
"Begin generating `content_id` #{} using `provider_id` #{provider_id}.",
source.content_id
);
let title = llm
.chat_completion(ChatCompletionRequest::new(&arg.llm_model).message(
Message::user(format!("{}\n{}", arg.llm_message, source.title)),
))
.await?;
let title =
llm.chat_completion(ChatCompletionRequest::new(&arg.llm_model).message(
Message::user(format!("{}\n{}", arg.llm_message, source.title)),
))
.await?;
let description = llm
.chat_completion(ChatCompletionRequest::new(&arg.llm_model).message(
Message::user(format!("{}\n{}", arg.llm_message, source.description)),
))
.await?;
let description =
llm.chat_completion(ChatCompletionRequest::new(&arg.llm_model).message(
Message::user(format!("{}\n{}", arg.llm_message, source.description)),
))
.await?;
let content_id = db.insert_content(
source.channel_item_id,
Some(provider_id),
&title.choices[0].message.content,
&description.choices[0].message.content,
)?;
let content_id = tx.insert_content(
source.channel_item_id,
Some(provider_id),
&title.choices[0].message.content,
&description.choices[0].message.content,
)?;
debug!(
"Created `content_id` #{content_id} using `content_id` #{} source by `provider_id` #{provider_id}.",
source.content_id
)
}
db.commit()?
debug!(
"Created `content_id` #{content_id} using `content_id` #{} source by `provider_id` #{provider_id}.",
source.content_id
)
}
tx.commit()?;
debug!("Queue completed");
if let Some(update) = arg.update {
debug!("Wait {update} seconds to continue...");
@ -105,15 +105,3 @@ async fn main() -> Result<()> {
}
}
}
// in fact, there is no need for a transaction at this moment,
// as there are no related table updates, but who knows what the future holds
fn tx(arg: &Argument) -> Result<Transactional> {
Ok(Transactional::connect(
&arg.mysql_host,
arg.mysql_port,
&arg.mysql_username,
&arg.mysql_password,
&arg.mysql_database,
)?)
}