Worker refactors (#253)

Author: ChuckHend

.github/workflows/extension_upgrade.yml

@@ -50,18 +50,17 @@ jobs:
       - name: Install project dependencies
         run: |
           make setup
-          cargo install pg-trunk
       - name: Test previous version (v0.20.0)
         env:
           HF_API_KEY: ${{ secrets.HF_API_KEY }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           CO_API_KEY: ${{ secrets.CO_API_KEY }}
         run: |
-          trunk install vectorize --version 0.20.0 --pg-config $(cargo pgrx info pg-config pg17)
           git fetch --tags
           git checkout tags/v0.20.0
           # pgrx=0.12.5 required for v0.20.0
           cargo install cargo-pgrx --version 0.12.5 --locked
+          cargo pgrx install --pg-config $(cargo pgrx info pg-config pg17)
           # use integration tests for v0.20.0
           make test-integration
       - name: Test branch's version

Cargo.lock

@@ -2,13 +2,14 @@
  <b>pg_vectorize: a VectorDB on Postgres</b>
 </h1>
 
+[![PostgreSQL](https://img.shields.io/badge/PostgreSQL-13%20%7C%2014%20%7C%2015%20%7C%2016%20%7C%2017%20%7C%2018-336791?logo=postgresql&logoColor=white)](https://www.postgresql.org/)
+
 A Postgres server and extension that automates the transformation and orchestration of text to embeddings and provides hooks into the most popular LLMs. This allows you to do get up and running and automate maintenance for vector search, full text search, and hybrid search, which enables you to quickly build RAG and search engines on Postgres.
 
 This project relies heavily on the work by [pgvector](https://github.com/pgvector/pgvector) for vector similarity search, [pgmq](https://github.com/pgmq/pgmq) for orchestration in background workers, and [SentenceTransformers](https://huggingface.co/sentence-transformers).
 
 ---
 
-[![PostgreSQL](https://img.shields.io/badge/PostgreSQL-13%20%7C%2014%20%7C%2015%20%7C%2016%20%7C%2017%20%7C%2018-336791?logo=postgresql&logoColor=white)](https://www.postgresql.org/)
 
 **API Documentation**: https://chuckhend.github.io/pg_vectorize/
 
@@ -64,7 +65,12 @@ curl -X POST http://localhost:8080/api/v1/table -d '{
 Search using the HTTP API:
 
 ```bash
-curl -X GET "http://localhost:8080/api/v1/search?job_name=my_job&query=camping%20backpack&limit=1" | jq .
+curl -G \
+  "http://localhost:8080/api/v1/search" \
+  --data-urlencode "job_name=my_job" \
+  --data-urlencode "query=camping backpack" \
+  --data-urlencode "limit=1" \
+  | jq .
 ```
 
 ```json
@@ -79,7 +85,7 @@ curl -X GET "http://localhost:8080/api/v1/search?job_name=my_job&query=camping%2
     "rrf_score": 0.03278688524590164,
     "semantic_rank": 1,
     "similarity_score": 0.6296013593673706,
-    "updated_at": "2025-10-04T14:45:16.152526+00:00"
+    "updated_at": "2025-10-05T00:14:39.220893+00:00"
   }
 ]
 ```

README.md

@@ -30,10 +30,35 @@ pub struct Config {
     pub max_retries: i32,
     pub webserver_port: u16,
     pub num_server_workers: usize,
+    pub database_pool_max: u32,
+    pub database_cache_pool_max: u32,
 }
 
 impl Config {
     pub fn from_env() -> Config {
+        // read server worker count first so we can derive sensible defaults
+        let num_server_workers: usize =
+            from_env_default("NUM_SERVER_WORKERS", "8").parse().unwrap();
+
+        // derive a default DB pool size from num_server_workers: 2 connections per worker + 2 extra,
+        // clamped between 4 and 64. This gives headroom for background tasks and short bursts.
+        let derived_db_pool_default: u32 = ((num_server_workers as u32).saturating_mul(2))
+            .saturating_add(2)
+            .clamp(4, 64);
+
+        // allow environment override; fall back to derived default when not set or invalid.
+        let database_pool_max: u32 = env::var("DATABASE_POOL_MAX")
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(derived_db_pool_default);
+
+        // cache pool is typically small; default to max(2, num_server_workers / 4)
+        let derived_cache_pool_default: u32 = ((num_server_workers as u32) / 4).max(2).clamp(2, 16);
+        let database_cache_pool_max: u32 = env::var("DATABASE_CACHE_POOL_MAX")
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(derived_cache_pool_default);
+
         Config {
             proxy_enabled: env::var("VECTORIZE_PROXY_ENABLED")
                 .map(|v| parse_bool_flexible(&v))
@@ -63,7 +88,9 @@ impl Config {
                 .unwrap(),
             max_retries: from_env_default("MAX_RETRIES", "2").parse().unwrap(),
             webserver_port: from_env_default("WEBSERVER_PORT", "8080").parse().unwrap(),
-            num_server_workers: from_env_default("NUM_SERVER_WORKERS", "8").parse().unwrap(),
+            num_server_workers,
+            database_pool_max,
+            database_cache_pool_max,
         }
     }
 }

core/src/config.rs

@@ -327,13 +327,16 @@ BEGIN
         1000 -- default batch size
     );
     FOR batch_result IN SELECT batch FROM vectorize.batch_texts(record_ids, batch_size) LOOP
-        job_messages := array_append(
-            job_messages,
-            jsonb_build_object(
-                'job_name', job_name,
-                'record_ids', batch_result.batch
-            )
-        );
+        -- only append non-null, non-empty batches
+        IF array_length(batch_result.batch, 1) > 0 THEN
+            job_messages := array_append(
+                job_messages,
+                jsonb_build_object(
+                    'job_name', job_name,
+                    'record_ids', batch_result.batch
+                )
+            );
+        END IF;
     END LOOP;
 
     PERFORM pgmq.send_batch(

core/src/query.rs

@@ -70,13 +70,16 @@ BEGIN
     -- create jobs of size batch_size
     batch_size := current_setting('vectorize.batch_size')::integer;
     FOR batch_result IN SELECT batch FROM vectorize.batch_texts(record_ids, batch_size) LOOP
-        job_messages := array_append(
-            job_messages,
-            jsonb_build_object(
-                'job_name', job_name,
-                'record_ids', batch_result.batch
-            )
-        );
+        -- only append non-null, non-empty batches
+        IF array_length(batch_result.batch, 1) > 0 THEN
+            job_messages := array_append(
+                job_messages,
+                jsonb_build_object(
+                    'job_name', job_name,
+                    'record_ids', batch_result.batch
+                )
+            );
+        END IF;
     END LOOP;
 
     PERFORM pgmq.send_batch(

extension/sql/meta.sql

@@ -130,21 +130,20 @@ pub fn resolve_prepared_embed_calls(
     parameters: &[String],
 ) -> Result<Vec<EmbedCall>, VectorizeError> {
     for call in &mut embed_calls {
-        if call.is_prepared {
-            if let (Some(query_idx), Some(project_idx)) =
+        if call.is_prepared
+            && let (Some(query_idx), Some(project_idx)) =
                 (call.query_param_index, call.project_param_index)
-            {
-                if query_idx >= parameters.len() || project_idx >= parameters.len() {
-                    return Err(VectorizeError::EmbeddingGenerationFailed(format!(
-                        "Parameter index out of bounds: query_idx={}, project_idx={}, params_len={}",
-                        query_idx,
-                        project_idx,
-                        parameters.len()
-                    )));
-                }
-                call.query = parameters[query_idx].clone();
-                call.project_name = parameters[project_idx].clone();
+        {
+            if query_idx >= parameters.len() || project_idx >= parameters.len() {
+                return Err(VectorizeError::EmbeddingGenerationFailed(format!(
+                    "Parameter index out of bounds: query_idx={}, project_idx={}, params_len={}",
+                    query_idx,
+                    project_idx,
+                    parameters.len()
+                )));
             }
+            call.query = parameters[query_idx].clone();
+            call.project_name = parameters[project_idx].clone();
         }
     }
     Ok(embed_calls)