Skip to content
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
13 changes: 12 additions & 1 deletion examples/text_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from dotenv import load_dotenv
from psycopg_pool import ConnectionPool
from pgvector.psycopg import register_vector
from typing import Any
import cocoindex
import os
import functools
Expand Down Expand Up @@ -72,6 +71,18 @@ def text_embedding_flow(
metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
)
],
attachments=[
cocoindex.targets.PostgresSqlCommand(
name="new_column_foo",
setup_sql="ALTER TABLE TextEmbedding__doc_embeddings DROP COLUMN IF EXISTS foo; ALTER TABLE TextEmbedding__doc_embeddings ADD COLUMN foo TEXT",
teardown_sql="ALTER TABLE TextEmbedding__doc_embeddings DROP COLUMN IF EXISTS foo",
),
cocoindex.targets.PostgresSqlCommand(
name="new_column_bar",
setup_sql="ALTER TABLE TextEmbedding__doc_embeddings ADD COLUMN bar TEXT",
teardown_sql="ALTER TABLE TextEmbedding__doc_embeddings DROP COLUMN IF EXISTS bar",
),
],
)


Expand Down
2 changes: 1 addition & 1 deletion python/cocoindex/targets/_engine_builtin_specs.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ class Postgres(op.TargetSpec):
schema: str | None = None


class PostgresSqlAttachment(op.TargetAttachmentSpec):
class PostgresSqlCommand(op.TargetAttachmentSpec):
"""Attachment to execute specified SQL statements for Postgres targets."""

name: str
Expand Down
34 changes: 17 additions & 17 deletions src/ops/targets/postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -791,26 +791,26 @@ impl TargetFactoryBase for TargetFactory {
////////////////////////////////////////////////////////////

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SqlStatementAttachmentSpec {
pub struct SqlCommandSpec {
name: String,
setup_sql: String,
teardown_sql: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SqlStatementAttachmentState {
pub struct SqlCommandState {
setup_sql: String,
teardown_sql: Option<String>,
}

pub struct SqlStatementAttachmentSetupChange {
pub struct SqlCommandSetupChange {
db_pool: PgPool,
setup_sql_to_run: Option<String>,
teardown_sql_to_run: IndexSet<String>,
}

#[async_trait]
impl AttachmentSetupChange for SqlStatementAttachmentSetupChange {
impl AttachmentSetupChange for SqlCommandSetupChange {
fn describe_changes(&self) -> Vec<String> {
let mut result = vec![];
for teardown_sql in self.teardown_sql_to_run.iter() {
Expand All @@ -833,30 +833,30 @@ impl AttachmentSetupChange for SqlStatementAttachmentSetupChange {
}
}

struct SqlAttachmentFactory;
struct SqlCommandFactory;

#[async_trait]
impl TargetSpecificAttachmentFactoryBase for SqlAttachmentFactory {
impl TargetSpecificAttachmentFactoryBase for SqlCommandFactory {
type TargetKey = TableId;
type TargetSpec = Spec;
type Spec = SqlStatementAttachmentSpec;
type Spec = SqlCommandSpec;
type SetupKey = String;
type SetupState = SqlStatementAttachmentState;
type SetupChange = SqlStatementAttachmentSetupChange;
type SetupState = SqlCommandState;
type SetupChange = SqlCommandSetupChange;

fn name(&self) -> &str {
"PostgresSqlAttachment"
"PostgresSqlCommand"
}

fn get_state(
&self,
_target_name: &str,
_target_spec: &Spec,
attachment_spec: SqlStatementAttachmentSpec,
attachment_spec: SqlCommandSpec,
) -> Result<TypedTargetAttachmentState<Self>> {
Ok(TypedTargetAttachmentState {
setup_key: attachment_spec.name,
setup_state: SqlStatementAttachmentState {
setup_state: SqlCommandState {
setup_sql: attachment_spec.setup_sql,
teardown_sql: attachment_spec.teardown_sql,
},
Expand All @@ -867,10 +867,10 @@ impl TargetSpecificAttachmentFactoryBase for SqlAttachmentFactory {
&self,
target_key: &TableId,
_attachment_key: &String,
new_state: Option<SqlStatementAttachmentState>,
existing_states: setup::CombinedState<SqlStatementAttachmentState>,
new_state: Option<SqlCommandState>,
existing_states: setup::CombinedState<SqlCommandState>,
context: &interface::FlowInstanceContext,
) -> Result<Option<SqlStatementAttachmentSetupChange>> {
) -> Result<Option<SqlCommandSetupChange>> {
let teardown_sql_to_run: IndexSet<String> = if new_state.is_none() {
existing_states
.possible_versions()
Expand All @@ -888,7 +888,7 @@ impl TargetSpecificAttachmentFactoryBase for SqlAttachmentFactory {
};
let change = if setup_sql_to_run.is_some() || !teardown_sql_to_run.is_empty() {
let db_pool = get_db_pool(target_key.database.as_ref(), &context.auth_registry).await?;
Some(SqlStatementAttachmentSetupChange {
Some(SqlCommandSetupChange {
db_pool,
setup_sql_to_run,
teardown_sql_to_run,
Expand All @@ -902,6 +902,6 @@ impl TargetSpecificAttachmentFactoryBase for SqlAttachmentFactory {

pub fn register(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
TargetFactory.register(registry)?;
SqlAttachmentFactory.register(registry)?;
SqlCommandFactory.register(registry)?;
Ok(())
}