Use pk in producer_worker of embedding

lanterndata · Mar 26, 2024 · 810fc87 · 810fc87
1 parent 81a5dbc
commit 810fc87
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/lantern_cli/src/embeddings/mod.rs b/lantern_cli/src/embeddings/mod.rs
@@ -46,6 +46,7 @@ fn producer_worker(
         let column = &args.column;
         let schema = &args.schema;
         let table = &args.table;
+        let pk = &args.pk;
         let full_table_name = get_full_table_name(schema, table);
 
         let filter_sql = if args.filter.is_some() {
@@ -105,8 +106,9 @@ fn producer_worker(
         // With portal we can execute a query and poll values from it in chunks
         let portal = transaction.bind(
             &format!(
-                "SELECT id::text, {column}::text FROM {full_table_name} {filter_sql} {limit_sql};",
+                "SELECT {pk}::text, {column}::text FROM {full_table_name} {filter_sql} {limit_sql};",
                 column = quote_ident(column),
+                pk = quote_ident(pk)
             ),
             &[],
         )?;

diff --git a/lantern_cli/tests/daemon_test_with_db.rs b/lantern_cli/tests/daemon_test_with_db.rs
@@ -10,7 +10,7 @@ use lantern_cli::daemon::{
 };
 use tokio_postgres::{Client, NoTls};
 
-static EMB_LOCK_TABLE_NAME: &'static str = "_lantern_emb_job_locks";
+// static EMB_LOCK_TABLE_NAME: &'static str = "_lantern_emb_job_locks";
 static EMBEDDING_JOBS_TABLE_NAME: &'static str = "_daemon_embedding_jobs";
 static AUTOTUNE_JOBS_TABLE_NAME: &'static str = "_daemon_autotune_jobs";
 static INDEX_JOBS_TABLE_NAME: &'static str = "_daemon_index_jobs";