Skip to content

style: rename ResourceSetupStatusCheck to ResourceSetupStatus #458

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions python/cocoindex/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,15 +85,15 @@ def setup():
Check and apply backend setup changes for flows, including the internal and target storage
(to export).
"""
status_check = sync_setup()
click.echo(status_check)
if status_check.is_up_to_date():
setup_status = sync_setup()
click.echo(setup_status)
if setup_status.is_up_to_date():
click.echo("No changes need to be pushed.")
return
if not click.confirm(
"Changes need to be pushed. Continue? [yes/N]", default=False, show_default=False):
return
apply_setup_changes(status_check)
apply_setup_changes(setup_status)

@cli.command()
@click.argument("flow_name", type=str, nargs=-1)
Expand All @@ -112,15 +112,15 @@ def drop(flow_name: tuple[str, ...], drop_all: bool):
flow_names = [fl.name for fl in flow.flows()]
else:
flow_names = list(flow_name)
status_check = drop_setup(flow_names)
click.echo(status_check)
if status_check.is_up_to_date():
setup_status = drop_setup(flow_names)
click.echo(setup_status)
if setup_status.is_up_to_date():
click.echo("No flows need to be dropped.")
return
if not click.confirm(
"Changes need to be pushed. Continue? [yes/N]", default=False, show_default=False):
return
apply_setup_changes(status_check)
apply_setup_changes(setup_status)

@cli.command()
@click.argument("flow_name", type=str, required=False)
Expand Down
8 changes: 4 additions & 4 deletions python/cocoindex/setup.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
from . import flow
from . import _engine

def sync_setup() -> _engine.SetupStatusCheck:
def sync_setup() -> _engine.SetupStatus:
flow.ensure_all_flows_built()
return _engine.sync_setup()

def drop_setup(flow_names: list[str]) -> _engine.SetupStatusCheck:
def drop_setup(flow_names: list[str]) -> _engine.SetupStatus:
flow.ensure_all_flows_built()
return _engine.drop_setup(flow_names)

def flow_names_with_setup() -> list[str]:
return _engine.flow_names_with_setup()

def apply_setup_changes(status_check: _engine.SetupStatusCheck):
_engine.apply_setup_changes(status_check)
def apply_setup_changes(setup_status: _engine.SetupStatus):
_engine.apply_setup_changes(setup_status)
6 changes: 3 additions & 3 deletions src/builder/analyzed_flow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use super::{analyzer, plan};
use crate::{
ops::registry::ExecutorFactoryRegistry,
service::error::{shared_ok, SharedError, SharedResultExt},
setup::{self, ObjectSetupStatusCheck},
setup::{self, ObjectSetupStatus},
};

pub struct AnalyzedFlow {
Expand All @@ -29,9 +29,9 @@ impl AnalyzedFlow {
existing_flow_ss,
registry,
)?;
let setup_status_check =
let setup_status =
setup::check_flow_setup_status(Some(&desired_state), existing_flow_ss).await?;
let execution_plan = if setup_status_check.is_up_to_date() {
let execution_plan = if setup_status.is_up_to_date() {
Some(
async move {
shared_ok(Arc::new(
Expand Down
12 changes: 6 additions & 6 deletions src/execution/db_tracking_setup.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use crate::prelude::*;

use crate::setup::{CombinedState, ResourceSetupInfo, ResourceSetupStatusCheck, SetupChangeType};
use crate::setup::{CombinedState, ResourceSetupInfo, ResourceSetupStatus, SetupChangeType};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;

Expand Down Expand Up @@ -53,7 +53,7 @@ pub struct TrackingTableSetupState {
}

#[derive(Debug)]
pub struct TrackingTableSetupStatusCheck {
pub struct TrackingTableSetupStatus {
pub desired_state: Option<TrackingTableSetupState>,

pub legacy_table_names: Vec<String>,
Expand All @@ -62,7 +62,7 @@ pub struct TrackingTableSetupStatusCheck {
pub source_ids_to_delete: Vec<i32>,
}

impl TrackingTableSetupStatusCheck {
impl TrackingTableSetupStatus {
pub fn new(
desired: Option<&TrackingTableSetupState>,
existing: &CombinedState<TrackingTableSetupState>,
Expand Down Expand Up @@ -91,19 +91,19 @@ impl TrackingTableSetupStatusCheck {

pub fn into_setup_info(
self,
) -> ResourceSetupInfo<(), TrackingTableSetupState, TrackingTableSetupStatusCheck> {
) -> ResourceSetupInfo<(), TrackingTableSetupState, TrackingTableSetupStatus> {
ResourceSetupInfo {
key: (),
state: self.desired_state.clone(),
description: "Tracking Table".to_string(),
status_check: Some(self),
setup_status: Some(self),
legacy_key: None,
}
}
}

#[async_trait]
impl ResourceSetupStatusCheck for TrackingTableSetupStatusCheck {
impl ResourceSetupStatus for TrackingTableSetupStatus {
fn describe_changes(&self) -> Vec<String> {
let mut changes: Vec<String> = vec![];
if self.desired_state.is_some() && !self.legacy_table_names.is_empty() {
Expand Down
26 changes: 16 additions & 10 deletions src/llm/anthropic.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
use crate::llm::{
LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, LlmSpec, OutputFormat,
ToJsonSchemaOptions,
};
use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use crate::llm::{LlmGenerationClient, LlmSpec, LlmGenerateRequest, LlmGenerateResponse, ToJsonSchemaOptions, OutputFormat};
use anyhow::{Result, bail, Context};
use serde_json::Value;
use json5;
use serde_json::Value;

use crate::api_bail;
use urlencoding::encode;
Expand Down Expand Up @@ -64,7 +67,8 @@ impl LlmGenerationClient for Client {

let encoded_api_key = encode(&self.api_key);

let resp = self.client
let resp = self
.client
.post(url)
.header("x-api-key", encoded_api_key.as_ref())
.header("anthropic-version", "2023-06-01")
Expand All @@ -81,7 +85,7 @@ impl LlmGenerationClient for Client {
// println!("Anthropic API full response: {resp_json:?}");

let resp_content = &resp_json["content"];
let tool_name = "report_result";
let tool_name = "report_result";
let mut extracted_json: Option<Value> = None;
if let Some(array) = resp_content.as_array() {
for item in array {
Expand Down Expand Up @@ -116,14 +120,16 @@ impl LlmGenerationClient for Client {
}
}
}
},
_ => return Err(anyhow::anyhow!("No structured tool output or text found in response")),
}
_ => {
return Err(anyhow::anyhow!(
"No structured tool output or text found in response"
))
}
}
};

Ok(LlmGenerateResponse {
text,
})
Ok(LlmGenerateResponse { text })
}

fn json_schema_options(&self) -> ToJsonSchemaOptions {
Expand Down
18 changes: 12 additions & 6 deletions src/llm/gemini.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
use crate::api_bail;
use crate::llm::{
LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, LlmSpec, OutputFormat,
ToJsonSchemaOptions,
};
use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use crate::llm::{LlmGenerationClient, LlmSpec, LlmGenerateRequest, LlmGenerateResponse, ToJsonSchemaOptions, OutputFormat};
use anyhow::{Result, bail, Context};
use serde_json::Value;
use crate::api_bail;
use urlencoding::encode;

pub struct Client {
Expand Down Expand Up @@ -76,10 +79,13 @@ impl LlmGenerationClient for Client {
let api_key = &self.api_key;
let url = format!(
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
encode(&self.model), encode(api_key)
encode(&self.model),
encode(api_key)
);

let resp = self.client.post(&url)
let resp = self
.client
.post(&url)
.json(&payload)
.send()
.await
Expand Down Expand Up @@ -107,4 +113,4 @@ impl LlmGenerationClient for Client {
top_level_must_be_object: true,
}
}
}
}
4 changes: 2 additions & 2 deletions src/llm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,10 +52,10 @@ pub trait LlmGenerationClient: Send + Sync {
fn json_schema_options(&self) -> ToJsonSchemaOptions;
}

mod anthropic;
mod gemini;
mod ollama;
mod openai;
mod gemini;
mod anthropic;

pub async fn new_llm_generation_client(spec: LlmSpec) -> Result<Box<dyn LlmGenerationClient>> {
let client = match spec.api_type {
Expand Down
8 changes: 4 additions & 4 deletions src/ops/factory_bases.rs
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
desired_state: Option<Self::SetupState>,
existing_states: setup::CombinedState<Self::SetupState>,
auth_registry: &Arc<AuthRegistry>,
) -> Result<impl setup::ResourceSetupStatusCheck + 'static>;
) -> Result<impl setup::ResourceSetupStatus + 'static>;

fn check_state_compatibility(
&self,
Expand Down Expand Up @@ -398,21 +398,21 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
desired_state: Option<serde_json::Value>,
existing_states: setup::CombinedState<serde_json::Value>,
auth_registry: &Arc<AuthRegistry>,
) -> Result<Box<dyn setup::ResourceSetupStatusCheck>> {
) -> Result<Box<dyn setup::ResourceSetupStatus>> {
let key: T::Key = serde_json::from_value(key.clone())?;
let desired_state: Option<T::SetupState> = desired_state
.map(|v| serde_json::from_value(v.clone()))
.transpose()?;
let existing_states = from_json_combined_state(existing_states)?;
let status_check = StorageFactoryBase::check_setup_status(
let setup_status = StorageFactoryBase::check_setup_status(
self,
key,
desired_state,
existing_states,
auth_registry,
)
.await?;
Ok(Box::new(status_check))
Ok(Box::new(setup_status))
}

fn describe_resource(&self, key: &serde_json::Value) -> Result<String> {
Expand Down
2 changes: 1 addition & 1 deletion src/ops/interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ pub trait ExportTargetFactory: Send + Sync {
desired_state: Option<serde_json::Value>,
existing_states: setup::CombinedState<serde_json::Value>,
auth_registry: &Arc<AuthRegistry>,
) -> Result<Box<dyn setup::ResourceSetupStatusCheck>>;
) -> Result<Box<dyn setup::ResourceSetupStatus>>;

/// Normalize the key. e.g. the JSON format may change (after code change, e.g. new optional field or field ordering), even if the underlying value is not changed.
/// This should always return the canonical serialized form.
Expand Down
16 changes: 8 additions & 8 deletions src/ops/storages/neo4j.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::prelude::*;

use super::spec::{GraphDeclaration, GraphElementMapping, NodeFromFieldsSpec, TargetFieldMapping};
use crate::setup::components::{self, State};
use crate::setup::{ResourceSetupStatusCheck, SetupChangeType};
use crate::setup::{ResourceSetupStatus, SetupChangeType};
use crate::{ops::sdk::*, setup::CombinedState};

use indoc::formatdoc;
Expand Down Expand Up @@ -855,7 +855,7 @@ fn build_composite_field_names(qualifier: &str, field_names: &[String]) -> Strin
}
#[derive(Derivative)]
#[derivative(Debug)]
struct SetupStatusCheck {
struct SetupStatus {
key: GraphElement,
#[derivative(Debug = "ignore")]
graph_pool: Arc<GraphPool>,
Expand All @@ -864,7 +864,7 @@ struct SetupStatusCheck {
change_type: SetupChangeType,
}

impl SetupStatusCheck {
impl SetupStatus {
fn new(
key: GraphElement,
graph_pool: Arc<GraphPool>,
Expand Down Expand Up @@ -908,7 +908,7 @@ impl SetupStatusCheck {
}

#[async_trait]
impl ResourceSetupStatusCheck for SetupStatusCheck {
impl ResourceSetupStatus for SetupStatus {
fn describe_changes(&self) -> Vec<String> {
let mut result = vec![];
if let Some(data_clear) = &self.data_clear {
Expand Down Expand Up @@ -1255,24 +1255,24 @@ impl StorageFactoryBase for Factory {
desired: Option<SetupState>,
existing: CombinedState<SetupState>,
auth_registry: &Arc<AuthRegistry>,
) -> Result<impl ResourceSetupStatusCheck + 'static> {
) -> Result<impl ResourceSetupStatus + 'static> {
let conn_spec = auth_registry.get::<ConnectionSpec>(&key.connection)?;
let base = SetupStatusCheck::new(
let base = SetupStatus::new(
key,
self.graph_pool.clone(),
conn_spec.clone(),
desired.as_ref(),
&existing,
);
let comp = components::StatusCheck::create(
let comp = components::Status::create(
SetupComponentOperator {
graph_pool: self.graph_pool.clone(),
conn_spec: conn_spec.clone(),
},
desired,
existing,
)?;
Ok(components::combine_status_checks(base, comp))
Ok(components::combine_setup_statuss(base, comp))
}

fn check_state_compatibility(
Expand Down
10 changes: 5 additions & 5 deletions src/ops/storages/postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,7 @@ impl TableSetupAction {
}

#[derive(Debug)]
pub struct SetupStatusCheck {
pub struct SetupStatus {
db_pool: PgPool,
table_name: String,

Expand All @@ -585,7 +585,7 @@ pub struct SetupStatusCheck {
desired_table_setup: Option<TableSetupAction>,
}

impl SetupStatusCheck {
impl SetupStatus {
fn new(
db_pool: PgPool,
table_name: String,
Expand Down Expand Up @@ -739,7 +739,7 @@ fn describe_index_spec(index_name: &str, index_spec: &VectorIndexDef) -> String
}

#[async_trait]
impl setup::ResourceSetupStatusCheck for SetupStatusCheck {
impl setup::ResourceSetupStatus for SetupStatus {
fn describe_changes(&self) -> Vec<String> {
let mut descriptions = vec![];
if self.drop_existing {
Expand Down Expand Up @@ -976,8 +976,8 @@ impl StorageFactoryBase for Factory {
desired: Option<SetupState>,
existing: setup::CombinedState<SetupState>,
auth_registry: &Arc<AuthRegistry>,
) -> Result<impl setup::ResourceSetupStatusCheck + 'static> {
Ok(SetupStatusCheck::new(
) -> Result<impl setup::ResourceSetupStatus + 'static> {
Ok(SetupStatus::new(
get_db_pool(key.database.as_ref(), auth_registry).await?,
key.table_name,
desired,
Expand Down
2 changes: 1 addition & 1 deletion src/ops/storages/qdrant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,7 @@ impl StorageFactoryBase for Arc<Factory> {
_desired: Option<()>,
_existing: setup::CombinedState<()>,
_auth_registry: &Arc<AuthRegistry>,
) -> Result<impl setup::ResourceSetupStatusCheck + 'static> {
) -> Result<impl setup::ResourceSetupStatus + 'static> {
Err(anyhow!("Set `setup_by_user` to `true` to export to Qdrant")) as Result<Infallible, _>
}

Expand Down
Loading