sharding: complete tenant-routed data and recovery paths

Author: NasitSony

src/main/java/com/veriprotocol/springAI/core/DocumentService.java

@@ -219,6 +219,7 @@ public void afterCommit() {
                                 );
 
                                 markPublishFailed(
+                                        tenantId,
                                         docId,
                                         safeMsg(ex)
                                 );
@@ -254,8 +255,16 @@ public void afterCommit() {
 
         return existing.id();
     }
-    private void markPublishFailed(String docId, String err) {
-        docRepo.updateLastError(docId, "PUBLISH_FAILED: " + err);
+    private void markPublishFailed(
+            String tenantId,
+            String docId,
+            String err) {
+
+        shardedDocumentWriteDao.updateLastError(
+                tenantId,
+                docId,
+                "PUBLISH_FAILED: " + err
+        );
     }
 
     public List<ChunkSearchDao.ChunkHit> semanticSearchChunks(
@@ -274,19 +283,9 @@ public List<ChunkSearchDao.ChunkHit> semanticSearchChunks(
         );
     }
 
-    public boolean claimProcessingLease(String docId, String workerId) {
-       return documentWriteDao.claimProcessingLease(docId, workerId);
-    }
 
-    public void markReadyDb(String docId) {
-        int updated = documentWriteDao.markReady(docId);
 
-        if (updated != 1) {
-            throw new IllegalStateException(
-                    "Invalid READY transition for docId=" + docId
-            );
-        }
-    }
+
 
 
     public int markRetryCycleExhausted(

src/main/java/com/veriprotocol/springAI/core/ShardedDocumentWriteDao.java

@@ -76,6 +76,27 @@ public Optional<ExistingDocument> findByTenantAndRequestId(
         return rows.stream().findFirst();
     }
 
+    public int updateLastError(
+            String tenantId,
+            String docId,
+            String error) {
+
+        JdbcTemplate jdbc =
+                shardJdbcRouter.jdbcFor(tenantId);
+
+        return jdbc.update("""
+        UPDATE documents
+        SET last_error = ?,
+            updated_at = now()
+        WHERE tenant_id = ?
+          AND id = ?
+        """,
+                error,
+                tenantId,
+                docId
+        );
+    }
+
     public record ExistingDocument(
             String id,
             String contentHash) {}

src/main/java/com/veriprotocol/springAI/persistence/DocumentReadDao.java

@@ -47,23 +47,7 @@ public DocumentStatusDto mapRow(ResultSet rs, int rowNum) throws SQLException {
 
     // Internal worker lookup.
 // Safe because document IDs are globally unique and this is not user-facing.
-    public Optional<DocumentStatusDto> findStatusById(String id) {
 
-        var list = jdbcTemplate.query("""
-        SELECT id, status,
-               COALESCE(retry_count, 0) AS retry_count,
-               created_at, updated_at,
-               last_error, worker_id,
-               processing_started_at, next_retry_at
-        FROM documents
-        WHERE id = ?
-        """,
-                MAPPER,
-                id
-        );
-
-        return list.stream().findFirst();
-    }
     public Optional<DocumentStatusDto> findStatusByTenantAndId(
             String tenantId,
             String id) {
@@ -89,59 +73,17 @@ public Optional<DocumentStatusDto> findStatusByTenantAndId(
         return list.stream().findFirst();
     }
 
-    public List<RetryableDoc> findRetryableFailedDocs(
-            int limit,
-            int maxRetries) {
 
-        return jdbcTemplate.query("""
-        SELECT tenant_id, id
-        FROM documents
-        WHERE status = 'FAILED'
-          AND retry_count < ?
-          AND (next_retry_at IS NULL OR next_retry_at <= now())
-        ORDER BY updated_at ASC
-        LIMIT ?
-        """,
-                (rs, rowNum) -> new RetryableDoc(
-                        rs.getString("tenant_id"),
-                        rs.getString("id")
-                ),
-                maxRetries,
-                limit
-        );
-    }
     public record RetryableDoc(
             String tenantId,
             String docId
     ) {}
 
-    public int resetFailedToPending(String docId) {
-        return jdbcTemplate.update("""
-            UPDATE documents
-            SET status = 'PENDING',
-                last_error = NULL,
-                processing_started_at = NULL,
-                worker_id = NULL,
-                updated_at = now()
-            WHERE id = ?
-              AND status = 'FAILED'
-        """, docId);
-    }
+
 
     public record DocPayload(String id, String tenantId, String text, String contentHash) {}
 
-    public DocPayload loadDocPayload(String id) {
-        return jdbcTemplate.queryForObject("""
-            SELECT id, tenant_id, text, content_hash
-                        FROM documents
-                        WHERE id = ?
-        """, (rs, rowNum) -> new DocPayload(
-                rs.getString("id"),
-                rs.getString("tenant_id"),
-                rs.getString("text"),
-                rs.getString("content_hash")
-        ), id);
-    }
+
 
     public List<String> findOldPendingDocIdsForTenant(
             String tenantId,

src/main/java/com/veriprotocol/springAI/persistence/DocumentSearchDao.java

@@ -78,104 +78,6 @@ public int updateStatusAndError(String docId, DocumentStatus status, String msg)
      }
 
 
-   /* public boolean claimProcessingLease(String docId, String workerId) {
-      var rows = jdbcTemplate.queryForList("""
-        INSERT INTO documents (id, status, worker_id, processing_started_at, retry_count, updated_at)
-        VALUES (?, 'PROCESSING', ?, now(), 0, now())
-        ON CONFLICT (id) DO UPDATE
-        SET status = 'PROCESSING',
-            worker_id = EXCLUDED.worker_id,
-            processing_started_at = now(),
-            updated_at = now()
-        WHERE documents.status IN ('PENDING', 'FAILED')
-        RETURNING id
-    """, docId, workerId);
-
-    return !rows.isEmpty();
-    }*/
-
-
-    public boolean claimProcessingLease(String docId, String workerId) {
-
-        var rows = jdbcTemplate.queryForList("""
-        UPDATE documents
-        SET status = 'PROCESSING',
-            worker_id = ?,
-            processing_started_at =
-                CASE
-                    WHEN status = 'PROCESSING'
-                    THEN processing_started_at
-                    ELSE now()
-                END,
-            updated_at = now()
-        WHERE id = ?
-          AND (
-                status IN ('PENDING', 'FAILED')
-                OR
-                (status = 'PROCESSING' AND worker_id = ?)
-              )
-        RETURNING id
-        """,
-                String.class,
-                workerId,
-                docId,
-                workerId
-        );
-
-        return !rows.isEmpty();
-    }
-
-    public int markReady(String docId) {
-        return jdbcTemplate.update("""
-        UPDATE documents
-        SET status = 'READY',
-            last_error = NULL,
-            next_retry_at = NULL,
-            processing_started_at = NULL,
-            worker_id = NULL,
-            updated_at = now()
-        WHERE id = ?
-          AND status = 'PROCESSING'
-        """, docId);
-    }
-
-    public String insertPendingIfAbsent(
-            String id,
-            String tenantId,
-            String requestId,
-            String text,
-            String contentHash) {
-
-        var ids = jdbcTemplate.queryForList("""
-    INSERT INTO documents (
-        id,
-        tenant_id,
-        request_id,
-        text,
-        content_hash,
-        status,
-        retry_count,
-        created_at,
-        updated_at
-    )
-    VALUES (?, ?, ?, ?, ?, 'PENDING', 0, now(), now())
-    ON CONFLICT (tenant_id, request_id) DO NOTHING
-    RETURNING id
-    """,
-                String.class,
-                id,
-                tenantId,
-                requestId,
-                text,
-                contentHash
-        );
-
-        return ids.isEmpty() ? null : ids.get(0);
-    }
-
-
-
-
 
 
     public int markStuckProcessingAsFailed(int minutes) {
@@ -192,23 +94,7 @@ public int markStuckProcessingAsFailed(int minutes) {
               AND processing_started_at < now() - (? || ' minutes')::interval
         """, minutes);
     }
-
-
-
-    public int resetFailedToPending(String docId) {
-        return jdbcTemplate.update("""
-            UPDATE documents
-            SET status = 'PENDING',
-              last_error = NULL,
-              next_retry_at = NULL,
-              processing_started_at = NULL,
-              worker_id = NULL,
-              updated_at = now()
-            WHERE id = ?
-            AND status = 'FAILED'
-            AND (next_retry_at IS NULL OR next_retry_at <= now())
-           """, docId);
-    }
+    
 
 
     public int forceToPending(