chore: address review feedback

Author: JamesGuthrie

projects/pgai/db/sql/idempotent/011-vectorizer-int.sql

@@ -507,6 +507,7 @@ begin
       ( %s
       , created_at pg_catalog.timestamptz not null default now()
       , failure_step pg_catalog.text not null default ''
+      , attempts pg_catalog.int4 not null default 0
       )
       $sql$
     , queue_schema, queue_failed_table
@@ -1299,6 +1300,7 @@ begin
     ) as _;
 
     -- TODO: for very small batch sizes (<10), an array _may_ be faster
+    drop table if exists seen_lock_ids;
     create temporary table seen_lock_ids (lock_id bigint);
     create index on seen_lock_ids (lock_id);
 

projects/pgai/db/sql/incremental/032-split-transaction-support.sql

@@ -8,5 +8,9 @@ begin
         execute format('alter table %I.%I rename column loading_retries to attempts', _vectorizer.queue_schema, _vectorizer.queue_table);
         execute format('alter table %I.%I rename column loading_retry_after to retry_after', _vectorizer.queue_schema, _vectorizer.queue_table);
     end loop;
+    for _vectorizer in select queue_schema, queue_failed_table from ai.vectorizer
+    loop
+        execute format('alter table %I.%I add column attempts pg_catalog.int4 not null default 0', _vectorizer.queue_schema, _vectorizer.queue_failed_table);
+    end loop;
 end;
 $block$;

projects/pgai/db/tests/golden/failed-queue-table-16.expected

@@ -0,0 +1,11 @@
+Table "ai._vectorizer_q_failed_1"
+    Column    |           Type           | Collation | Nullable | Default  | Storage  | Compression | Stats target | Description 
+--------------+--------------------------+-----------+----------+----------+----------+-------------+--------------+-------------
+ title        | text                     |           | not null |          | extended |             |              | 
+ published    | timestamp with time zone |           | not null |          | plain    |             |              | 
+ created_at   | timestamp with time zone |           | not null | now()    | plain    |             |              | 
+ failure_step | text                     |           | not null | ''::text | extended |             |              | 
+ attempts     | integer                  |           | not null | 0        | plain    |             |              | 
+Indexes:
+    "_vectorizer_q_failed_1_title_published_idx" btree (title, published)
+Access method: heap

projects/pgai/db/tests/golden/failed-queue-table-17.expected

@@ -0,0 +1,11 @@
+Table "ai._vectorizer_q_failed_1"
+    Column    |           Type           | Collation | Nullable | Default  | Storage  | Compression | Stats target | Description 
+--------------+--------------------------+-----------+----------+----------+----------+-------------+--------------+-------------
+ title        | text                     |           | not null |          | extended |             |              | 
+ published    | timestamp with time zone |           | not null |          | plain    |             |              | 
+ created_at   | timestamp with time zone |           | not null | now()    | plain    |             |              | 
+ failure_step | text                     |           | not null | ''::text | extended |             |              | 
+ attempts     | integer                  |           | not null | 0        | plain    |             |              | 
+Indexes:
+    "_vectorizer_q_failed_1_title_published_idx" btree (title, published)
+Access method: heap

projects/pgai/db/tests/vectorizer/test_vectorizer.py

@@ -524,6 +524,10 @@ def test_vectorizer_timescaledb():
     actual = psql_cmd(r"\d+ ai._vectorizer_q_1")
     golden_check("queue-table", actual)
 
+    # does the queue failed table look right?
+    actual = psql_cmd(r"\d+ ai._vectorizer_q_failed_1")
+    golden_check("failed-queue-table", actual)
+
     # does the view look right?
     actual = psql_cmd(r"\d+ website.blog_embedding")
     golden_check("view", actual)

projects/pgai/pgai/data/ai.sql

@@ -1243,6 +1243,10 @@ begin
         execute format('alter table %I.%I rename column loading_retries to attempts', _vectorizer.queue_schema, _vectorizer.queue_table);
         execute format('alter table %I.%I rename column loading_retry_after to retry_after', _vectorizer.queue_schema, _vectorizer.queue_table);
     end loop;
+    for _vectorizer in select queue_schema, queue_failed_table from ai.vectorizer
+    loop
+        execute format('alter table %I.%I add column attempts pg_catalog.int4 not null default 0', _vectorizer.queue_schema, _vectorizer.queue_failed_table);
+    end loop;
 end;
 $block$;
 
@@ -2826,6 +2830,7 @@ begin
       ( %s
       , created_at pg_catalog.timestamptz not null default now()
       , failure_step pg_catalog.text not null default ''
+      , attempts pg_catalog.int4 not null default 0
       )
       $sql$
     , queue_schema, queue_failed_table
@@ -3618,6 +3623,7 @@ begin
     ) as _;
 
     -- TODO: for very small batch sizes (<10), an array _may_ be faster
+    drop table if exists seen_lock_ids;
     create temporary table seen_lock_ids (lock_id bigint);
     create index on seen_lock_ids (lock_id);
 

projects/pgai/pgai/vectorizer/vectorizer.py

@@ -635,8 +635,8 @@ def fetch_work_query_with_split_transaction(self) -> sql.Composed:
                 ),
                 delete_duplicate_rows AS (
                    DELETE FROM {queue_table} q
-                   WHERE ({pk_fields}) IN (SELECT {pk_fields} FROM locked_queue_rows)
-                     AND q.ctid NOT IN (SELECT _ctid FROM locked_queue_rows)
+                   WHERE ({pk_fields}) IN (SELECT {pk_fields} FROM update_locked_rows)
+                     AND q.ctid NOT IN (SELECT _ctid FROM update_locked_rows)
                 )
                 SELECT s.*, {attempts}
                 FROM update_locked_rows l
@@ -846,10 +846,10 @@ def update_retry_after_and_move_permanent_failures_to_dlq_query(
                 DELETE FROM {queue_table}
                 WHERE attempts > {max_attempts}
                   AND ({pk_fields}) IN (SELECT {pk_fields} FROM erroring_pk_and_step)
-                RETURNING {pk_fields}
+                RETURNING {pk_fields}, attempts
             )
-            INSERT INTO {queue_failed_table} ({pk_fields}, failure_step)
-                SELECT {pk_fields}, failure_step FROM delete_over_attempts JOIN erroring_pk_and_step USING ({pk_fields})
+            INSERT INTO {queue_failed_table} ({pk_fields}, failure_step, attempts)
+                SELECT {pk_fields}, failure_step, attempts FROM delete_over_attempts JOIN erroring_pk_and_step USING ({pk_fields})
             """).format(
             error_values=self._placeholders_tuples(
                 len(self.vectorizer.source_pk) + 1, len(items)

projects/pgai/tests/vectorizer/cli/test_vectorizer_split_transaction_semantics.py

@@ -346,6 +346,7 @@ async def test_vectorizer_error_handling_moves_failed_items_to_dlq_and_logs_erro
                 "id": 1,
                 "failure_step": "loading",
                 "created_at": now,
+                "attempts": 7,
             },
         )