diff --git a/CLAUDE.md b/CLAUDE.md index 48256769..fe980cae 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -286,6 +286,7 @@ com.bablsoft.accessflow/ | `ACCESSFLOW_PROXY_IDLE_TIMEOUT` | HikariCP `idleTimeout` (default `10m`). | | `ACCESSFLOW_PROXY_MAX_LIFETIME` | HikariCP `maxLifetime` (default `30m`). | | `ACCESSFLOW_PROXY_LEAK_DETECTION_THRESHOLD` | HikariCP leak-detection threshold (default `0s` = disabled). | +| `ACCESSFLOW_PROXY_ESTIMATE_TIMEOUT` | ISO-8601 duration. Statement timeout for the automatic pre-flight cost estimate (AF-624) — the async EXPLAIN + affected-row COUNT run after each submission use this tighter bound instead of the full execution statement timeout (default `PT5S`). | | `ACCESSFLOW_PROXY_EXECUTION_MAX_ROWS` | Hard cap on rows returned by a single query (default `10000`). | | `ACCESSFLOW_PROXY_EXECUTION_STATEMENT_TIMEOUT` | Statement-level timeout applied to customer-DB JDBC statements (default `30s`). | | `ACCESSFLOW_PROXY_EXECUTION_DEFAULT_FETCH_SIZE` | Default JDBC fetch size (default `1000`). | @@ -861,18 +862,20 @@ The `AiAnalyzerStrategy` interface must be implemented by all three adapters: ```java public interface AiAnalyzerStrategy { AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, - String language, UUID organizationId); + String costEstimateContext, String language, UUID aiConfigId); } ``` `schemaContext` may be `null` or empty when introspection is unavailable; the prompt template -substitutes `(no schema introspection available)` in that case. +substitutes `(no schema introspection available)` in that case. `costEstimateContext` (AF-624) +carries the same contract for the pre-flight cost estimate — `null`/empty substitutes +`(no cost estimate available)`. - Adapters route their HTTP calls through **Spring AI 2.0** (`spring-ai-bom:2.0.0`). The autowired `AiAnalyzerStrategy` is `AiAnalyzerStrategyHolder` — it builds `AnthropicChatModel` / `OpenAiChatModel` / `OllamaChatModel` per-org from the `ai_config` row, caches the delegate, and evicts on `AiConfigUpdatedEvent` (no Spring context refresh, no restart). - The same Spring AI BOM also pins `spring-ai-starter-mcp-server-webmvc` (`2.0.0`) — the stateless MCP server starter used by the `mcp` module. The starter ships with `spring-ai-autoconfigure-mcp-server-common` as an explicit dependency in `backend/pom.xml` to work around a missing transitive in the starter POM (`McpServerStdioDisabledCondition` lives in the common artifact). See `docs/13-mcp.md` and `docs/05-backend.md` → "MCP server and user API keys". - Active provider per org is the `ai_config.provider` column. There is no `accessflow.ai.provider` property and no `@ConditionalOnProperty` on the strategy classes — they are plain classes, not Spring beans. - Connection settings (API key, base URL, model, max-tokens, timeout) come from `ai_config`, not from `spring.ai.*` properties. `application.yml` sets `spring.ai.model.{chat,embedding,image,audio.speech,audio.transcription,moderation}=none` so no `ChatModel` is auto-built at startup. -- The default system prompt template (`SystemPromptRenderer.DEFAULT_TEMPLATE`) is in `docs/05-backend.md` — use it verbatim; do not invent a different prompt. It uses named placeholders `{{db_type}}`, `{{schema_context}}`, `{{sql}}`, `{{language}}` (substituted at render time; `{{sql}}` replaced last). Admins may override the prompt **per `ai_config` row** via the nullable `system_prompt_template` column (blank ⇒ default); a custom template **must contain `{{sql}}`** or the service throws `AiConfigInvalidPromptException` (HTTP 400 `AI_CONFIG_INVALID_PROMPT`). The default is served to the admin UI via `GET /admin/ai-configs/prompt-default`. +- The default system prompt template (`SystemPromptRenderer.DEFAULT_TEMPLATE`) is in `docs/05-backend.md` — use it verbatim; do not invent a different prompt. It uses named placeholders `{{db_type}}`, `{{schema_context}}`, `{{rag_context}}`, `{{cost_estimate}}` (AF-624 — the pre-flight estimate summary), `{{sql}}`, `{{language}}` (substituted at render time; `{{sql}}` replaced last). Admins may override the prompt **per `ai_config` row** via the nullable `system_prompt_template` column (blank ⇒ default); a custom template **must contain `{{sql}}`** or the service throws `AiConfigInvalidPromptException` (HTTP 400 `AI_CONFIG_INVALID_PROMPT`). The default is served to the admin UI via `GET /admin/ai-configs/prompt-default`. - AI calls are asynchronous — publish a `QuerySubmittedEvent` and handle in the strategy asynchronously using virtual threads. - The response must be parsed strictly as JSON matching the `AiAnalysisResult` schema. If the AI returns non-JSON or an unexpected schema, log and mark the analysis as failed; do not propagate the exception to the query request. - For Anthropic: use `claude-sonnet-4-20250514` as the default model. diff --git a/README.md b/README.md index c877428b..325682e6 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ A glance at the day-to-day flows engineers and approvers actually use. - **Multi-model orchestration, voting & guardrails** — an AI configuration can run several models in parallel (the primary plus weighted `{provider, model, weight}` members — e.g. a fast local Ollama pre-score alongside a deep Claude analysis) and combine their verdicts by a configurable voting strategy (**weighted average**, **highest risk**, or **majority**); issues are merged and risk aggregated. **Guardrails** block configured prompt patterns (case-insensitive regex) before any model is called, and the strict output-schema validation rejects malformed responses without failing the query. Per-model **cost (tokens) and latency** are recorded for every analysis and charted on the admin dashboard. - **Query-optimization suggestions** — alongside the risk verdict, the analyzer returns concrete, dialect-aware optimization suggestions: index recommendations (`CREATE INDEX …`) and query rewrites. Each suggestion has a one-click **"Apply as draft"** that pre-fills the editor with the suggested statement and routes it through the normal review pipeline — nothing auto-executes, and the draft is audited as AI-suggested. - **Query playground / dry-run sandbox** — preview a query's impact before submitting it for review. A **Dry run** action returns the engine's execution plan (node type, estimated rows, cost, filters) and a best-effort estimated row impact **without executing or mutating data**, shown in the editor next to the AI panel. It runs through the same governance as a real execution (datasource access, table allow-list, row-level security) but creates no review. Dialect-aware across every engine with a plan concept — relational `EXPLAIN` (PostgreSQL / MySQL / MariaDB / Oracle / SQL Server), MongoDB `explain`, Couchbase / Neo4j `EXPLAIN`, Elasticsearch / OpenSearch query validation; engines without one degrade gracefully. SELECT dry-runs prefer the read replica when configured. +- **Pre-flight cost & blast-radius estimation (AF-624)** — every submitted query automatically gets a persisted cost estimate before review: the engine's own `EXPLAIN` plan (estimated rows, scan type, cost) plus, for UPDATE/DELETE, an **exact affected-row count** computed with a governed, non-mutating count (relational `SELECT COUNT(*)` rewrite; MongoDB `countDocuments`, SQL++ `COUNT(*)`, Cypher `count(*)`, Elasticsearch `_count`). Reviewers see it on the query detail page ("this DELETE touches ~2.4M rows via Seq Scan"), routing policies can match on `estimated_rows` / `scan_type` (e.g. auto-escalate full-scan DELETEs over 100k rows), and the estimate is folded into the AI analyzer's prompt so risk scoring reflects the actual blast radius. Engines without a plan concept degrade gracefully to an "estimate unavailable" state that never blocks submission. - **Text-to-query** — opt-in per datasource. Users describe what they want in plain language ("order numbers for the last 5 days") and the AI drafts a schema-grounded query into the editor — in the datasource engine's **native query language** (SQL for relational engines, plus MongoDB shell/JSON, Cypher, CQL, Elasticsearch Query DSL, redis-cli, SQL++, and PartiQL for the NoSQL engines). Reuses the datasource's AI configuration and is grounded in its introspected schema (restricted columns are never referenced). The generated query is only a draft — it's still submitted through the normal pipeline, so AI risk analysis and human review always apply. - **RAG knowledge base** — opt-in per AI configuration. Attach knowledge documents (data-governance policies, naming conventions, schema notes); AccessFlow embeds them with a dedicated embedding model and, at analysis / text-to-SQL time, retrieves the most relevant chunks and injects them into the prompt — so analysis reflects your house rules. Pluggable vector store: in-app **pgvector** (self-contained, auto-provisioned where the DB role permits — and optional: AccessFlow still starts if the extension is absent, with the in-app store disabled) or external **Qdrant**. Retrieval is best-effort and never blocks analysis. - **Langfuse integration** — optional, per-organization. Send a trace of every AI analysis (input SQL, structured output, model, token usage, latency) to [Langfuse](https://langfuse.com) for LLM observability, and/or fetch analyzer prompts from Langfuse by name at runtime (prompt management) so you can iterate on prompts without redeploying. Best-effort and non-blocking — a Langfuse outage never affects query workflow. diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/api/AiAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/api/AiAnalyzerStrategy.java index a363df3c..75f457a3 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/api/AiAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/api/AiAnalyzerStrategy.java @@ -16,7 +16,9 @@ public interface AiAnalyzerStrategy { /** * Analyze {@code sql} for risks, returning a structured result. The {@code schemaContext} * argument is an opaque, provider-renderable description of the target schema (or - * {@code null}/empty if introspection is unavailable). + * {@code null}/empty if introspection is unavailable). {@code costEstimateContext} (AF-624) + * carries the same contract for the pre-flight cost / blast-radius estimate — an opaque, + * provider-renderable summary (or {@code null}/empty when no estimate is available). * *

{@code language} is a BCP-47 code (e.g. {@code "en"}, {@code "es"}, {@code "zh-CN"}) that * tells the AI which language to use for free-form fields. When {@code null} or unrecognised @@ -30,8 +32,14 @@ public interface AiAnalyzerStrategy { * @throws com.bablsoft.accessflow.ai.api.AiAnalysisParseException provider response did not * match the expected schema */ - AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId); + AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId); + + /** Convenience overload for callers with no cost-estimate context. */ + default AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String language, UUID aiConfigId) { + return analyze(sql, dbType, schemaContext, null, language, aiConfigId); + } /** * Translate the natural-language {@code prompt} into a single SQL statement for the target diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolder.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolder.java index 8a33b36f..5e8caa28 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolder.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolder.java @@ -85,18 +85,19 @@ class AiAnalyzerStrategyHolder implements AiAnalyzerStrategy { what a reviewer should check. Do not use markdown, headings, or bullet points."""; @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { if (aiConfigId == null) { throw notConfigured(); } try { - return delegateFor(aiConfigId).analyze(sql, dbType, schemaContext, language, aiConfigId); + return delegateFor(aiConfigId).analyze(sql, dbType, schemaContext, costEstimateContext, + language, aiConfigId); } catch (AiAnalysisException | AiAnalysisParseException primary) { rethrowIfNotFallbackEligible(primary); return failover(aiConfigId, primary, fallback -> - delegateFor(fallback.getId()).analyze(sql, dbType, schemaContext, language, - fallback.getId())); + delegateFor(fallback.getId()).analyze(sql, dbType, schemaContext, + costEstimateContext, language, fallback.getId())); } } diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AnthropicAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AnthropicAnalyzerStrategy.java index 0ec6da80..12060203 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AnthropicAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/AnthropicAnalyzerStrategy.java @@ -42,11 +42,11 @@ class AnthropicAnalyzerStrategy implements AiAnalyzerStrategy { private final RagRetriever ragRetriever; @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { var ragContext = ragRetriever.retrieve(sql); var userPrompt = promptRenderer.render(promptSource.template(), sql, dbType, schemaContext, - ragContext, language); + ragContext, costEstimateContext, language); var prompt = new Prompt(List.of( new SystemMessage(SYSTEM_PROMPT_PREAMBLE), new UserMessage(userPrompt))); diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerService.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerService.java index a41843bd..7f91db21 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerService.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerService.java @@ -27,6 +27,8 @@ import com.bablsoft.accessflow.core.events.AiAnalysisCompletedEvent; import com.bablsoft.accessflow.core.events.AiAnalysisFailedEvent; import com.bablsoft.accessflow.core.events.AiAnalysisSkippedEvent; +import com.bablsoft.accessflow.core.api.QueryEstimateSnapshot; +import com.bablsoft.accessflow.proxy.api.QueryCostEstimateService; import com.bablsoft.accessflow.proxy.api.SqlParserService; import io.micrometer.core.instrument.MeterRegistry; import io.micrometer.observation.Observation; @@ -69,6 +71,7 @@ class DefaultAiAnalyzerService implements AiAnalyzerService { private final ApplicationEventPublisher eventPublisher; private final AiRateLimiter aiRateLimiter; private final ObservationRegistry observationRegistry; + private final QueryCostEstimateService queryCostEstimateService; private final MeterRegistry meterRegistry; @Override @@ -143,9 +146,14 @@ public void analyzeSubmittedQuery(UUID queryRequestId) { } catch (RuntimeException e) { log.warn("Schema introspection failed for query {}: {}", queryRequestId, e.getMessage()); } + // AF-624: ensure the pre-flight cost estimate exists (idempotent — the proxy listener may + // already have computed it) and fold it into the prompt so risk scoring sees the actual + // blast radius. Best-effort: an estimate failure never blocks analysis. + String costEstimateContext = buildCostEstimateContext(queryRequestId); try { aiRateLimiter.enforce(snapshot.organizationId()); - var analysis = analyzeWithObservation(snapshot, descriptor, schemaContext); + var analysis = analyzeWithObservation(snapshot, descriptor, schemaContext, + costEstimateContext); var result = ClassificationRiskBooster.boost(analysis, ClassificationRiskBooster.bumpFor( referencedClassifications(snapshot.sqlText(), classificationTags))); var issuesJson = responseParser.issuesAsJson(result.issues()); @@ -180,13 +188,15 @@ public void analyzeSubmittedQuery(UUID queryRequestId) { */ private AiAnalysisResult analyzeWithObservation(QueryRequestSnapshot snapshot, DatasourceConnectionDescriptor descriptor, - String schemaContext) { + String schemaContext, + String costEstimateContext) { // Each tag is set exactly once per branch (no re-add) so the meter series carry the right // provider/risk on success and consistent keys on failure. Observation observation = Observation.start("accessflow.ai.analyze", observationRegistry); try (Observation.Scope ignored = observation.openScope()) { var analysis = strategy.analyze(snapshot.sqlText(), descriptor.dbType(), schemaContext, - resolveLanguage(snapshot.organizationId()), descriptor.aiConfigId()); + costEstimateContext, resolveLanguage(snapshot.organizationId()), + descriptor.aiConfigId()); observation.lowCardinalityKeyValue("provider", providerTag(analysis.aiProvider())) .lowCardinalityKeyValue("risk_level", analysis.riskLevel() != null ? analysis.riskLevel().name() : "none") @@ -243,6 +253,52 @@ private UUID requireBoundAiConfig(DatasourceConnectionDescriptor descriptor) { return aiConfigId; } + /** + * Best-effort AF-624 prompt context: triggers/reads the persisted estimate and renders it as a + * short plain-text summary. Null (→ the template's "(no cost estimate available)" fallback) + * when the estimate is absent, unsupported, or failed. + */ + private String buildCostEstimateContext(UUID queryRequestId) { + QueryEstimateSnapshot estimate; + try { + estimate = queryCostEstimateService.estimateSubmittedQuery(queryRequestId).orElse(null); + } catch (RuntimeException e) { + log.warn("Cost estimate unavailable for query {}: {}", queryRequestId, e.getMessage()); + return null; + } + if (estimate == null || estimate.failed() || !estimate.supported()) { + return null; + } + var sb = new StringBuilder(); + if (estimate.affectedRowCount() != null) { + sb.append("Exact affected-row count for this ").append(estimate.queryType()) + .append(" (governed SELECT COUNT(*)): ") + .append(String.format(Locale.ROOT, "%,d", estimate.affectedRowCount())) + .append(" rows."); + } + if (estimate.estimatedRows() != null) { + if (!sb.isEmpty()) { + sb.append(' '); + } + sb.append("The engine's execution plan estimates ~") + .append(String.format(Locale.ROOT, "%,d", estimate.estimatedRows())) + .append(" rows"); + if (estimate.scanType() != null) { + sb.append(" via ").append(estimate.scanType()); + } + if (estimate.estimatedCost() != null) { + sb.append(" (cost ").append(estimate.estimatedCost()).append(')'); + } + sb.append('.'); + } else if (estimate.scanType() != null) { + if (!sb.isEmpty()) { + sb.append(' '); + } + sb.append("Plan root operation: ").append(estimate.scanType()).append('.'); + } + return sb.isEmpty() ? null : sb.toString(); + } + private String resolveLanguage(UUID organizationId) { if (organizationId == null) { return SupportedLanguage.EN.code(); diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategy.java index f4b980c7..1430d194 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategy.java @@ -37,10 +37,11 @@ class GuardrailAiAnalyzerStrategy implements AiAnalyzerStrategy { } @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { guard(sql); - return delegate.analyze(sql, dbType, schemaContext, language, aiConfigId); + return delegate.analyze(sql, dbType, schemaContext, costEstimateContext, language, + aiConfigId); } @Override diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OllamaAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OllamaAnalyzerStrategy.java index c907578e..396456d6 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OllamaAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OllamaAnalyzerStrategy.java @@ -42,11 +42,11 @@ class OllamaAnalyzerStrategy implements AiAnalyzerStrategy { private final RagRetriever ragRetriever; @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { var ragContext = ragRetriever.retrieve(sql); var userPrompt = promptRenderer.render(promptSource.template(), sql, dbType, schemaContext, - ragContext, language); + ragContext, costEstimateContext, language); var prompt = new Prompt(List.of( new SystemMessage(SYSTEM_PROMPT_PREAMBLE), new UserMessage(userPrompt))); diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OpenAiAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OpenAiAnalyzerStrategy.java index f09062d0..dc8715d8 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OpenAiAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OpenAiAnalyzerStrategy.java @@ -46,11 +46,11 @@ class OpenAiAnalyzerStrategy implements AiAnalyzerStrategy { private final RagRetriever ragRetriever; @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { var ragContext = ragRetriever.retrieve(sql); var userPrompt = promptRenderer.render(promptSource.template(), sql, dbType, schemaContext, - ragContext, language); + ragContext, costEstimateContext, language); var prompt = new Prompt(List.of( new SystemMessage(SYSTEM_PROMPT_PREAMBLE), new UserMessage(userPrompt))); diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategy.java index 599e6a2e..ecab6a67 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategy.java @@ -54,9 +54,9 @@ record Member(AiAnalyzerStrategy strategy, AiProviderType provider, String model } @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { - var outcomes = invokeAll(sql, dbType, schemaContext, language, aiConfigId); + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { + var outcomes = invokeAll(sql, dbType, schemaContext, costEstimateContext, language, aiConfigId); var successes = new ArrayList(); var breakdown = new ArrayList(outcomes.size()); @@ -104,15 +104,17 @@ public GeneratedSqlResult generateSql(String prompt, DbType dbType, String schem } private List invokeAll(String sql, DbType dbType, String schemaContext, - String language, UUID aiConfigId) { + String costEstimateContext, String language, UUID aiConfigId) { if (members.size() == 1) { - return List.of(runMember(members.get(0), sql, dbType, schemaContext, language, aiConfigId)); + return List.of(runMember(members.get(0), sql, dbType, schemaContext, costEstimateContext, + language, aiConfigId)); } try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { var futures = new ArrayList>(members.size()); for (var member : members) { Callable task = - () -> runMember(member, sql, dbType, schemaContext, language, aiConfigId); + () -> runMember(member, sql, dbType, schemaContext, costEstimateContext, + language, aiConfigId); futures.add(executor.submit(task)); } var outcomes = new ArrayList(members.size()); @@ -124,10 +126,11 @@ private List invokeAll(String sql, DbType dbType, String schemaCo } private MemberOutcome runMember(Member member, String sql, DbType dbType, String schemaContext, - String language, UUID aiConfigId) { + String costEstimateContext, String language, UUID aiConfigId) { var start = clock.instant(); try { - var result = member.strategy().analyze(sql, dbType, schemaContext, language, aiConfigId); + var result = member.strategy().analyze(sql, dbType, schemaContext, costEstimateContext, + language, aiConfigId); return MemberOutcome.success(member, result, elapsedMs(start)); } catch (RuntimeException e) { log.warn("Orchestration member {}/{} failed: {}", member.provider(), member.model(), diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/SystemPromptRenderer.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/SystemPromptRenderer.java index 5b5c302e..c9ca2c03 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/SystemPromptRenderer.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/SystemPromptRenderer.java @@ -57,6 +57,8 @@ class SystemPromptRenderer { Database type: {{db_type}} Schema context: {{schema_context}} + Pre-flight cost estimate (from the database engine's own EXPLAIN / affected-row count — treat it as the authoritative blast radius and factor it into risk_score and risk_level): + {{cost_estimate}} Knowledge base context (authoritative organization-specific guidance retrieved for this query — prefer it over general assumptions when it applies): {{rag_context}} SQL to analyze: @@ -75,6 +77,11 @@ Knowledge base context (authoritative organization-specific guidance retrieved f private static final String NO_RAG_CONTEXT = "(no knowledge base context available)"; + /** The pre-flight cost-estimate token (AF-624); substituted with the summary or a fallback. */ + static final String COST_ESTIMATE_PLACEHOLDER = "{{cost_estimate}}"; + + private static final String NO_COST_ESTIMATE = "(no cost estimate available)"; + /** The token naming the engine's target query language (e.g. SQL, Cypher, CQL). */ static final String TARGET_LANGUAGE_PLACEHOLDER = "{{target_language}}"; @@ -211,11 +218,16 @@ String defaultTemplate() { } String render(String sql, DbType dbType, String schemaContext, String language) { - return render(null, sql, dbType, schemaContext, null, language); + return render(null, sql, dbType, schemaContext, null, null, language); + } + + String render(String template, String sql, DbType dbType, String schemaContext, + String ragContext, String language) { + return render(template, sql, dbType, schemaContext, ragContext, null, language); } String render(String template, String sql, DbType dbType, String schemaContext, String ragContext, - String language) { + String costEstimateContext, String language) { var effective = (template == null || template.isBlank()) ? DEFAULT_TEMPLATE : template; var schemaText = (schemaContext == null || schemaContext.isBlank()) ? "(no schema introspection available)" @@ -229,6 +241,7 @@ String render(String template, String sql, DbType dbType, String schemaContext, .replace("{{db_type}}", dbType.name()) .replace("{{schema_context}}", schemaText) .replace(RAG_CONTEXT_PLACEHOLDER, ragText(ragContext)) + .replace(COST_ESTIMATE_PLACEHOLDER, costEstimateText(costEstimateContext)) .replace("{{language}}", displayName) .replace(SQL_PLACEHOLDER, sql); } @@ -278,6 +291,12 @@ private static String ragText(String ragContext) { return (ragContext == null || ragContext.isBlank()) ? NO_RAG_CONTEXT : ragContext; } + private static String costEstimateText(String costEstimateContext) { + return (costEstimateContext == null || costEstimateContext.isBlank()) + ? NO_COST_ESTIMATE + : costEstimateContext; + } + String describeSchema(DatabaseSchemaView schema) { return describeSchema(schema, List.of()); } diff --git a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategy.java b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategy.java index 036c3360..8ba2a31d 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategy.java +++ b/backend/src/main/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategy.java @@ -40,11 +40,12 @@ class TracingAiAnalyzerStrategy implements AiAnalyzerStrategy { } @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { var start = clock.instant(); try { - var result = delegate.analyze(sql, dbType, schemaContext, language, aiConfigId); + var result = delegate.analyze(sql, dbType, schemaContext, costEstimateContext, + language, aiConfigId); safeTrace(() -> tracer.trace(LangfuseTraceContext.success( organizationId, aiConfigId, provider, model, sql, dbType, schemaContext, language, result, start, clock.instant()))); diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/PersistQueryEstimateCommand.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/PersistQueryEstimateCommand.java new file mode 100644 index 00000000..61539632 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/PersistQueryEstimateCommand.java @@ -0,0 +1,22 @@ +package com.bablsoft.accessflow.core.api; + +/** + * Everything needed to persist one {@code query_estimates} row (issue AF-624). Mirrors + * {@link PersistAiAnalysisCommand}: the proxy module computes the estimate and hands the values to + * {@link QueryEstimatePersistenceService}. All value fields are nullable best-effort signals. + */ +public record PersistQueryEstimateCommand( + String engineId, + QueryType queryType, + boolean supported, + Long estimatedRows, + Long affectedRowCount, + String scanType, + Double estimatedCost, + String planJson, + String rawPlan, + String unsupportedReason, + boolean failed, + String errorMessage, + Integer durationMs) { +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResult.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResult.java new file mode 100644 index 00000000..0f172612 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResult.java @@ -0,0 +1,35 @@ +package com.bablsoft.accessflow.core.api; + +import java.time.Duration; + +/** + * Result of a governed affected-row count for an UPDATE/DELETE (issue AF-624): how many rows the + * statement would touch, computed by the engine's native non-mutating count (e.g. MongoDB + * {@code countDocuments}, SQL++ / relational {@code SELECT COUNT(*)}, Cypher + * {@code MATCH … RETURN count(*)}, Elasticsearch {@code _count}) with the request's row-security + * directives applied. {@code supported} is {@code false} for engines with no count concept or + * statement shapes the engine cannot provably count; {@code unsupportedReason} then carries an + * optional engine-supplied explanation (the host localizes a generic one when absent). + */ +public record QueryAffectedRowsResult( + boolean supported, + String engineId, + Long affectedRows, + Duration duration, + String unsupportedReason) { + + /** Degraded result for an engine that cannot count affected rows. */ + public static QueryAffectedRowsResult unsupported(String engineId) { + return unsupported(engineId, null); + } + + /** Degraded result carrying an engine-supplied reason (e.g. "join shapes cannot be counted"). */ + public static QueryAffectedRowsResult unsupported(String engineId, String reason) { + return new QueryAffectedRowsResult(false, engineId, null, Duration.ZERO, reason); + } + + public static QueryAffectedRowsResult of(String engineId, long affectedRows, + Duration duration) { + return new QueryAffectedRowsResult(true, engineId, affectedRows, duration, null); + } +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryDetailView.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryDetailView.java index 77f572f0..ee7ad2aa 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryDetailView.java +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryDetailView.java @@ -22,6 +22,7 @@ public record QueryDetailView( QueryStatus status, String justification, AiAnalysisDetail aiAnalysis, + CostEstimateDetail costEstimate, Long rowsAffected, Integer durationMs, String errorMessage, @@ -34,6 +35,22 @@ public record QueryDetailView( Instant createdAt, Instant updatedAt) { + /** Backward-compatible constructor without the AF-624 cost estimate (defaults to absent). */ + public QueryDetailView(UUID id, UUID datasourceId, String datasourceName, DbType dbType, + UUID organizationId, UUID submittedByUserId, String submittedByEmail, + String submittedByDisplayName, String sqlText, QueryType queryType, + QueryStatus status, String justification, AiAnalysisDetail aiAnalysis, + Long rowsAffected, Integer durationMs, String errorMessage, + UUID previousRunId, UUID approvedByGrantId, String reviewPlanName, + Integer approvalTimeoutHours, List reviewDecisions, + Instant scheduledFor, Instant createdAt, Instant updatedAt) { + this(id, datasourceId, datasourceName, dbType, organizationId, submittedByUserId, + submittedByEmail, submittedByDisplayName, sqlText, queryType, status, justification, + aiAnalysis, null, rowsAffected, durationMs, errorMessage, previousRunId, + approvedByGrantId, reviewPlanName, approvalTimeoutHours, reviewDecisions, + scheduledFor, createdAt, updatedAt); + } + public record AiAnalysisDetail( UUID id, RiskLevel riskLevel, @@ -51,6 +68,24 @@ public record AiAnalysisDetail( String errorMessage) { } + /** The persisted pre-flight cost / blast-radius estimate (AF-624), when computed. */ + public record CostEstimateDetail( + UUID id, + String engineId, + QueryType queryType, + boolean supported, + Long estimatedRows, + Long affectedRowCount, + String scanType, + Double estimatedCost, + String planJson, + String rawPlan, + String unsupportedReason, + boolean failed, + String errorMessage, + Integer durationMs) { + } + public record ReviewDecisionView( UUID id, ReviewerRef reviewer, diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEngine.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEngine.java index b86b320f..1db9f58e 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEngine.java +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEngine.java @@ -78,6 +78,18 @@ default QueryDryRunResult dryRun(QueryEngineDryRunRequest request) { return QueryDryRunResult.unsupported(engineId()); } + /** + * Count the rows the request's UPDATE/DELETE would affect (issue AF-624), applying the + * request's row-security directives, without executing or mutating data — the + * engine's native non-mutating count (e.g. MongoDB {@code countDocuments}, SQL++ + * {@code SELECT COUNT(*)}, Cypher {@code MATCH … RETURN count(*)}, Elasticsearch + * {@code _count}). Engines with no count concept inherit the default, which returns + * {@link QueryAffectedRowsResult#unsupported(String)} so the host degrades gracefully. + */ + default QueryAffectedRowsResult countAffectedRows(QueryEngineDryRunRequest request) { + return QueryAffectedRowsResult.unsupported(engineId()); + } + /** * Short-lived connectivity probe (the engine analogue of the JDBC {@code SELECT 1}). * diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateLookupService.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateLookupService.java new file mode 100644 index 00000000..ab8ae090 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateLookupService.java @@ -0,0 +1,14 @@ +package com.bablsoft.accessflow.core.api; + +import java.util.Optional; +import java.util.UUID; + +/** + * Read-only lookup of a query's persisted pre-flight cost estimate (issue AF-624). Consumed by the + * workflow module's routing-condition context builder and the read API — pure read, no + * computation. + */ +public interface QueryEstimateLookupService { + + Optional findByQueryRequestId(UUID queryRequestId); +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimatePersistenceService.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimatePersistenceService.java new file mode 100644 index 00000000..5589e4e7 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimatePersistenceService.java @@ -0,0 +1,19 @@ +package com.bablsoft.accessflow.core.api; + +import java.util.UUID; + +/** + * Persists pre-flight cost estimates (issue AF-624) into {@code query_estimates} and links the row + * from {@code query_requests.query_estimate_id}, mirroring {@link AiAnalysisPersistenceService}. + */ +public interface QueryEstimatePersistenceService { + + /** + * Inserts the estimate row and sets the {@code query_requests.query_estimate_id} back-pointer + * in the same transaction. When a row already exists for the query request (a concurrent + * trigger raced this one) the existing row's id is returned and nothing is written. + * + * @return the persisted (or pre-existing) estimate row id + */ + UUID persist(UUID queryRequestId, PersistQueryEstimateCommand command); +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateSnapshot.java b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateSnapshot.java new file mode 100644 index 00000000..bd118ddd --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/api/QueryEstimateSnapshot.java @@ -0,0 +1,32 @@ +package com.bablsoft.accessflow.core.api; + +import java.time.Instant; +import java.util.UUID; + +/** + * Cross-module DTO carrying a persisted pre-flight cost / blast-radius estimate (issue AF-624) — + * the {@code query_estimates} row computed asynchronously after submission from the engine's + * non-committing dry-run plus, for UPDATE/DELETE, a governed affected-row count. {@code supported} + * is {@code false} when the engine has no plan concept; {@code failed} marks an unexpected + * computation error. {@code estimatedRows}, {@code affectedRowCount}, {@code scanType} and + * {@code estimatedCost} are individually nullable — best-effort signals filled from what the + * engine's plan exposes. + */ +public record QueryEstimateSnapshot( + UUID id, + UUID queryRequestId, + String engineId, + QueryType queryType, + boolean supported, + Long estimatedRows, + Long affectedRowCount, + String scanType, + Double estimatedCost, + String planJson, + String rawPlan, + String unsupportedReason, + boolean failed, + String errorMessage, + Integer durationMs, + Instant createdAt) { +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateCompletedEvent.java b/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateCompletedEvent.java new file mode 100644 index 00000000..693e857c --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateCompletedEvent.java @@ -0,0 +1,13 @@ +package com.bablsoft.accessflow.core.events; + +import java.util.UUID; + +/** + * Published after a pre-flight cost estimate (issue AF-624) has been computed and persisted for a + * submitted query — on both the supported and the gracefully-degraded (engine has no plan concept) + * paths. Consumed by the realtime module to push the {@code query.estimate_complete} WebSocket + * event so open detail views refetch. + */ +public record QueryEstimateCompletedEvent(UUID queryRequestId, UUID queryEstimateId, + boolean supported) { +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateFailedEvent.java b/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateFailedEvent.java new file mode 100644 index 00000000..408719a3 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/events/QueryEstimateFailedEvent.java @@ -0,0 +1,11 @@ +package com.bablsoft.accessflow.core.events; + +import java.util.UUID; + +/** + * Published when the pre-flight cost estimate (issue AF-624) hit an unexpected error — the + * sentinel {@code query_estimates} row is persisted with {@code failed=true} and this event lets + * the realtime module refresh open detail views so the failure surface renders. + */ +public record QueryEstimateFailedEvent(UUID queryRequestId, String reason) { +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateService.java b/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateService.java new file mode 100644 index 00000000..642303f4 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateService.java @@ -0,0 +1,83 @@ +package com.bablsoft.accessflow.core.internal; + +import com.bablsoft.accessflow.core.api.PersistQueryEstimateCommand; +import com.bablsoft.accessflow.core.api.QueryEstimateLookupService; +import com.bablsoft.accessflow.core.api.QueryEstimatePersistenceService; +import com.bablsoft.accessflow.core.api.QueryEstimateSnapshot; +import com.bablsoft.accessflow.core.internal.persistence.entity.QueryEstimateEntity; +import com.bablsoft.accessflow.core.internal.persistence.repo.QueryEstimateRepository; +import com.bablsoft.accessflow.core.internal.persistence.repo.QueryRequestRepository; +import lombok.RequiredArgsConstructor; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; + +import java.time.Instant; +import java.util.Optional; +import java.util.UUID; + +@Service +@RequiredArgsConstructor +class DefaultQueryEstimateService implements QueryEstimatePersistenceService, + QueryEstimateLookupService { + + private final QueryEstimateRepository queryEstimateRepository; + private final QueryRequestRepository queryRequestRepository; + + @Override + @Transactional + public UUID persist(UUID queryRequestId, PersistQueryEstimateCommand command) { + var existing = queryEstimateRepository.findByQueryRequestId(queryRequestId); + if (existing.isPresent()) { + return existing.get().getId(); + } + var queryRequest = queryRequestRepository.findById(queryRequestId) + .orElseThrow(() -> new IllegalStateException( + "Query request not found: " + queryRequestId)); + var entity = new QueryEstimateEntity(); + entity.setId(UUID.randomUUID()); + entity.setQueryRequest(queryRequest); + entity.setEngineId(command.engineId()); + entity.setQueryType(command.queryType()); + entity.setSupported(command.supported()); + entity.setEstimatedRows(command.estimatedRows()); + entity.setAffectedRowCount(command.affectedRowCount()); + entity.setScanType(command.scanType()); + entity.setEstimatedCost(command.estimatedCost()); + entity.setPlan(command.planJson()); + entity.setRawPlan(command.rawPlan()); + entity.setUnsupportedReason(command.unsupportedReason()); + entity.setFailed(command.failed()); + entity.setErrorMessage(command.errorMessage()); + entity.setDurationMs(command.durationMs()); + entity.setCreatedAt(Instant.now()); + var saved = queryEstimateRepository.save(entity); + queryRequest.setQueryEstimateId(saved.getId()); + return saved.getId(); + } + + @Override + @Transactional(readOnly = true) + public Optional findByQueryRequestId(UUID queryRequestId) { + return queryEstimateRepository.findByQueryRequestId(queryRequestId).map(this::toSnapshot); + } + + private QueryEstimateSnapshot toSnapshot(QueryEstimateEntity entity) { + return new QueryEstimateSnapshot( + entity.getId(), + entity.getQueryRequest().getId(), + entity.getEngineId(), + entity.getQueryType(), + entity.isSupported(), + entity.getEstimatedRows(), + entity.getAffectedRowCount(), + entity.getScanType(), + entity.getEstimatedCost(), + entity.getPlan(), + entity.getRawPlan(), + entity.getUnsupportedReason(), + entity.isFailed(), + entity.getErrorMessage(), + entity.getDurationMs(), + entity.getCreatedAt()); + } +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryRequestLookupService.java b/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryRequestLookupService.java index 90d028c4..cf593859 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryRequestLookupService.java +++ b/backend/src/main/java/com/bablsoft/accessflow/core/internal/DefaultQueryRequestLookupService.java @@ -10,9 +10,11 @@ import com.bablsoft.accessflow.core.api.QueryRequestSnapshot; import com.bablsoft.accessflow.core.api.QueryStatus; import com.bablsoft.accessflow.core.internal.persistence.entity.AiAnalysisEntity; +import com.bablsoft.accessflow.core.internal.persistence.entity.QueryEstimateEntity; import com.bablsoft.accessflow.core.internal.persistence.entity.QueryRequestEntity; import com.bablsoft.accessflow.core.internal.persistence.entity.ReviewDecisionEntity; import com.bablsoft.accessflow.core.internal.persistence.repo.AiAnalysisRepository; +import com.bablsoft.accessflow.core.internal.persistence.repo.QueryEstimateRepository; import com.bablsoft.accessflow.core.internal.persistence.repo.QueryRequestRepository; import com.bablsoft.accessflow.core.internal.persistence.repo.ReviewDecisionRepository; import lombok.RequiredArgsConstructor; @@ -32,6 +34,7 @@ class DefaultQueryRequestLookupService implements QueryRequestLookupService { private final QueryRequestRepository queryRequestRepository; private final AiAnalysisRepository aiAnalysisRepository; private final ReviewDecisionRepository reviewDecisionRepository; + private final QueryEstimateRepository queryEstimateRepository; @Override @Transactional(readOnly = true) @@ -211,6 +214,9 @@ private QueryDetailView toDetailView(QueryRequestEntity entity) { entity.getStatus(), entity.getJustification(), toAnalysisDetail(aiAnalysis), + toCostEstimateDetail(entity.getQueryEstimateId() != null + ? queryEstimateRepository.findById(entity.getQueryEstimateId()).orElse(null) + : null), entity.getRowsAffected(), entity.getExecutionDurationMs(), entity.getErrorMessage(), @@ -259,6 +265,28 @@ private static QueryDetailView.AiAnalysisDetail toAnalysisDetail(AiAnalysisEntit entity.getErrorMessage()); } + private static QueryDetailView.CostEstimateDetail toCostEstimateDetail( + QueryEstimateEntity entity) { + if (entity == null) { + return null; + } + return new QueryDetailView.CostEstimateDetail( + entity.getId(), + entity.getEngineId(), + entity.getQueryType(), + entity.isSupported(), + entity.getEstimatedRows(), + entity.getAffectedRowCount(), + entity.getScanType(), + entity.getEstimatedCost(), + entity.getPlan(), + entity.getRawPlan(), + entity.getUnsupportedReason(), + entity.isFailed(), + entity.getErrorMessage(), + entity.getDurationMs()); + } + private static QueryRequestSnapshot toSnapshot(QueryRequestEntity entity) { return new QueryRequestSnapshot( entity.getId(), diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryEstimateEntity.java b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryEstimateEntity.java new file mode 100644 index 00000000..c59b8f1a --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryEstimateEntity.java @@ -0,0 +1,82 @@ +package com.bablsoft.accessflow.core.internal.persistence.entity; + +import com.bablsoft.accessflow.core.api.QueryType; +import jakarta.persistence.Column; +import jakarta.persistence.Entity; +import jakarta.persistence.EnumType; +import jakarta.persistence.Enumerated; +import jakarta.persistence.FetchType; +import jakarta.persistence.Id; +import jakarta.persistence.JoinColumn; +import jakarta.persistence.ManyToOne; +import jakarta.persistence.Table; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; +import org.hibernate.annotations.JdbcType; +import org.hibernate.annotations.JdbcTypeCode; +import org.hibernate.dialect.type.PostgreSQLEnumJdbcType; +import org.hibernate.type.SqlTypes; + +import java.time.Instant; +import java.util.UUID; + +@Entity +@Table(name = "query_estimates") +@Getter +@Setter +@NoArgsConstructor +public class QueryEstimateEntity { + + @Id + private UUID id; + + @ManyToOne(fetch = FetchType.LAZY) + @JoinColumn(name = "query_request_id", nullable = false) + private QueryRequestEntity queryRequest; + + @Column(name = "engine_id", length = 64) + private String engineId; + + @Enumerated(EnumType.STRING) + @JdbcType(PostgreSQLEnumJdbcType.class) + @Column(name = "query_type", columnDefinition = "query_type") + private QueryType queryType; + + @Column(nullable = false) + private boolean supported = false; + + @Column(name = "estimated_rows") + private Long estimatedRows; + + @Column(name = "affected_row_count") + private Long affectedRowCount; + + @Column(name = "scan_type", length = 128) + private String scanType; + + @Column(name = "estimated_cost") + private Double estimatedCost; + + @JdbcTypeCode(SqlTypes.JSON) + @Column(columnDefinition = "jsonb") + private String plan; + + @Column(name = "raw_plan", columnDefinition = "text") + private String rawPlan; + + @Column(name = "unsupported_reason", length = 500) + private String unsupportedReason; + + @Column(nullable = false) + private boolean failed = false; + + @Column(name = "error_message", length = 500) + private String errorMessage; + + @Column(name = "duration_ms") + private Integer durationMs; + + @Column(name = "created_at", nullable = false, updatable = false) + private Instant createdAt = Instant.now(); +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryRequestEntity.java b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryRequestEntity.java index 435d0fe6..c0fc286e 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryRequestEntity.java +++ b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/entity/QueryRequestEntity.java @@ -66,6 +66,11 @@ public class QueryRequestEntity { @Column(name = "ai_analysis_id") private UUID aiAnalysisId; + // FK to query_estimates (AF-624) — bare UUID for the same circular-reference reason; + // the DB constraint is added in V128 after query_estimates is created. + @Column(name = "query_estimate_id") + private UUID queryEstimateId; + // Bare UUID to access_grant_request (#582) — the grant row's lifecycle (expiry, revocation) // is independent of this query's audit trail, so no FK. @Column(name = "approved_by_grant_id") diff --git a/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/repo/QueryEstimateRepository.java b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/repo/QueryEstimateRepository.java new file mode 100644 index 00000000..857fcc24 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/core/internal/persistence/repo/QueryEstimateRepository.java @@ -0,0 +1,12 @@ +package com.bablsoft.accessflow.core.internal.persistence.repo; + +import com.bablsoft.accessflow.core.internal.persistence.entity.QueryEstimateEntity; +import org.springframework.data.jpa.repository.JpaRepository; + +import java.util.Optional; +import java.util.UUID; + +public interface QueryEstimateRepository extends JpaRepository { + + Optional findByQueryRequestId(UUID queryRequestId); +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryCostEstimateService.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryCostEstimateService.java new file mode 100644 index 00000000..e93c831d --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryCostEstimateService.java @@ -0,0 +1,25 @@ +package com.bablsoft.accessflow.proxy.api; + +import com.bablsoft.accessflow.core.api.QueryEstimateSnapshot; + +import java.util.Optional; +import java.util.UUID; + +/** + * Computes and persists the pre-flight cost / blast-radius estimate for a submitted query (issue + * AF-624): the engine's non-committing dry-run plan (estimated rows, scan type, cost) plus, for + * UPDATE/DELETE, a governed affected-row count. Runs asynchronously right after submission and is + * consumed by reviewers ({@code GET /queries/{id}}), the routing-policy engine + * ({@code estimated_rows} / {@code scan_type} conditions), and the AI analyzer's prompt context. + */ +public interface QueryCostEstimateService { + + /** + * Idempotent: returns the existing estimate when one is already persisted for the query, + * otherwise computes, persists, and publishes {@code QueryEstimateCompletedEvent} (or + * {@code QueryEstimateFailedEvent} plus a sentinel row on unexpected errors). Never throws — + * failures, timeouts, and unsupported engines persist a degraded row instead of blocking the + * caller. Empty only when the query request itself does not exist. + */ + Optional estimateSubmittedQuery(UUID queryRequestId); +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryExecutor.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryExecutor.java index 8a9ec599..2022027b 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryExecutor.java +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/api/QueryExecutor.java @@ -1,5 +1,6 @@ package com.bablsoft.accessflow.proxy.api; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; @@ -10,6 +11,16 @@ public interface QueryExecutor { QueryExecutionResult execute(QueryExecutionRequest request); + /** + * Count the rows the request's UPDATE/DELETE would affect (issue AF-624) without + * mutating data, applying the request's row-security directives. For relational datasources + * the executor rewrites the statement into a governed {@code SELECT COUNT(*)}; for + * engine-managed datasources it delegates to {@code QueryEngine.countAffectedRows}. Statement + * shapes that cannot be provably counted (joins, {@code USING} lists) and engines with no + * count concept return a {@link QueryAffectedRowsResult} with {@code supported=false}. + */ + QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request); + /** * Produce a non-committing execution plan + best-effort estimated row impact for the request's * query (issue AF-445) without mutating data. The request's row-security directives are diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutor.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutor.java index ba80972e..91aa2d1b 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutor.java +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutor.java @@ -1,5 +1,6 @@ package com.bablsoft.accessflow.proxy.internal; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; @@ -54,6 +55,13 @@ public QueryDryRunResult dryRun(QueryExecutionRequest request) { return delegate.dryRun(request); } + // Like dryRun: a COUNT(*) materializes a single row, so the heap-protecting permit budget + // does not apply; per-datasource concurrency stays HikariCP's job. + @Override + public QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request) { + return delegate.countAffectedRows(request); + } + @Override public SelectExecutionResult sampleTable(SampleTableRequest request) { return withPermit(() -> delegate.sampleTable(request)); diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateService.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateService.java new file mode 100644 index 00000000..dca1f52a --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateService.java @@ -0,0 +1,243 @@ +package com.bablsoft.accessflow.proxy.internal; + +import com.bablsoft.accessflow.core.api.PersistQueryEstimateCommand; +import com.bablsoft.accessflow.core.api.QueryDryRunResult; +import com.bablsoft.accessflow.core.api.QueryEstimateLookupService; +import com.bablsoft.accessflow.core.api.QueryEstimatePersistenceService; +import com.bablsoft.accessflow.core.api.QueryEstimateSnapshot; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryPlanNode; +import com.bablsoft.accessflow.core.api.QueryRequestLookupService; +import com.bablsoft.accessflow.core.api.QueryRequestSnapshot; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.RowSecurityDirective; +import com.bablsoft.accessflow.core.api.RowSecurityResolutionService; +import com.bablsoft.accessflow.core.events.QueryEstimateCompletedEvent; +import com.bablsoft.accessflow.core.events.QueryEstimateFailedEvent; +import com.bablsoft.accessflow.proxy.api.QueryCostEstimateService; +import com.bablsoft.accessflow.proxy.api.QueryExecutor; +import lombok.RequiredArgsConstructor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.context.ApplicationEventPublisher; +import org.springframework.context.MessageSource; +import org.springframework.context.i18n.LocaleContextHolder; +import org.springframework.stereotype.Service; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.node.ObjectNode; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Optional; +import java.util.UUID; + +/** + * Reuses the AF-445 dry-run machinery ({@code QueryExecutor.dryRun} — dialect EXPLAIN planners and + * engine {@code dryRun} overrides) plus the AF-624 affected-row count, bounded by the tighter + * {@code accessflow.proxy.estimate-timeout}. Two independent triggers call this (the + * submission-event listener and the AI analyzer, which needs the estimate for its prompt); the + * persistence layer's insert-once semantics make the race harmless. + */ +@Service +@RequiredArgsConstructor +class DefaultQueryCostEstimateService implements QueryCostEstimateService { + + private static final Logger log = LoggerFactory.getLogger(DefaultQueryCostEstimateService.class); + + private final QueryEstimateLookupService queryEstimateLookupService; + private final QueryEstimatePersistenceService queryEstimatePersistenceService; + private final QueryRequestLookupService queryRequestLookupService; + private final RowSecurityResolutionService rowSecurityResolutionService; + private final QueryExecutor queryExecutor; + private final ProxyPoolProperties properties; + private final ObjectMapper objectMapper; + private final ApplicationEventPublisher eventPublisher; + private final MessageSource messageSource; + private final Clock clock; + + @Override + public Optional estimateSubmittedQuery(UUID queryRequestId) { + var existing = queryEstimateLookupService.findByQueryRequestId(queryRequestId); + if (existing.isPresent()) { + return existing; + } + var snapshot = queryRequestLookupService.findById(queryRequestId).orElse(null); + if (snapshot == null) { + log.warn("Cost estimate skipped: query request {} not found", queryRequestId); + return Optional.empty(); + } + if (snapshot.transactional()) { + return persistAndPublish(queryRequestId, unsupportedCommand(snapshot, null, + msg("error.estimate.transactional_unsupported")), true); + } + Instant start = clock.instant(); + try { + var request = buildRequest(snapshot); + var dryRun = queryExecutor.dryRun(request); + Long affectedRows = countAffectedRows(request, snapshot.queryType()); + return persistAndPublish(queryRequestId, + toCommand(snapshot, dryRun, affectedRows, durationMs(start)), true); + } catch (RuntimeException ex) { + log.warn("Cost estimate failed for query {}: {}", queryRequestId, ex.getMessage()); + var command = new PersistQueryEstimateCommand(null, snapshot.queryType(), false, null, + null, null, null, null, null, null, true, truncate(ex.getMessage()), + durationMs(start)); + var result = persistAndPublish(queryRequestId, command, false); + eventPublisher.publishEvent( + new QueryEstimateFailedEvent(queryRequestId, truncate(ex.getMessage()))); + return result; + } + } + + private QueryExecutionRequest buildRequest(QueryRequestSnapshot snapshot) { + var rowSecurityPredicates = rowSecurityResolutionService + .resolveApplicable(snapshot.organizationId(), snapshot.datasourceId(), + snapshot.submittedByUserId()).stream() + .map(p -> new RowSecurityDirective(p.policyId(), p.tableRef(), p.columnName(), + p.operator(), p.values())) + .toList(); + return new QueryExecutionRequest(snapshot.datasourceId(), snapshot.sqlText(), + snapshot.queryType(), null, properties.estimateTimeout(), List.of(), List.of(), + rowSecurityPredicates, false, null, List.of()); + } + + private Long countAffectedRows(QueryExecutionRequest request, QueryType queryType) { + if (queryType != QueryType.UPDATE && queryType != QueryType.DELETE) { + return null; + } + try { + var count = queryExecutor.countAffectedRows(request); + return count.supported() ? count.affectedRows() : null; + } catch (RuntimeException ex) { + log.debug("Affected-row count failed for datasource {}: {}", + request.datasourceId(), ex.getMessage()); + return null; + } + } + + private PersistQueryEstimateCommand toCommand(QueryRequestSnapshot snapshot, + QueryDryRunResult dryRun, Long affectedRows, + int durationMs) { + if (!dryRun.supported()) { + var reason = dryRun.unsupportedReason() != null + ? dryRun.unsupportedReason() + : msg("error.dry_run.unsupported", dryRun.engineId()); + var command = unsupportedCommand(snapshot, dryRun.engineId(), reason); + // A degraded plan may still carry an exact write count (e.g. INSERT has no plan but + // the COUNT rewrite worked) — keep it. + return affectedRows == null ? command + : withAffectedRows(command, affectedRows, durationMs); + } + var root = dryRun.plan(); + var access = accessNode(root, snapshot.queryType()); + Long estimatedRows = dryRun.estimatedRows(); + if ((estimatedRows == null || estimatedRows == 0) && access != null + && access != root && access.estimatedRows() != null) { + estimatedRows = Math.round(access.estimatedRows()); + } + return new PersistQueryEstimateCommand( + dryRun.engineId(), snapshot.queryType(), true, estimatedRows, + affectedRows, + access != null ? truncateTo(access.operation(), 128) : null, + root != null ? root.estimatedCost() : null, + planJson(root), + dryRun.rawPlan(), null, false, null, durationMs); + } + + /** + * Relational write plans wrap the access path in a modify node (e.g. PostgreSQL's + * {@code ModifyTable} with {@code Plan Rows: 0}) — the scan type and row estimate a reviewer + * or routing policy cares about live in its first child, so descend one level for writes. + */ + private static QueryPlanNode accessNode(QueryPlanNode root, QueryType queryType) { + if (root == null) { + return null; + } + if ((queryType == QueryType.UPDATE || queryType == QueryType.DELETE + || queryType == QueryType.INSERT) && !root.children().isEmpty()) { + return root.children().get(0); + } + return root; + } + + private PersistQueryEstimateCommand unsupportedCommand(QueryRequestSnapshot snapshot, + String engineId, String reason) { + return new PersistQueryEstimateCommand(engineId, snapshot.queryType(), false, null, null, + null, null, null, null, truncate(reason), false, null, null); + } + + private static PersistQueryEstimateCommand withAffectedRows(PersistQueryEstimateCommand command, + Long affectedRows, int durationMs) { + return new PersistQueryEstimateCommand(command.engineId(), command.queryType(), + command.supported(), command.estimatedRows(), affectedRows, command.scanType(), + command.estimatedCost(), command.planJson(), command.rawPlan(), + command.unsupportedReason(), command.failed(), command.errorMessage(), durationMs); + } + + /** Serializes the plan tree with explicit snake_case keys — the frontend's PlanTree shape. */ + private String planJson(QueryPlanNode root) { + if (root == null) { + return null; + } + return objectMapper.writeValueAsString(planNode(root)); + } + + private ObjectNode planNode(QueryPlanNode node) { + var out = objectMapper.createObjectNode(); + out.put("operation", node.operation()); + out.put("target", node.target()); + if (node.estimatedRows() != null) { + out.put("estimated_rows", node.estimatedRows()); + } else { + out.putNull("estimated_rows"); + } + if (node.estimatedCost() != null) { + out.put("estimated_cost", node.estimatedCost()); + } else { + out.putNull("estimated_cost"); + } + out.put("detail", node.detail()); + var children = out.putArray("children"); + for (var child : node.children()) { + children.add(planNode(child)); + } + return out; + } + + private Optional persistAndPublish(UUID queryRequestId, + PersistQueryEstimateCommand command, + boolean publishCompleted) { + var estimateId = queryEstimatePersistenceService.persist(queryRequestId, command); + if (publishCompleted) { + eventPublisher.publishEvent(new QueryEstimateCompletedEvent(queryRequestId, estimateId, + command.supported())); + } + return queryEstimateLookupService.findByQueryRequestId(queryRequestId); + } + + private int durationMs(Instant start) { + return (int) Math.min(Duration.between(start, clock.instant()).toMillis(), + Integer.MAX_VALUE); + } + + private static String truncate(String message) { + if (message == null) { + return null; + } + return message.length() <= 500 ? message : message.substring(0, 500); + } + + private static String truncateTo(String value, int max) { + if (value == null) { + return null; + } + return value.length() <= max ? value : value.substring(0, max); + } + + private String msg(String key, Object... args) { + return messageSource.getMessage(key, args.length == 0 ? null : args, + LocaleContextHolder.getLocale()); + } +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutor.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutor.java index 167f0254..76b07c89 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutor.java +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutor.java @@ -6,6 +6,7 @@ import com.bablsoft.accessflow.core.api.QueryType; import com.bablsoft.accessflow.core.api.ColumnMaskDirective; import com.bablsoft.accessflow.proxy.api.DatasourceUnavailableException; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngineCatalog; import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; @@ -14,6 +15,7 @@ import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; import com.bablsoft.accessflow.proxy.api.QueryExecutor; +import com.bablsoft.accessflow.proxy.internal.dryrun.AffectedRowCounter; import com.bablsoft.accessflow.proxy.internal.dryrun.DryRunPlanRequest; import com.bablsoft.accessflow.proxy.internal.dryrun.DryRunPlanner; import com.bablsoft.accessflow.proxy.internal.dryrun.DryRunPlannerRegistry; @@ -231,6 +233,41 @@ public QueryDryRunResult dryRun(QueryExecutionRequest request) { } } + @Override + public QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request) { + var descriptor = datasourceLookupService.findById(request.datasourceId()) + .orElseThrow(() -> new DatasourceUnavailableException( + msg("error.datasource_unavailable_not_found"))); + Duration effectiveTimeout = request.statementTimeoutOverride() != null + ? request.statementTimeoutOverride() + : properties.execution().statementTimeout(); + String engineId = descriptor.connectorId() != null + ? descriptor.connectorId() + : descriptor.dbType().name().toLowerCase(java.util.Locale.ROOT); + + if (engineCatalog.isEngineManaged(descriptor.dbType())) { + return engineCatalog.engineFor(descriptor.dbType()) + .countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + effectiveTimeout)); + } + + var countSql = AffectedRowCounter.toCountSql(request.sql()); + if (countSql.isEmpty()) { + return QueryAffectedRowsResult.unsupported(engineId); + } + Instant start = clock.instant(); + var result = execute(new QueryExecutionRequest(request.datasourceId(), countSql.get(), + QueryType.SELECT, null, effectiveTimeout, List.of(), List.of(), + request.rowSecurityPredicates(), false, null, List.of())); + if (result instanceof SelectExecutionResult select + && !select.rows().isEmpty() + && !select.rows().get(0).isEmpty() + && select.rows().get(0).get(0) instanceof Number count) { + return QueryAffectedRowsResult.of(engineId, count.longValue(), durationSince(start)); + } + return QueryAffectedRowsResult.unsupported(engineId); + } + private QueryExecutionResult executeTransactional(QueryExecutionRequest request, DbType dbType, Duration effectiveTimeout, Instant start) { var appliedPolicyIds = new java.util.LinkedHashSet(); diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolProperties.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolProperties.java index 58d5cfd2..aca05d14 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolProperties.java +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolProperties.java @@ -11,7 +11,8 @@ record ProxyPoolProperties( Duration maxLifetime, Duration leakDetectionThreshold, String poolNamePrefix, - Execution execution) { + Execution execution, + Duration estimateTimeout) { ProxyPoolProperties { if (connectionTimeout == null) { @@ -32,6 +33,9 @@ record ProxyPoolProperties( if (execution == null) { execution = new Execution(null, null, null, null, null, null, null); } + if (estimateTimeout == null) { + estimateTimeout = Duration.ofSeconds(5); + } } record Execution(Integer maxRows, Duration statementTimeout, Integer defaultFetchSize, diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListener.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListener.java new file mode 100644 index 00000000..072ad767 --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListener.java @@ -0,0 +1,25 @@ +package com.bablsoft.accessflow.proxy.internal; + +import com.bablsoft.accessflow.core.events.QuerySubmittedEvent; +import com.bablsoft.accessflow.proxy.api.QueryCostEstimateService; +import lombok.RequiredArgsConstructor; +import org.springframework.modulith.events.ApplicationModuleListener; +import org.springframework.stereotype.Component; + +/** + * Kicks off the pre-flight cost estimate (issue AF-624) asynchronously after submission — the + * proxy-side analogue of the AI module's {@code AiAnalysisListener}. Runs unconditionally + * (regardless of {@code ai_analysis_enabled}); the AI analyzer independently requests the same + * estimate for its prompt, and the service's insert-once persistence makes the race harmless. + */ +@Component +@RequiredArgsConstructor +class QueryCostEstimateListener { + + private final QueryCostEstimateService queryCostEstimateService; + + @ApplicationModuleListener + void onSubmitted(QuerySubmittedEvent event) { + queryCostEstimateService.estimateSubmittedQuery(event.queryRequestId()); + } +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounter.java b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounter.java new file mode 100644 index 00000000..6441624d --- /dev/null +++ b/backend/src/main/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounter.java @@ -0,0 +1,74 @@ +package com.bablsoft.accessflow.proxy.internal.dryrun; + +import net.sf.jsqlparser.JSQLParserException; +import net.sf.jsqlparser.expression.Function; +import net.sf.jsqlparser.expression.operators.relational.ExpressionList; +import net.sf.jsqlparser.parser.CCJSqlParserUtil; +import net.sf.jsqlparser.statement.delete.Delete; +import net.sf.jsqlparser.statement.select.AllColumns; +import net.sf.jsqlparser.statement.select.PlainSelect; +import net.sf.jsqlparser.statement.select.SelectItem; +import net.sf.jsqlparser.statement.update.Update; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; + +/** + * Rewrites a plain single-table UPDATE/DELETE into the equivalent {@code SELECT COUNT(*) FROM + * [WHERE …]} so the proxy can report the exact number of rows the write would affect + * (issue AF-624) without mutating data. Shapes whose count semantics are not provably identical — + * joins, {@code UPDATE … FROM}, {@code DELETE … USING} — return empty so the caller degrades to + * the EXPLAIN estimate. The rewritten SELECT still flows through the row-security rewriter, so the + * count reflects the governed statement. + */ +public final class AffectedRowCounter { + + private AffectedRowCounter() { + } + + public static Optional toCountSql(String sql) { + net.sf.jsqlparser.statement.Statement statement; + try { + statement = CCJSqlParserUtil.parse(sql); + } catch (JSQLParserException ex) { + return Optional.empty(); + } + return switch (statement) { + case Update update -> toCountSql(update); + case Delete delete -> toCountSql(delete); + default -> Optional.empty(); + }; + } + + private static Optional toCountSql(Update update) { + if ((update.getJoins() != null && !update.getJoins().isEmpty()) + || (update.getStartJoins() != null && !update.getStartJoins().isEmpty()) + || update.getFromItem() != null || update.getTable() == null) { + return Optional.empty(); + } + return Optional.of(countSelect(update.getTable(), + update.getWhere())); + } + + private static Optional toCountSql(Delete delete) { + if ((delete.getJoins() != null && !delete.getJoins().isEmpty()) + || (delete.getUsingList() != null && !delete.getUsingList().isEmpty()) + || delete.getTable() == null) { + return Optional.empty(); + } + return Optional.of(countSelect(delete.getTable(), delete.getWhere())); + } + + private static String countSelect(net.sf.jsqlparser.schema.Table table, + net.sf.jsqlparser.expression.Expression where) { + var count = new Function(); + count.setName("COUNT"); + count.setParameters(new ExpressionList<>(new AllColumns())); + var select = new PlainSelect(); + select.setSelectItems(new ArrayList<>(List.of(new SelectItem<>(count)))); + select.setFromItem(table); + select.setWhere(where); + return select.toString(); + } +} diff --git a/backend/src/main/java/com/bablsoft/accessflow/realtime/internal/RealtimeEventDispatcher.java b/backend/src/main/java/com/bablsoft/accessflow/realtime/internal/RealtimeEventDispatcher.java index 6d884289..78e6cc2b 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/realtime/internal/RealtimeEventDispatcher.java +++ b/backend/src/main/java/com/bablsoft/accessflow/realtime/internal/RealtimeEventDispatcher.java @@ -15,6 +15,8 @@ import com.bablsoft.accessflow.core.api.UserRoleType; import com.bablsoft.accessflow.core.api.UserView; import com.bablsoft.accessflow.core.events.AiAnalysisCompletedEvent; +import com.bablsoft.accessflow.core.events.QueryEstimateCompletedEvent; +import com.bablsoft.accessflow.core.events.QueryEstimateFailedEvent; import com.bablsoft.accessflow.core.events.AnomalyDetectedEvent; import com.bablsoft.accessflow.core.events.QueryReadyForReviewEvent; import com.bablsoft.accessflow.core.events.QueryStatusChangedEvent; @@ -180,6 +182,29 @@ void onAiAnalysisCompleted(AiAnalysisCompletedEvent event) { }); } + @ApplicationModuleListener + void onQueryEstimateCompleted(QueryEstimateCompletedEvent event) { + sendEstimateComplete(event.queryRequestId(), event.supported()); + } + + @ApplicationModuleListener + void onQueryEstimateFailed(QueryEstimateFailedEvent event) { + sendEstimateComplete(event.queryRequestId(), false); + } + + private void sendEstimateComplete(java.util.UUID queryRequestId, boolean supported) { + safe("query.estimate_complete", queryRequestId, () -> { + var snapshot = queryRequestLookupService.findById(queryRequestId).orElse(null); + if (snapshot == null) { + return; + } + ObjectNode data = objectMapper.createObjectNode(); + data.put("query_id", queryRequestId.toString()); + data.put("supported", supported); + sendTo(snapshot.submittedByUserId(), "query.estimate_complete", data); + }); + } + @ApplicationModuleListener void onQueryReadyForReview(QueryReadyForReviewEvent event) { safe("review.new_request", event.queryRequestId(), () -> { diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ComparisonOperator.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ComparisonOperator.java index 5b07b062..a6629483 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ComparisonOperator.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ComparisonOperator.java @@ -14,6 +14,13 @@ public enum ComparisonOperator { * @return {@code true} if {@code left right} holds. */ public boolean test(int left, int right) { + return test((long) left, (long) right); + } + + /** + * @return {@code true} if {@code left right} holds. + */ + public boolean test(long left, long right) { return switch (this) { case LT -> left < right; case LTE -> left <= right; diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionContext.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionContext.java index c72b7969..88ae5919 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionContext.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionContext.java @@ -28,6 +28,12 @@ * anomaly (UBA, AF-383) on this datasource at submission time. Because anomaly detection is a * periodic batch over past audit data, it cannot mutate an already-executed query — instead it * raises this signal so a routing policy can ESCALATE the requester's next query. + * + *

{@code estimatedRows} and {@code scanType} carry the pre-flight cost estimate (AF-624), read + * live at routing time from the persisted {@code query_estimates} row: {@code estimatedRows} is + * the exact affected-row count for UPDATE/DELETE when available, otherwise the EXPLAIN estimate; + * {@code scanType} is the plan's root operation (e.g. {@code Seq Scan}). Both are {@code null} + * when the estimate is absent, unsupported, or failed — the matching conditions fail closed. */ public record ConditionContext( QueryType queryType, @@ -44,15 +50,36 @@ public record ConditionContext( String requesterUserAgent, boolean ciCdOrigin, Integer minutesSinceLastApproval, - boolean anomalyActive) { + boolean anomalyActive, + Long estimatedRows, + String scanType) { public ConditionContext { referencedTables = Set.copyOf(referencedTables == null ? Set.of() : referencedTables); requesterGroupIds = Set.copyOf(requesterGroupIds == null ? Set.of() : requesterGroupIds); } + /** Backward-compatible constructor without the AF-624 estimate signals (defaults to absent). */ + public ConditionContext(QueryType queryType, Set referencedTables, RiskLevel riskLevel, + int riskScore, String requesterRoleName, Set requesterGroupIds, + LocalDateTime evaluatedAt, boolean hasWhereClause, + boolean hasLimitClause, boolean transactional, + String requesterIpAddress, String requesterUserAgent, + boolean ciCdOrigin, Integer minutesSinceLastApproval, + boolean anomalyActive) { + this(queryType, referencedTables, riskLevel, riskScore, requesterRoleName, + requesterGroupIds, evaluatedAt, hasWhereClause, hasLimitClause, transactional, + requesterIpAddress, requesterUserAgent, ciCdOrigin, minutesSinceLastApproval, + anomalyActive, null, null); + } + /** @return {@code true} when an AI risk level / score signal is present. */ public boolean hasRiskSignal() { return riskLevel != null && riskScore >= 0; } + + /** @return {@code true} when a pre-flight estimated-row signal is present (AF-624). */ + public boolean hasEstimateSignal() { + return estimatedRows != null; + } } diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionNode.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionNode.java index 67bdef4e..a95e783e 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionNode.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/api/ConditionNode.java @@ -185,6 +185,36 @@ record TimeSinceLastApproval(ComparisonOperator operator, int minutes) implement record CiCdOrigin(boolean expected) implements ConditionNode { } + /** + * Compares the pre-flight estimated row impact (AF-624 — for UPDATE/DELETE the exact + * affected-row count when available, otherwise the EXPLAIN estimate) with {@code value}. + * Fails closed: evaluates to {@code false} when no estimate signal exists — + * the estimate has not been computed yet, the engine has no plan concept, or the computation + * failed — so a permissive policy keyed on a small estimate never fires on missing context. + */ + record EstimatedRows(ComparisonOperator operator, long value) implements ConditionNode { + public EstimatedRows { + if (operator == null) { + throw new IllegalArgumentException("EstimatedRows condition requires an operator"); + } + if (value < 0) { + throw new IllegalArgumentException("EstimatedRows value must be >= 0"); + } + } + } + + /** + * Matches when the pre-flight plan's root scan/operation type (AF-624 — e.g. {@code Seq Scan}, + * {@code Index Scan}, {@code COLLSCAN}) matches any glob in {@code patterns} ({@code *} = any + * run of characters, case-insensitive — same matcher as referenced tables). + * Fails closed: evaluates to {@code false} when no plan was captured. + */ + record ScanTypeMatches(List patterns) implements ConditionNode { + public ScanTypeMatches { + patterns = List.copyOf(patterns == null ? List.of() : patterns); + } + } + /** * Matches when the requester has an active (OPEN) behavioural anomaly on this datasource * (UBA, AF-383) equal to {@code expected}. Pair {@code AnomalyDetected(true)} with an diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachine.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachine.java index ed12d018..a7e3c1f0 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachine.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachine.java @@ -3,6 +3,7 @@ import com.bablsoft.accessflow.access.api.AccessGrantLookupService; import com.bablsoft.accessflow.access.api.AccessGrantView; import com.bablsoft.accessflow.ai.api.BehaviorAnomalyLookupService; +import com.bablsoft.accessflow.core.api.QueryEstimateLookupService; import com.bablsoft.accessflow.core.api.QueryRequestLookupService; import com.bablsoft.accessflow.core.api.QueryRequestSnapshot; import com.bablsoft.accessflow.core.api.QueryRequestStateService; @@ -73,6 +74,7 @@ class QueryReviewStateMachine { private final RoutingDecisionService routingDecisionService; private final BehaviorAnomalyLookupService behaviorAnomalyLookupService; private final AccessGrantLookupService accessGrantLookupService; + private final QueryEstimateLookupService queryEstimateLookupService; private final MessageSource messageSource; private final ApplicationEventPublisher eventPublisher; @@ -294,10 +296,22 @@ private ConditionContext buildContext(QueryRequestSnapshot query, RiskLevel risk .orElse(null); boolean anomalyActive = behaviorAnomalyLookupService.hasActiveAnomaly( query.organizationId(), query.submittedByUserId(), query.datasourceId()); + // AF-624: live, fail-closed lookup of the pre-flight estimate — the estimate pipeline runs + // independently of AI analysis, so whatever is persisted right now is the signal; absent / + // unsupported / failed rows leave both fields null and the matching conditions false. + Long estimatedRows = null; + String scanType = null; + var estimate = queryEstimateLookupService.findByQueryRequestId(query.id()).orElse(null); + if (estimate != null && !estimate.failed()) { + estimatedRows = estimate.affectedRowCount() != null + ? estimate.affectedRowCount() + : estimate.estimatedRows(); + scanType = estimate.scanType(); + } return new ConditionContext(query.queryType(), referencedTables, riskLevel, riskScore, roleName, groupIds, LocalDateTime.now(clock), hasWhere, hasLimit, transactional, query.submittedIp(), query.submittedUserAgent(), query.ciCdOrigin(), - minutesSinceLastApproval, anomalyActive); + minutesSinceLastApproval, anomalyActive, estimatedRows, scanType); } private QueryStatus decideNextStatus(ReviewPlanSnapshot plan, QueryRequestSnapshot query, diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/ConditionNodeMixin.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/ConditionNodeMixin.java index 1bc11a91..577888b8 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/ConditionNodeMixin.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/ConditionNodeMixin.java @@ -31,7 +31,9 @@ @JsonSubTypes.Type(value = ConditionNode.UserAgentMatches.class, name = "user_agent"), @JsonSubTypes.Type(value = ConditionNode.TimeSinceLastApproval.class, name = "time_since_last_approval"), @JsonSubTypes.Type(value = ConditionNode.CiCdOrigin.class, name = "cicd_origin"), - @JsonSubTypes.Type(value = ConditionNode.AnomalyDetected.class, name = "anomaly_detected") + @JsonSubTypes.Type(value = ConditionNode.AnomalyDetected.class, name = "anomaly_detected"), + @JsonSubTypes.Type(value = ConditionNode.EstimatedRows.class, name = "estimated_rows"), + @JsonSubTypes.Type(value = ConditionNode.ScanTypeMatches.class, name = "scan_type") }) interface ConditionNodeMixin { } diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluator.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluator.java index 0fa8b66a..87f9f5dd 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluator.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluator.java @@ -48,9 +48,20 @@ public boolean matches(ConditionNode node, ConditionContext ctx) { && c.operator().test(ctx.minutesSinceLastApproval(), c.minutes()); case ConditionNode.CiCdOrigin c -> ctx.ciCdOrigin() == c.expected(); case ConditionNode.AnomalyDetected c -> ctx.anomalyActive() == c.expected(); + case ConditionNode.EstimatedRows c -> + ctx.hasEstimateSignal() && c.operator().test(ctx.estimatedRows(), c.value()); + case ConditionNode.ScanTypeMatches c -> matchesScanType(c, ctx); }; } + private static boolean matchesScanType(ConditionNode.ScanTypeMatches c, ConditionContext ctx) { + String scanType = ctx.scanType(); + if (scanType == null) { + return false; + } + return c.patterns().stream().anyMatch(pattern -> GlobMatcher.matches(pattern, scanType)); + } + private static boolean matchesUserAgent(ConditionNode.UserAgentMatches c, ConditionContext ctx) { String userAgent = ctx.requesterUserAgent(); if (userAgent == null) { diff --git a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/web/QueryDetailResponse.java b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/web/QueryDetailResponse.java index 393ebc12..2c710f9f 100644 --- a/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/web/QueryDetailResponse.java +++ b/backend/src/main/java/com/bablsoft/accessflow/workflow/internal/web/QueryDetailResponse.java @@ -28,6 +28,7 @@ public record QueryDetailResponse( QueryStatus status, String justification, AiAnalysisDetail aiAnalysis, + CostEstimateDetail costEstimate, Long rowsAffected, Integer durationMs, String errorMessage, @@ -68,6 +69,7 @@ public static QueryDetailResponse from(QueryDetailView view, MatchedRoutingPolic view.status(), view.justification(), AiAnalysisDetail.from(view.aiAnalysis()), + CostEstimateDetail.from(view.costEstimate()), view.rowsAffected(), view.durationMs(), view.errorMessage(), @@ -188,6 +190,48 @@ static AiAnalysisDetail from(QueryDetailView.AiAnalysisDetail src) { } } + /** + * The persisted pre-flight cost / blast-radius estimate (AF-624). {@code plan} embeds the + * stored snake_case plan-tree JSON raw (same shape as the dry-run endpoint's {@code plan}). + */ + public record CostEstimateDetail( + UUID id, + String engineId, + QueryType queryType, + boolean supported, + Long estimatedRows, + Long affectedRowCount, + String scanType, + Double estimatedCost, + @JsonRawValue String plan, + String rawPlan, + String unsupportedReason, + boolean failed, + String errorMessage, + Integer durationMs) { + + static CostEstimateDetail from(QueryDetailView.CostEstimateDetail src) { + if (src == null) { + return null; + } + return new CostEstimateDetail( + src.id(), + src.engineId(), + src.queryType(), + src.supported(), + src.estimatedRows(), + src.affectedRowCount(), + src.scanType(), + src.estimatedCost(), + src.planJson(), + src.rawPlan(), + src.unsupportedReason(), + src.failed(), + src.errorMessage(), + src.durationMs()); + } + } + public record ReviewDecisionDetail( UUID id, ReviewerRef reviewer, diff --git a/backend/src/main/resources/application.yml b/backend/src/main/resources/application.yml index 2a367f5b..f540d7cc 100644 --- a/backend/src/main/resources/application.yml +++ b/backend/src/main/resources/application.yml @@ -219,6 +219,9 @@ accessflow: max-lifetime: ${ACCESSFLOW_PROXY_MAX_LIFETIME:30m} leak-detection-threshold: ${ACCESSFLOW_PROXY_LEAK_DETECTION_THRESHOLD:0s} pool-name-prefix: accessflow-ds- + # Tight bound on the async pre-flight cost estimate (AF-624): EXPLAIN + affected-row COUNT + # each run under this statement timeout instead of the full execution statement-timeout. + estimate-timeout: ${ACCESSFLOW_PROXY_ESTIMATE_TIMEOUT:PT5S} execution: max-rows: ${ACCESSFLOW_PROXY_EXECUTION_MAX_ROWS:10000} statement-timeout: ${ACCESSFLOW_PROXY_EXECUTION_STATEMENT_TIMEOUT:30s} diff --git a/backend/src/main/resources/db/migration/V127__create_query_estimates.sql b/backend/src/main/resources/db/migration/V127__create_query_estimates.sql new file mode 100644 index 00000000..1523a637 --- /dev/null +++ b/backend/src/main/resources/db/migration/V127__create_query_estimates.sql @@ -0,0 +1,18 @@ +CREATE TABLE query_estimates ( + id UUID PRIMARY KEY, + query_request_id UUID NOT NULL UNIQUE REFERENCES query_requests(id) ON DELETE CASCADE, + engine_id VARCHAR(64), + query_type query_type, + supported BOOLEAN NOT NULL DEFAULT false, + estimated_rows BIGINT, + affected_row_count BIGINT, + scan_type VARCHAR(128), + estimated_cost DOUBLE PRECISION, + plan JSONB, + raw_plan TEXT, + unsupported_reason VARCHAR(500), + failed BOOLEAN NOT NULL DEFAULT false, + error_message VARCHAR(500), + duration_ms INTEGER, + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP +); diff --git a/backend/src/main/resources/db/migration/V128__add_query_estimate_id_to_query_requests.sql b/backend/src/main/resources/db/migration/V128__add_query_estimate_id_to_query_requests.sql new file mode 100644 index 00000000..eaef5e6a --- /dev/null +++ b/backend/src/main/resources/db/migration/V128__add_query_estimate_id_to_query_requests.sql @@ -0,0 +1,6 @@ +ALTER TABLE query_requests + ADD COLUMN query_estimate_id UUID; + +ALTER TABLE query_requests + ADD CONSTRAINT fk_query_requests_query_estimate + FOREIGN KEY (query_estimate_id) REFERENCES query_estimates(id); diff --git a/backend/src/main/resources/i18n/messages.properties b/backend/src/main/resources/i18n/messages.properties index 576560b2..71c70d9d 100644 --- a/backend/src/main/resources/i18n/messages.properties +++ b/backend/src/main/resources/i18n/messages.properties @@ -86,6 +86,7 @@ error.unauthorized=Authentication required error.forbidden=Access denied error.permission.table_not_allowed=Query references one or more tables that are not in the user''s allow-list: {0} error.dry_run.unsupported=Dry-run is not supported for the {0} engine +error.estimate.transactional_unsupported=Cost estimation is not available for transactional batch envelopes # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=Setup has already been completed diff --git a/backend/src/main/resources/i18n/messages_de.properties b/backend/src/main/resources/i18n/messages_de.properties index 078a5ed9..53debd7a 100644 --- a/backend/src/main/resources/i18n/messages_de.properties +++ b/backend/src/main/resources/i18n/messages_de.properties @@ -73,6 +73,7 @@ error.unauthorized=Authentifizierung erforderlich error.forbidden=Zugriff verweigert error.permission.table_not_allowed=Die Abfrage verweist auf eine oder mehrere Tabellen, die nicht in der Zulassungsliste des Benutzers stehen: {0} error.dry_run.unsupported=Probelauf wird für die {0}-Engine nicht unterstützt +error.estimate.transactional_unsupported=Die Kostenschätzung ist für transaktionale Batch-Umschläge nicht verfügbar # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=Die Ersteinrichtung wurde bereits abgeschlossen diff --git a/backend/src/main/resources/i18n/messages_es.properties b/backend/src/main/resources/i18n/messages_es.properties index 2583ed21..4bc4fb86 100644 --- a/backend/src/main/resources/i18n/messages_es.properties +++ b/backend/src/main/resources/i18n/messages_es.properties @@ -73,6 +73,7 @@ error.unauthorized=Se requiere autenticación error.forbidden=Acceso denegado error.permission.table_not_allowed=La consulta hace referencia a una o más tablas que no están en la lista de permitidos del usuario: {0} error.dry_run.unsupported=La simulación no es compatible con el motor {0} +error.estimate.transactional_unsupported=La estimación de coste no está disponible para envoltorios de lotes transaccionales # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=La configuración inicial ya se ha completado diff --git a/backend/src/main/resources/i18n/messages_fr.properties b/backend/src/main/resources/i18n/messages_fr.properties index 7636f138..29f86512 100644 --- a/backend/src/main/resources/i18n/messages_fr.properties +++ b/backend/src/main/resources/i18n/messages_fr.properties @@ -73,6 +73,7 @@ error.unauthorized=Authentification requise error.forbidden=Accès refusé error.permission.table_not_allowed=La requête référence une ou plusieurs tables absentes de la liste d''autorisation de l''utilisateur : {0} error.dry_run.unsupported=La simulation n''est pas prise en charge pour le moteur {0} +error.estimate.transactional_unsupported=L''estimation du coût n''est pas disponible pour les enveloppes de lots transactionnels # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=La configuration initiale est déjà terminée diff --git a/backend/src/main/resources/i18n/messages_hy.properties b/backend/src/main/resources/i18n/messages_hy.properties index 263f437e..18be8df4 100644 --- a/backend/src/main/resources/i18n/messages_hy.properties +++ b/backend/src/main/resources/i18n/messages_hy.properties @@ -73,6 +73,7 @@ error.unauthorized=Անհրաժեշտ է վավերացում error.forbidden=Մուտքն արգելված է error.permission.table_not_allowed=Հարցումը հղում է կատարում մեկ կամ մի քանի աղյուսակների, որոնք օգտատիրոջ թույլատրված ցանկում չեն՝ {0} error.dry_run.unsupported={0} շարժիչի համար փորձնական գործարկումը չի աջակցվում +error.estimate.transactional_unsupported=Գնահատումը հասանելի չէ գործարքային փաթեթային հարցումների համար # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=Սկզբնական կարգավորումն արդեն ավարտված է diff --git a/backend/src/main/resources/i18n/messages_ru.properties b/backend/src/main/resources/i18n/messages_ru.properties index 8e4ad457..b94c207d 100644 --- a/backend/src/main/resources/i18n/messages_ru.properties +++ b/backend/src/main/resources/i18n/messages_ru.properties @@ -73,6 +73,7 @@ error.unauthorized=Требуется аутентификация error.forbidden=Доступ запрещён error.permission.table_not_allowed=Запрос обращается к одной или нескольким таблицам, отсутствующим в списке разрешённых пользователя: {0} error.dry_run.unsupported=Пробный запуск не поддерживается для движка {0} +error.estimate.transactional_unsupported=Оценка стоимости недоступна для транзакционных пакетных конвертов # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=Первоначальная настройка уже завершена diff --git a/backend/src/main/resources/i18n/messages_zh_CN.properties b/backend/src/main/resources/i18n/messages_zh_CN.properties index 5518298c..aee423b4 100644 --- a/backend/src/main/resources/i18n/messages_zh_CN.properties +++ b/backend/src/main/resources/i18n/messages_zh_CN.properties @@ -73,6 +73,7 @@ error.unauthorized=需要身份验证 error.forbidden=访问被拒绝 error.permission.table_not_allowed=查询引用了一个或多个不在用户允许列表中的表：{0} error.dry_run.unsupported={0} 引擎不支持试运行 +error.estimate.transactional_unsupported=事务批处理语句不支持成本预估 # ── Auth / setup ────────────────────────────────────────────────────────────── error.setup_already_completed=系统初始化已完成 diff --git a/backend/src/test/java/com/bablsoft/accessflow/TestSystemRoleSeeder.java b/backend/src/test/java/com/bablsoft/accessflow/TestSystemRoleSeeder.java new file mode 100644 index 00000000..dacd37ae --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/TestSystemRoleSeeder.java @@ -0,0 +1,61 @@ +package com.bablsoft.accessflow; + +import com.bablsoft.accessflow.core.api.SystemRolePermissions; +import com.bablsoft.accessflow.core.api.UserRoleType; +import org.springframework.jdbc.core.JdbcTemplate; + +import java.util.Map; + +/** + * Restores the V114 global system-role seed rows after a test wipes them. Integration tests that + * clean the shared Testcontainers database with {@code TRUNCATE TABLE organizations CASCADE} + * truncate every transitively-referencing table — including {@code roles} (its + * {@code organization_id} FK makes the whole table a cascade target, system rows included). Any + * later context that resolves a role name (row-security reveal roles, API-masking reveal roles, + * requester-role routing conditions) then fails on the missing seed. Call + * {@link #reseedSystemRoles(JdbcTemplate)} immediately after such a truncate so suite ordering + * never matters. + * + *

Ids and descriptions mirror {@code V114__create_roles_and_permissions.sql}; the permission + * sets come from {@link SystemRolePermissions}, which the seed-parity test guarantees is identical + * to the migration. Inserts are idempotent ({@code ON CONFLICT DO NOTHING}). + */ +public final class TestSystemRoleSeeder { + + private static final Map SEED_IDS = Map.of( + UserRoleType.ADMIN, "c0000000-0000-0000-0000-000000000001", + UserRoleType.REVIEWER, "c0000000-0000-0000-0000-000000000002", + UserRoleType.ANALYST, "c0000000-0000-0000-0000-000000000003", + UserRoleType.READONLY, "c0000000-0000-0000-0000-000000000004", + UserRoleType.AUDITOR, "c0000000-0000-0000-0000-000000000005"); + + private static final Map DESCRIPTIONS = Map.of( + UserRoleType.ADMIN, "Full administrative access to every AccessFlow capability.", + UserRoleType.REVIEWER, "Submits SELECT/DML queries and reviews queries, access requests, " + + "API requests, erasure requests, and attestation items.", + UserRoleType.ANALYST, "Submits SELECT and DML queries and views own history.", + UserRoleType.READONLY, "Submits SELECT queries only.", + UserRoleType.AUDITOR, "Read-only compliance role: compliance reports, attestation " + + "evidence, break-glass events, and anomalies."); + + private TestSystemRoleSeeder() { + } + + public static void reseedSystemRoles(JdbcTemplate jdbcTemplate) { + for (var role : UserRoleType.values()) { + var id = SEED_IDS.get(role); + jdbcTemplate.update(""" + INSERT INTO roles (id, organization_id, name, description, is_system) + VALUES (?::uuid, NULL, ?, ?, TRUE) + ON CONFLICT (id) DO NOTHING""", + id, role.name(), DESCRIPTIONS.get(role)); + for (var permission : SystemRolePermissions.of(role)) { + jdbcTemplate.update(""" + INSERT INTO role_permissions (role_id, permission) + VALUES (?::uuid, ?) + ON CONFLICT DO NOTHING""", + id, permission.name()); + } + } + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalysisListenerIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalysisListenerIntegrationTest.java index 3591faba..6018b77a 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalysisListenerIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalysisListenerIntegrationTest.java @@ -160,7 +160,7 @@ void onSubmittedPersistsAnalysisAndLinksQueryRequest() { List.of(new AiIssue(RiskLevel.HIGH, "DELETE_WITHOUT_WHERE", "msg", "fix")), false, 1000L, AiProviderType.ANTHROPIC, "claude-sonnet-4-20250514", 100, 50, List.of()); - when(strategy.analyze(any(), any(), any(), any(), any())).thenReturn(result); + when(strategy.analyze(any(), any(), any(), any(), any(), any())).thenReturn(result); new TransactionTemplate(transactionManager).executeWithoutResult(status -> eventPublisher.publishEvent(new QuerySubmittedEvent(queryRequestId))); @@ -209,7 +209,7 @@ void onSubmittedSkipsAnalysisAndAdvancesWhenAiDisabled() { @Test void onSubmittedPersistsSentinelOnStrategyFailure() { - when(strategy.analyze(any(), any(), any(), any(), any())) + when(strategy.analyze(any(), any(), any(), any(), any(), any())) .thenThrow(new com.bablsoft.accessflow.ai.api.AiAnalysisException("provider down")); new TransactionTemplate(transactionManager).executeWithoutResult(status -> diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolderTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolderTest.java index 016b6c81..c7743f5a 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolderTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/AiAnalyzerStrategyHolderTest.java @@ -586,7 +586,7 @@ private static class ThrowingStrategy implements com.bablsoft.accessflow.ai.api. @Override public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, - String language, UUID aiConfigId) { + String costEstimateContext, String language, UUID aiConfigId) { calls++; throw failure; } @@ -605,7 +605,7 @@ private static class SucceedingStrategy implements com.bablsoft.accessflow.ai.ap @Override public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, - String language, UUID aiConfigId) { + String costEstimateContext, String language, UUID aiConfigId) { calls++; return new AiAnalysisResult(10, com.bablsoft.accessflow.core.api.RiskLevel.LOW, "ok", List.of(), false, null, AiProviderType.OLLAMA, "test-model", 10, 5, List.of()); @@ -685,8 +685,8 @@ private static class StubStrategy implements com.bablsoft.accessflow.ai.api.AiAn int calls = 0; @Override - public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, String language, - UUID aiConfigId) { + public AiAnalysisResult analyze(String sql, DbType dbType, String schemaContext, + String costEstimateContext, String language, UUID aiConfigId) { calls++; return null; } diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/BehaviorAnomalyDetectionIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/BehaviorAnomalyDetectionIntegrationTest.java index b9293f83..5202bc92 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/BehaviorAnomalyDetectionIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/BehaviorAnomalyDetectionIntegrationTest.java @@ -1,6 +1,7 @@ package com.bablsoft.accessflow.ai.internal; import com.bablsoft.accessflow.TestcontainersConfig; +import com.bablsoft.accessflow.TestSystemRoleSeeder; import com.bablsoft.accessflow.ai.api.BehaviorAnomalyDetectionService; import com.bablsoft.accessflow.ai.api.BehaviorAnomalyStatus; import com.bablsoft.accessflow.ai.internal.persistence.repo.BehaviorAnomalyRepository; @@ -134,6 +135,7 @@ void setUp() { @AfterEach void cleanup() { jdbcTemplate.execute("TRUNCATE TABLE organizations CASCADE"); + TestSystemRoleSeeder.reseedSystemRoles(jdbcTemplate); jdbcTemplate.execute("TRUNCATE TABLE audit_log CASCADE"); } diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerServiceTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerServiceTest.java index 8c4f0c5e..1612225b 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerServiceTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/DefaultAiAnalyzerServiceTest.java @@ -75,6 +75,7 @@ class DefaultAiAnalyzerServiceTest { @Mock SqlParserService sqlParserService; @Mock ApplicationEventPublisher eventPublisher; @Mock AiRateLimiter aiRateLimiter; + @Mock com.bablsoft.accessflow.proxy.api.QueryCostEstimateService queryCostEstimateService; private final SystemPromptRenderer promptRenderer = new SystemPromptRenderer(); private final AiResponseParser responseParser = new AiResponseParser(JsonMapper.builder().build()); @@ -97,7 +98,7 @@ void setUp() { datasourceLookupService, datasourceAdminService, queryRequestLookupService, permissionLookupService, aiAnalysisPersistenceService, localizationConfigService, dataClassificationQueryService, sqlParserService, eventPublisher, aiRateLimiter, - observationRegistry, meterRegistry); + observationRegistry, queryCostEstimateService, meterRegistry); org.mockito.Mockito.lenient().when(localizationConfigService.getOrDefault(any())) .thenReturn(new LocalizationConfigView(organizationId, List.of("en"), "en", "en")); org.mockito.Mockito.lenient().when(aiConfigRepository.findById(aiConfigId)) @@ -216,7 +217,7 @@ void analyzeSubmittedQueryPersistsResultAndPublishesCompleted() { when(queryRequestLookupService.findById(queryRequestId)).thenReturn(Optional.of(snapshot)); when(datasourceLookupService.findById(datasourceId)).thenReturn(Optional.of(descriptor(DbType.MYSQL))); when(datasourceAdminService.introspectSchemaForSystem(datasourceId, organizationId)).thenReturn(schemaView()); - when(strategy.analyze(eq("SELECT 1"), eq(DbType.MYSQL), any(), any(), eq(aiConfigId))) + when(strategy.analyze(eq("SELECT 1"), eq(DbType.MYSQL), any(), any(), any(), eq(aiConfigId))) .thenReturn(sampleResult()); var newAnalysisId = UUID.randomUUID(); when(aiAnalysisPersistenceService.persist(eq(queryRequestId), any())).thenReturn(newAnalysisId); @@ -265,7 +266,7 @@ void analyzeSubmittedQueryPersistsPerModelBreakdown() { 100, 40, 900L, false, null), new AiModelResult(AiProviderType.OLLAMA, "llama3", 50, RiskLevel.MEDIUM, 2.0, 50, 20, 200L, false, null))); - when(strategy.analyze(eq("SELECT 1"), eq(DbType.MYSQL), any(), any(), eq(aiConfigId))) + when(strategy.analyze(eq("SELECT 1"), eq(DbType.MYSQL), any(), any(), any(), eq(aiConfigId))) .thenReturn(multiModel); when(aiAnalysisPersistenceService.persist(eq(queryRequestId), any())).thenReturn(UUID.randomUUID()); @@ -295,7 +296,7 @@ void analyzeSubmittedQueryRaisesRiskWhenReferencedTableIsClassified() { .thenReturn(new SqlParseResult(QueryType.SELECT, false, List.of("SELECT * FROM users"), Set.of("users"))); // LLM verdict is MEDIUM/60; PCI adds +30 → 90 → CRITICAL. - when(strategy.analyze(any(), any(), any(), any(), eq(aiConfigId))) + when(strategy.analyze(any(), any(), any(), any(), any(), eq(aiConfigId))) .thenReturn(new AiAnalysisResult(60, RiskLevel.MEDIUM, "s", List.of(), false, null, AiProviderType.ANTHROPIC, "model-x", 1, 1, List.of())); when(aiAnalysisPersistenceService.persist(eq(queryRequestId), any())).thenReturn(UUID.randomUUID()); @@ -328,7 +329,7 @@ void analyzeSubmittedQueryAnnotatesSchemaContextWithClassifications() { .thenReturn(new SqlParseResult(QueryType.SELECT, false, List.of("SELECT id FROM users"), Set.of("users"))); ArgumentCaptor contextCaptor = ArgumentCaptor.forClass(String.class); - when(strategy.analyze(eq("SELECT id FROM users"), eq(DbType.POSTGRESQL), contextCaptor.capture(), + when(strategy.analyze(eq("SELECT id FROM users"), eq(DbType.POSTGRESQL), contextCaptor.capture(), any(), any(), eq(aiConfigId))).thenReturn(sampleResult()); when(aiAnalysisPersistenceService.persist(eq(queryRequestId), any())).thenReturn(UUID.randomUUID()); @@ -378,7 +379,7 @@ void analyzeSubmittedQueryPersistsSentinelOnStrategyFailure() { when(queryRequestLookupService.findById(queryRequestId)).thenReturn(Optional.of(snapshot)); when(datasourceLookupService.findById(datasourceId)).thenReturn(Optional.of(descriptor(DbType.POSTGRESQL))); when(datasourceAdminService.introspectSchemaForSystem(datasourceId, organizationId)).thenReturn(schemaView()); - when(strategy.analyze(any(), any(), any(), any(), eq(aiConfigId))) + when(strategy.analyze(any(), any(), any(), any(), any(), eq(aiConfigId))) .thenThrow(new AiAnalysisException("provider down")); service.analyzeSubmittedQuery(queryRequestId); @@ -401,7 +402,7 @@ void analyzeSubmittedQueryPersistsSentinelOnParseFailure() { when(queryRequestLookupService.findById(queryRequestId)).thenReturn(Optional.of(snapshot)); when(datasourceLookupService.findById(datasourceId)).thenReturn(Optional.of(descriptor(DbType.POSTGRESQL))); when(datasourceAdminService.introspectSchemaForSystem(datasourceId, organizationId)).thenReturn(schemaView()); - when(strategy.analyze(any(), any(), any(), any(), eq(aiConfigId))) + when(strategy.analyze(any(), any(), any(), any(), any(), eq(aiConfigId))) .thenThrow(new AiAnalysisParseException("bad json")); service.analyzeSubmittedQuery(queryRequestId); @@ -417,13 +418,13 @@ void analyzeSubmittedQueryFallsBackToNullSchemaWhenIntrospectionFails() { when(datasourceLookupService.findById(datasourceId)).thenReturn(Optional.of(descriptor(DbType.POSTGRESQL))); when(datasourceAdminService.introspectSchemaForSystem(datasourceId, organizationId)) .thenThrow(new RuntimeException("connection refused")); - when(strategy.analyze(eq("SELECT 1"), eq(DbType.POSTGRESQL), eq(null), any(), eq(aiConfigId))) + when(strategy.analyze(eq("SELECT 1"), eq(DbType.POSTGRESQL), eq(null), any(), any(), eq(aiConfigId))) .thenReturn(sampleResult()); when(aiAnalysisPersistenceService.persist(eq(queryRequestId), any())).thenReturn(UUID.randomUUID()); service.analyzeSubmittedQuery(queryRequestId); - verify(strategy).analyze(eq("SELECT 1"), eq(DbType.POSTGRESQL), eq(null), any(), eq(aiConfigId)); + verify(strategy).analyze(eq("SELECT 1"), eq(DbType.POSTGRESQL), eq(null), any(), any(), eq(aiConfigId)); verify(eventPublisher).publishEvent(any(AiAnalysisCompletedEvent.class)); } @@ -459,7 +460,7 @@ void sentinelFallsBackToUnknownWhenAiConfigLookupFails() { when(queryRequestLookupService.findById(queryRequestId)).thenReturn(Optional.of(snapshot)); when(datasourceLookupService.findById(datasourceId)).thenReturn(Optional.of(descriptor(DbType.POSTGRESQL))); when(datasourceAdminService.introspectSchemaForSystem(datasourceId, organizationId)).thenReturn(schemaView()); - when(strategy.analyze(any(), any(), any(), any(), eq(aiConfigId))) + when(strategy.analyze(any(), any(), any(), any(), any(), eq(aiConfigId))) .thenThrow(new AiAnalysisException("provider down")); when(aiConfigRepository.findById(aiConfigId)).thenThrow(new RuntimeException("db unreachable")); @@ -483,6 +484,7 @@ void analyzePreviewPropagatesRateLimitExceeded() { assertThatThrownBy(() -> service.analyzePreview(datasourceId, "SELECT 1", userId, organizationId, false)) .isInstanceOf(AiRateLimitExceededException.class); verify(strategy, never()).analyze(any(), any(), any(), any(), any()); + verify(strategy, never()).analyze(any(), any(), any(), any(), any(), any()); } @Test @@ -505,6 +507,7 @@ void analyzeSubmittedQueryPersistsSentinelWhenBudgetExhausted() { assertThat(cmd.failed()).isTrue(); verify(eventPublisher).publishEvent(any(AiAnalysisFailedEvent.class)); verify(strategy, never()).analyze(any(), any(), any(), any(), any()); + verify(strategy, never()).analyze(any(), any(), any(), any(), any(), any()); } @Test @@ -525,5 +528,6 @@ void analyzeSubmittedQueryPersistsSentinelWhenRateLimited() { assertThat(captor.getValue().riskLevel()).isEqualTo(RiskLevel.CRITICAL); verify(eventPublisher).publishEvent(any(AiAnalysisFailedEvent.class)); verify(strategy, never()).analyze(any(), any(), any(), any(), any()); + verify(strategy, never()).analyze(any(), any(), any(), any(), any(), any()); } } diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategyTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategyTest.java index c0f698b3..9f684dba 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategyTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/GuardrailAiAnalyzerStrategyTest.java @@ -72,23 +72,23 @@ void matchingIsCaseInsensitiveAndSubstring() { @Test void allowsAnalyzeWhenNoPatternMatches() { - when(delegate.analyze(any(), any(), any(), any(), any())).thenReturn(result()); + when(delegate.analyze(any(), any(), any(), any(), any(), any())).thenReturn(result()); var returned = guardrail("drop\\s+table") .analyze("SELECT 1", DbType.POSTGRESQL, null, "en", CONFIG_ID); assertThat(returned).isNotNull(); - verify(delegate).analyze(any(), any(), any(), any(), any()); + verify(delegate).analyze(any(), any(), any(), any(), any(), any()); } @Test void emptyPatternsPassThrough() { - when(delegate.analyze(any(), any(), any(), any(), any())).thenReturn(result()); + when(delegate.analyze(any(), any(), any(), any(), any(), any())).thenReturn(result()); var returned = guardrail().analyze("DROP TABLE t", DbType.POSTGRESQL, null, "en", CONFIG_ID); assertThat(returned).isNotNull(); - verify(delegate).analyze(any(), any(), any(), any(), any()); + verify(delegate).analyze(any(), any(), any(), any(), any(), any()); } @Test diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategyTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategyTest.java index 8e2c114e..c5f7456a 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategyTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/OrchestratingAiAnalyzerStrategyTest.java @@ -48,9 +48,9 @@ private OrchestratingAiAnalyzerStrategy orchestrator(VotingStrategy strategy, @Test void aggregatesAllSuccessfulMembersAndSumsTokens() { - when(primaryStrategy.analyze(any(), any(), any(), any(), any())) + when(primaryStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenReturn(result(AiProviderType.ANTHROPIC, "claude", 40, RiskLevel.MEDIUM, 100, 20)); - when(memberStrategy.analyze(any(), any(), any(), any(), any())) + when(memberStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenReturn(result(AiProviderType.OLLAMA, "llama", 80, RiskLevel.CRITICAL, 50, 10)); var orchestrator = orchestrator(VotingStrategy.MAX_RISK, @@ -75,9 +75,9 @@ void aggregatesAllSuccessfulMembersAndSumsTokens() { @Test void partialFailureAggregatesSurvivorsAndRecordsFailedMember() { - when(primaryStrategy.analyze(any(), any(), any(), any(), any())) + when(primaryStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenReturn(result(AiProviderType.ANTHROPIC, "claude", 55, RiskLevel.HIGH, 100, 20)); - when(memberStrategy.analyze(any(), any(), any(), any(), any())) + when(memberStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenThrow(new AiAnalysisException("provider down")); var orchestrator = orchestrator(VotingStrategy.WEIGHTED_AVERAGE, @@ -97,9 +97,9 @@ void partialFailureAggregatesSurvivorsAndRecordsFailedMember() { @Test void rethrowsWhenEveryMemberFails() { - when(primaryStrategy.analyze(any(), any(), any(), any(), any())) + when(primaryStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenThrow(new AiAnalysisException("p1 down")); - when(memberStrategy.analyze(any(), any(), any(), any(), any())) + when(memberStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenThrow(new AiAnalysisException("p2 down")); var orchestrator = orchestrator(VotingStrategy.WEIGHTED_AVERAGE, @@ -112,7 +112,7 @@ void rethrowsWhenEveryMemberFails() { @Test void singleMemberDegeneratesToThatMemberWithOneBreakdownRow() { - when(primaryStrategy.analyze(any(), any(), any(), any(), any())) + when(primaryStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenReturn(result(AiProviderType.OPENAI, "gpt", 30, RiskLevel.MEDIUM, 10, 5)); var orchestrator = orchestrator(VotingStrategy.WEIGHTED_AVERAGE, @@ -149,7 +149,7 @@ void capturesPerMemberLatencyFromClock() { when(steppingClock.instant()).thenReturn( Instant.parse("2026-01-01T00:00:00Z"), Instant.parse("2026-01-01T00:00:00.250Z")); - when(primaryStrategy.analyze(any(), any(), any(), any(), any())) + when(primaryStrategy.analyze(any(), any(), any(), any(), any(), any())) .thenReturn(result(AiProviderType.OPENAI, "gpt", 30, RiskLevel.MEDIUM, 10, 5)); var orchestrator = new OrchestratingAiAnalyzerStrategy( diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/SystemPromptRendererTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/SystemPromptRendererTest.java index 835cf9a0..406252a0 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/SystemPromptRendererTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/SystemPromptRendererTest.java @@ -200,6 +200,32 @@ void blankCustomTemplateFallsBackToDefault() { assertThat(fromBlank).contains("Database type: POSTGRESQL"); } + @Test + void costEstimateContextIsSubstitutedIntoDefaultTemplate() { + var prompt = renderer.render(null, "DELETE FROM users", DbType.POSTGRESQL, null, null, + "Exact affected-row count for this DELETE: 90 rows.", "en"); + + assertThat(prompt).contains("Exact affected-row count for this DELETE: 90 rows."); + assertThat(prompt).doesNotContain("{{cost_estimate}}"); + } + + @Test + void costEstimateFallsBackWhenAbsent() { + var fromNull = renderer.render(null, "SELECT 1", DbType.POSTGRESQL, null, null, null, "en"); + var fromBlank = renderer.render(null, "SELECT 1", DbType.POSTGRESQL, null, null, " ", "en"); + + assertThat(fromNull).contains("(no cost estimate available)"); + assertThat(fromBlank).contains("(no cost estimate available)"); + } + + @Test + void customTemplateSubstitutesCostEstimatePlaceholder() { + var prompt = renderer.render("estimate={{cost_estimate}} sql={{sql}}", "SELECT 1", + DbType.POSTGRESQL, null, null, "~120 rows via Seq Scan", "en"); + + assertThat(prompt).isEqualTo("estimate=~120 rows via Seq Scan sql=SELECT 1"); + } + @Test void sqlIsSubstitutedLastSoTokenStringInSqlIsNotReSubstituted() { // The SQL value itself contains a "{{language}}" token. Because {{sql}} is replaced last, diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategyTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategyTest.java index 1c6918f0..ff7f3323 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategyTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/TracingAiAnalyzerStrategyTest.java @@ -46,7 +46,7 @@ private static AiAnalysisResult result() { @Test void returnsResultAndFiresSuccessTrace() { var result = result(); - when(delegate.analyze("SELECT 1", DbType.POSTGRESQL, "schema", "en", AI_CONFIG_ID)).thenReturn(result); + when(delegate.analyze("SELECT 1", DbType.POSTGRESQL, "schema", null, "en", AI_CONFIG_ID)).thenReturn(result); var returned = strategy().analyze("SELECT 1", DbType.POSTGRESQL, "schema", "en", AI_CONFIG_ID); @@ -60,7 +60,7 @@ void returnsResultAndFiresSuccessTrace() { @Test void firesErrorTraceAndRethrowsOnFailure() { - when(delegate.analyze(any(), any(), any(), any(), any())) + when(delegate.analyze(any(), any(), any(), any(), any(), any())) .thenThrow(new AiAnalysisException("boom")); assertThatThrownBy(() -> strategy().analyze("SELECT 1", DbType.POSTGRESQL, null, "en", AI_CONFIG_ID)) @@ -76,7 +76,7 @@ void firesErrorTraceAndRethrowsOnFailure() { @Test void tracerFailureDoesNotBreakAnalysis() { var result = result(); - when(delegate.analyze(any(), any(), any(), any(), any())).thenReturn(result); + when(delegate.analyze(any(), any(), any(), any(), any(), any())).thenReturn(result); org.mockito.Mockito.doThrow(new RuntimeException("tracer down")).when(tracer).trace(any()); var returned = strategy().analyze("SELECT 1", DbType.POSTGRESQL, null, "en", AI_CONFIG_ID); @@ -86,9 +86,9 @@ void tracerFailureDoesNotBreakAnalysis() { @Test void delegateInvokedExactlyOnce() { - when(delegate.analyze(any(), any(), any(), any(), any())).thenReturn(result()); + when(delegate.analyze(any(), any(), any(), any(), any(), any())).thenReturn(result()); strategy().analyze("SELECT 1", DbType.POSTGRESQL, null, "en", AI_CONFIG_ID); - verify(delegate).analyze(any(), any(), any(), any(), any()); + verify(delegate).analyze(any(), any(), any(), any(), any(), any()); } @Test diff --git a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/web/AdminBehaviorAnomalyControllerIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/web/AdminBehaviorAnomalyControllerIntegrationTest.java index 69e4665f..c21ccb80 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/ai/internal/web/AdminBehaviorAnomalyControllerIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/ai/internal/web/AdminBehaviorAnomalyControllerIntegrationTest.java @@ -1,6 +1,7 @@ package com.bablsoft.accessflow.ai.internal.web; import com.bablsoft.accessflow.TestcontainersConfig; +import com.bablsoft.accessflow.TestSystemRoleSeeder; import com.bablsoft.accessflow.ai.api.BehaviorAnomalyStatus; import com.bablsoft.accessflow.ai.internal.persistence.entity.BehaviorAnomalyEntity; import com.bablsoft.accessflow.ai.internal.persistence.repo.BehaviorAnomalyRepository; @@ -100,6 +101,7 @@ void setUp() { @AfterEach void cleanup() { jdbcTemplate.execute("TRUNCATE TABLE organizations CASCADE"); + TestSystemRoleSeeder.reseedSystemRoles(jdbcTemplate); } // ----- list ----- diff --git a/backend/src/test/java/com/bablsoft/accessflow/audit/internal/web/AdminAuditLogControllerIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/audit/internal/web/AdminAuditLogControllerIntegrationTest.java index abeb7608..ef1244bd 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/audit/internal/web/AdminAuditLogControllerIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/audit/internal/web/AdminAuditLogControllerIntegrationTest.java @@ -1,6 +1,7 @@ package com.bablsoft.accessflow.audit.internal.web; import com.bablsoft.accessflow.TestcontainersConfig; +import com.bablsoft.accessflow.TestSystemRoleSeeder; import com.bablsoft.accessflow.audit.api.AuditAction; import com.bablsoft.accessflow.audit.api.AuditEntry; import com.bablsoft.accessflow.audit.api.AuditLogService; @@ -90,6 +91,7 @@ void cleanup() { // A plain DELETE FROM organizations chain would FK-fail whenever a prior test class left // datasources or other dependents behind — that exact flake bit this class in CI. jdbcTemplate.execute("TRUNCATE TABLE organizations CASCADE"); + TestSystemRoleSeeder.reseedSystemRoles(jdbcTemplate); } @Test diff --git a/backend/src/test/java/com/bablsoft/accessflow/compliance/internal/web/ComplianceReportControllerIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/compliance/internal/web/ComplianceReportControllerIntegrationTest.java index b57f94b0..4eeb8cf9 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/compliance/internal/web/ComplianceReportControllerIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/compliance/internal/web/ComplianceReportControllerIntegrationTest.java @@ -1,6 +1,7 @@ package com.bablsoft.accessflow.compliance.internal.web; import com.bablsoft.accessflow.TestcontainersConfig; +import com.bablsoft.accessflow.TestSystemRoleSeeder; import com.bablsoft.accessflow.core.api.AuthProviderType; import com.bablsoft.accessflow.core.api.CreateDataClassificationTagCommand; import com.bablsoft.accessflow.core.api.CredentialEncryptionService; @@ -116,6 +117,7 @@ void setUp() { @AfterEach void cleanup() { jdbcTemplate.execute("TRUNCATE TABLE organizations CASCADE"); + TestSystemRoleSeeder.reseedSystemRoles(jdbcTemplate); } @Test diff --git a/backend/src/test/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResultTest.java b/backend/src/test/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResultTest.java new file mode 100644 index 00000000..3295b951 --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/core/api/QueryAffectedRowsResultTest.java @@ -0,0 +1,39 @@ +package com.bablsoft.accessflow.core.api; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; + +import static org.assertj.core.api.Assertions.assertThat; + +class QueryAffectedRowsResultTest { + + @Test + void unsupportedCarriesEngineIdAndNoCount() { + var result = QueryAffectedRowsResult.unsupported("redis"); + + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("redis"); + assertThat(result.affectedRows()).isNull(); + assertThat(result.duration()).isEqualTo(Duration.ZERO); + assertThat(result.unsupportedReason()).isNull(); + } + + @Test + void unsupportedWithReason() { + var result = QueryAffectedRowsResult.unsupported("couchbase", "join shapes"); + + assertThat(result.supported()).isFalse(); + assertThat(result.unsupportedReason()).isEqualTo("join shapes"); + } + + @Test + void ofCarriesCountAndDuration() { + var result = QueryAffectedRowsResult.of("postgresql", 90L, Duration.ofMillis(8)); + + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("postgresql"); + assertThat(result.affectedRows()).isEqualTo(90L); + assertThat(result.duration()).isEqualTo(Duration.ofMillis(8)); + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateServiceTest.java b/backend/src/test/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateServiceTest.java new file mode 100644 index 00000000..7ef7e671 --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/core/internal/DefaultQueryEstimateServiceTest.java @@ -0,0 +1,149 @@ +package com.bablsoft.accessflow.core.internal; + +import com.bablsoft.accessflow.core.api.PersistQueryEstimateCommand; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.internal.persistence.entity.QueryEstimateEntity; +import com.bablsoft.accessflow.core.internal.persistence.entity.QueryRequestEntity; +import com.bablsoft.accessflow.core.internal.persistence.repo.QueryEstimateRepository; +import com.bablsoft.accessflow.core.internal.persistence.repo.QueryRequestRepository; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.time.Instant; +import java.util.Optional; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class DefaultQueryEstimateServiceTest { + + @Mock QueryEstimateRepository queryEstimateRepository; + @Mock QueryRequestRepository queryRequestRepository; + @InjectMocks DefaultQueryEstimateService service; + + private static PersistQueryEstimateCommand command() { + return new PersistQueryEstimateCommand("postgresql", QueryType.DELETE, true, 120L, 90L, + "Seq Scan", 44.5, "{\"operation\":\"Seq Scan\"}", "[raw]", null, false, null, 12); + } + + @Test + void persistInsertsEstimateAndLinksQueryRequest() { + var queryRequestId = UUID.randomUUID(); + var queryRequest = new QueryRequestEntity(); + queryRequest.setId(queryRequestId); + when(queryEstimateRepository.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()); + when(queryRequestRepository.findById(queryRequestId)).thenReturn(Optional.of(queryRequest)); + when(queryEstimateRepository.save(any(QueryEstimateEntity.class))) + .thenAnswer(inv -> inv.getArgument(0)); + + var id = service.persist(queryRequestId, command()); + + var captor = ArgumentCaptor.forClass(QueryEstimateEntity.class); + verify(queryEstimateRepository).save(captor.capture()); + var saved = captor.getValue(); + assertThat(saved.getId()).isEqualTo(id); + assertThat(saved.getQueryRequest()).isSameAs(queryRequest); + assertThat(saved.getEngineId()).isEqualTo("postgresql"); + assertThat(saved.getQueryType()).isEqualTo(QueryType.DELETE); + assertThat(saved.isSupported()).isTrue(); + assertThat(saved.getEstimatedRows()).isEqualTo(120L); + assertThat(saved.getAffectedRowCount()).isEqualTo(90L); + assertThat(saved.getScanType()).isEqualTo("Seq Scan"); + assertThat(saved.getEstimatedCost()).isEqualTo(44.5); + assertThat(saved.getPlan()).isEqualTo("{\"operation\":\"Seq Scan\"}"); + assertThat(saved.getRawPlan()).isEqualTo("[raw]"); + assertThat(saved.isFailed()).isFalse(); + assertThat(saved.getDurationMs()).isEqualTo(12); + assertThat(queryRequest.getQueryEstimateId()).isEqualTo(id); + } + + @Test + void persistReturnsExistingRowIdWhenRaced() { + var queryRequestId = UUID.randomUUID(); + var existing = new QueryEstimateEntity(); + existing.setId(UUID.randomUUID()); + when(queryEstimateRepository.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.of(existing)); + + var id = service.persist(queryRequestId, command()); + + assertThat(id).isEqualTo(existing.getId()); + verify(queryEstimateRepository, never()).save(any()); + } + + @Test + void persistThrowsWhenQueryRequestMissing() { + var queryRequestId = UUID.randomUUID(); + when(queryEstimateRepository.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()); + when(queryRequestRepository.findById(queryRequestId)).thenReturn(Optional.empty()); + + assertThatThrownBy(() -> service.persist(queryRequestId, command())) + .isInstanceOf(IllegalStateException.class); + } + + @Test + void lookupMapsEntityToSnapshot() { + var queryRequestId = UUID.randomUUID(); + var queryRequest = new QueryRequestEntity(); + queryRequest.setId(queryRequestId); + var entity = new QueryEstimateEntity(); + entity.setId(UUID.randomUUID()); + entity.setQueryRequest(queryRequest); + entity.setEngineId("postgresql"); + entity.setQueryType(QueryType.UPDATE); + entity.setSupported(true); + entity.setEstimatedRows(10L); + entity.setAffectedRowCount(8L); + entity.setScanType("Index Scan"); + entity.setEstimatedCost(1.5); + entity.setPlan("{}"); + entity.setRawPlan("raw"); + entity.setUnsupportedReason(null); + entity.setFailed(false); + entity.setErrorMessage(null); + entity.setDurationMs(4); + entity.setCreatedAt(Instant.parse("2026-07-22T10:00:00Z")); + when(queryEstimateRepository.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.of(entity)); + + var snapshot = service.findByQueryRequestId(queryRequestId); + + assertThat(snapshot).hasValueSatisfying(s -> { + assertThat(s.id()).isEqualTo(entity.getId()); + assertThat(s.queryRequestId()).isEqualTo(queryRequestId); + assertThat(s.engineId()).isEqualTo("postgresql"); + assertThat(s.queryType()).isEqualTo(QueryType.UPDATE); + assertThat(s.supported()).isTrue(); + assertThat(s.estimatedRows()).isEqualTo(10L); + assertThat(s.affectedRowCount()).isEqualTo(8L); + assertThat(s.scanType()).isEqualTo("Index Scan"); + assertThat(s.estimatedCost()).isEqualTo(1.5); + assertThat(s.planJson()).isEqualTo("{}"); + assertThat(s.rawPlan()).isEqualTo("raw"); + assertThat(s.failed()).isFalse(); + assertThat(s.durationMs()).isEqualTo(4); + assertThat(s.createdAt()).isEqualTo(Instant.parse("2026-07-22T10:00:00Z")); + }); + } + + @Test + void lookupEmptyWhenNoRow() { + var queryRequestId = UUID.randomUUID(); + when(queryEstimateRepository.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()); + + assertThat(service.findByQueryRequestId(queryRequestId)).isEmpty(); + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/dashboard/internal/web/DashboardControllerIntegrationTest.java b/backend/src/test/java/com/bablsoft/accessflow/dashboard/internal/web/DashboardControllerIntegrationTest.java index da3bdb35..4a6ffa0a 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/dashboard/internal/web/DashboardControllerIntegrationTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/dashboard/internal/web/DashboardControllerIntegrationTest.java @@ -1,6 +1,7 @@ package com.bablsoft.accessflow.dashboard.internal.web; import com.bablsoft.accessflow.TestcontainersConfig; +import com.bablsoft.accessflow.TestSystemRoleSeeder; import com.bablsoft.accessflow.ai.api.BehaviorAnomalyStatus; import com.bablsoft.accessflow.ai.internal.persistence.entity.BehaviorAnomalyEntity; import com.bablsoft.accessflow.ai.internal.persistence.repo.BehaviorAnomalyRepository; @@ -107,6 +108,7 @@ void cleanup() { // organizations cascade — truncate it explicitly (CASCADE clears the ai_analyses back-refs). jdbcTemplate.execute("TRUNCATE TABLE api_requests CASCADE"); jdbcTemplate.execute("TRUNCATE TABLE organizations CASCADE"); + TestSystemRoleSeeder.reseedSystemRoles(jdbcTemplate); } @Test diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutorTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutorTest.java index 09e2a266..f12b6571 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutorTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ConcurrencyLimitingQueryExecutorTest.java @@ -37,7 +37,7 @@ private ConcurrencyLimitingQueryExecutor executor(int maxConcurrent, Duration ac .thenReturn("concurrency limit reached"); var properties = new ProxyPoolProperties(null, null, null, null, null, new ProxyPoolProperties.Execution(null, null, null, null, null, maxConcurrent, - acquireTimeout)); + acquireTimeout), null); return new ConcurrencyLimitingQueryExecutor(delegate, properties, messageSource); } diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DatasourcePoolFactoryTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DatasourcePoolFactoryTest.java index cf9a1e89..bbf65cdf 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DatasourcePoolFactoryTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DatasourcePoolFactoryTest.java @@ -60,7 +60,7 @@ void setUp() { customJdbcDriverService = mock(com.bablsoft.accessflow.core.api.CustomJdbcDriverService.class); properties = new ProxyPoolProperties( Duration.ofSeconds(30), Duration.ofMinutes(10), Duration.ofMinutes(30), - Duration.ZERO, "accessflow-ds-", null); + Duration.ZERO, "accessflow-ds-", null, null); when(coordinatesFactory.from(any(), anyString(), anyInt(), anyString(), anyString(), any())) .thenReturn(new JdbcCoordinates( "jdbc:postgresql://h:5432/appdb?sslmode=disable", @@ -128,7 +128,7 @@ void createPoolSkipsLeakDetectionWhenZero() { void createPoolSetsLeakDetectionWhenPositive() { properties = new ProxyPoolProperties( Duration.ofSeconds(30), Duration.ofMinutes(10), Duration.ofMinutes(30), - Duration.ofSeconds(2), "accessflow-ds-", null); + Duration.ofSeconds(2), "accessflow-ds-", null, null); factory = new DatasourcePoolFactory(secretResolutionService, coordinatesFactory, properties, driverCatalog, customJdbcDriverService); diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateServiceTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateServiceTest.java new file mode 100644 index 00000000..09c4274f --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryCostEstimateServiceTest.java @@ -0,0 +1,299 @@ +package com.bablsoft.accessflow.proxy.internal; + +import com.bablsoft.accessflow.core.api.PersistQueryEstimateCommand; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; +import com.bablsoft.accessflow.core.api.QueryDryRunResult; +import com.bablsoft.accessflow.core.api.QueryEstimateLookupService; +import com.bablsoft.accessflow.core.api.QueryEstimatePersistenceService; +import com.bablsoft.accessflow.core.api.QueryEstimateSnapshot; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryPlanNode; +import com.bablsoft.accessflow.core.api.QueryRequestLookupService; +import com.bablsoft.accessflow.core.api.QueryRequestSnapshot; +import com.bablsoft.accessflow.core.api.QueryStatus; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.RowSecurityResolutionService; +import com.bablsoft.accessflow.core.events.QueryEstimateCompletedEvent; +import com.bablsoft.accessflow.core.events.QueryEstimateFailedEvent; +import com.bablsoft.accessflow.proxy.api.QueryExecutor; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.context.ApplicationEventPublisher; +import org.springframework.context.support.StaticMessageSource; +import tools.jackson.databind.json.JsonMapper; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class DefaultQueryCostEstimateServiceTest { + + @Mock QueryEstimateLookupService lookupService; + @Mock QueryEstimatePersistenceService persistenceService; + @Mock QueryRequestLookupService queryRequestLookupService; + @Mock RowSecurityResolutionService rowSecurityResolutionService; + @Mock QueryExecutor queryExecutor; + @Mock ApplicationEventPublisher eventPublisher; + + private DefaultQueryCostEstimateService service; + + private final UUID queryRequestId = UUID.randomUUID(); + private final UUID datasourceId = UUID.randomUUID(); + private final UUID organizationId = UUID.randomUUID(); + private final UUID userId = UUID.randomUUID(); + private final UUID estimateId = UUID.randomUUID(); + + @BeforeEach + void setUp() { + var properties = new ProxyPoolProperties(null, null, null, null, null, null, + Duration.ofSeconds(5)); + var messageSource = new StaticMessageSource(); + messageSource.setUseCodeAsDefaultMessage(true); + service = new DefaultQueryCostEstimateService(lookupService, persistenceService, + queryRequestLookupService, rowSecurityResolutionService, queryExecutor, properties, + JsonMapper.builder().build(), eventPublisher, messageSource, + Clock.fixed(Instant.parse("2026-07-22T10:00:00Z"), ZoneOffset.UTC)); + } + + private QueryRequestSnapshot snapshot(QueryType type, boolean transactional) { + return new QueryRequestSnapshot(queryRequestId, datasourceId, organizationId, userId, + "DELETE FROM users WHERE active = false", type, transactional, + QueryStatus.PENDING_AI, null, null, null, false); + } + + private QueryEstimateSnapshot persistedSnapshot() { + return new QueryEstimateSnapshot(estimateId, queryRequestId, "postgresql", + QueryType.DELETE, true, 100L, 90L, "Seq Scan", 12.5, null, "raw", null, + false, null, 3, Instant.now()); + } + + @Test + void returnsExistingEstimateWithoutRecomputing() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.of(persistedSnapshot())); + + var result = service.estimateSubmittedQuery(queryRequestId); + + assertThat(result).hasValueSatisfying(s -> assertThat(s.id()).isEqualTo(estimateId)); + verifyNoInteractions(queryExecutor, persistenceService, eventPublisher); + } + + @Test + void emptyWhenQueryRequestMissing() { + when(lookupService.findByQueryRequestId(queryRequestId)).thenReturn(Optional.empty()); + when(queryRequestLookupService.findById(queryRequestId)).thenReturn(Optional.empty()); + + assertThat(service.estimateSubmittedQuery(queryRequestId)).isEmpty(); + verifyNoInteractions(queryExecutor, persistenceService); + } + + @Test + void transactionalEnvelopePersistsUnsupportedRow() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.DELETE, true))); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + var result = service.estimateSubmittedQuery(queryRequestId); + + assertThat(result).isPresent(); + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().supported()).isFalse(); + assertThat(captor.getValue().unsupportedReason()) + .isEqualTo("error.estimate.transactional_unsupported"); + verify(eventPublisher).publishEvent(any(QueryEstimateCompletedEvent.class)); + verifyNoInteractions(queryExecutor); + } + + @Test + void supportedDryRunWithAffectedCountPersistsFullRow() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.DELETE, false))); + when(rowSecurityResolutionService.resolveApplicable(organizationId, datasourceId, userId)) + .thenReturn(List.of()); + var plan = new QueryPlanNode("Seq Scan", "users", 120.0, 44.5, null); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.of( + "postgresql", QueryType.DELETE, 120L, plan, "[raw]", Set.of(), + Duration.ofMillis(12))); + when(queryExecutor.countAffectedRows(any())) + .thenReturn(QueryAffectedRowsResult.of("postgresql", 90L, Duration.ofMillis(8))); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + var command = captor.getValue(); + assertThat(command.supported()).isTrue(); + assertThat(command.estimatedRows()).isEqualTo(120L); + assertThat(command.affectedRowCount()).isEqualTo(90L); + assertThat(command.scanType()).isEqualTo("Seq Scan"); + assertThat(command.estimatedCost()).isEqualTo(44.5); + assertThat(command.planJson()).contains("\"operation\":\"Seq Scan\"") + .contains("\"estimated_rows\":120.0"); + verify(eventPublisher).publishEvent(any(QueryEstimateCompletedEvent.class)); + var requestCaptor = ArgumentCaptor.forClass(QueryExecutionRequest.class); + verify(queryExecutor).dryRun(requestCaptor.capture()); + assertThat(requestCaptor.getValue().statementTimeoutOverride()) + .isEqualTo(Duration.ofSeconds(5)); + } + + @Test + void writePlanDescendsIntoAccessNodeForScanTypeAndEstimate() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.UPDATE, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + // PostgreSQL-style write plan: root ModifyTable with Plan Rows 0, child Seq Scan. + var scan = new QueryPlanNode("Seq Scan", "users", 2_400_000.0, 44_543.5, null); + var root = new QueryPlanNode("ModifyTable", "users", 0.0, 44_543.5, null, List.of(scan)); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.of( + "postgresql", QueryType.UPDATE, 0L, root, "[raw]", Set.of(), Duration.ZERO)); + when(queryExecutor.countAffectedRows(any())) + .thenReturn(QueryAffectedRowsResult.unsupported("postgresql")); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().scanType()).isEqualTo("Seq Scan"); + assertThat(captor.getValue().estimatedRows()).isEqualTo(2_400_000L); + } + + @Test + void selectSkipsAffectedRowCount() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.SELECT, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.of( + "postgresql", QueryType.SELECT, 10L, null, null, Set.of(), Duration.ZERO)); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + verify(queryExecutor, never()).countAffectedRows(any()); + } + + @Test + void unsupportedDryRunPersistsLocalizedReason() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.SELECT, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.unsupported("redis")); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().supported()).isFalse(); + assertThat(captor.getValue().unsupportedReason()).isEqualTo("error.dry_run.unsupported"); + verify(eventPublisher).publishEvent(any(QueryEstimateCompletedEvent.class)); + } + + @Test + void unsupportedDryRunKeepsAffectedCountWhenCountWorked() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.UPDATE, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.unsupported("custom")); + when(queryExecutor.countAffectedRows(any())) + .thenReturn(QueryAffectedRowsResult.of("custom", 7L, Duration.ofMillis(4))); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().supported()).isFalse(); + assertThat(captor.getValue().affectedRowCount()).isEqualTo(7L); + } + + @Test + void dryRunFailurePersistsFailedSentinelAndPublishesFailedEvent() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.DELETE, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + when(queryExecutor.dryRun(any())).thenThrow(new IllegalStateException("boom")); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + var result = service.estimateSubmittedQuery(queryRequestId); + + assertThat(result).isPresent(); + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().failed()).isTrue(); + assertThat(captor.getValue().errorMessage()).isEqualTo("boom"); + verify(eventPublisher).publishEvent(any(QueryEstimateFailedEvent.class)); + verify(eventPublisher, never()).publishEvent(any(QueryEstimateCompletedEvent.class)); + } + + @Test + void countFailureLeavesAffectedRowsNull() { + when(lookupService.findByQueryRequestId(queryRequestId)) + .thenReturn(Optional.empty()) + .thenReturn(Optional.of(persistedSnapshot())); + when(queryRequestLookupService.findById(queryRequestId)) + .thenReturn(Optional.of(snapshot(QueryType.DELETE, false))); + when(rowSecurityResolutionService.resolveApplicable(any(), any(), any())) + .thenReturn(List.of()); + when(queryExecutor.dryRun(any())).thenReturn(QueryDryRunResult.of( + "postgresql", QueryType.DELETE, 120L, null, null, Set.of(), Duration.ZERO)); + lenient().when(queryExecutor.countAffectedRows(any())) + .thenThrow(new IllegalStateException("count failed")); + when(persistenceService.persist(eq(queryRequestId), any())).thenReturn(estimateId); + + service.estimateSubmittedQuery(queryRequestId); + + var captor = ArgumentCaptor.forClass(PersistQueryEstimateCommand.class); + verify(persistenceService).persist(eq(queryRequestId), captor.capture()); + assertThat(captor.getValue().supported()).isTrue(); + assertThat(captor.getValue().affectedRowCount()).isNull(); + verify(eventPublisher).publishEvent(any(QueryEstimateCompletedEvent.class)); + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutorTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutorTest.java index d71211aa..da8fc88f 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutorTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/DefaultQueryExecutorTest.java @@ -76,7 +76,7 @@ class DefaultQueryExecutorTest { private final ProxyPoolProperties properties = new ProxyPoolProperties( null, null, null, null, null, new ProxyPoolProperties.Execution(10_000, Duration.ofSeconds(30), 1_000, null, - null, null, null)); + null, null, null), null); private final AtomicLong nanos = new AtomicLong(); private final Clock clock = Clock.fixed(Instant.parse("2026-05-05T12:00:00Z"), ZoneOffset.UTC); @@ -358,7 +358,7 @@ void transactionalBatchFlushesEveryChunkSizeRows() throws SQLException { var chunkedProperties = new ProxyPoolProperties( null, null, null, null, null, new ProxyPoolProperties.Execution(10_000, Duration.ofSeconds(30), 1_000, 2, - null, null, null)); + null, null, null), null); var healthRegistry = new ReplicaHealthRegistry(clock, new ProxyReplicaProperties(null, null, null)); var router = new RoutingDataSourceResolver(poolManager, healthRegistry, lookupService, @@ -678,6 +678,68 @@ private DatasourceConnectionDescriptor descriptor(int maxRows) { false, null, false, null, null, null, null, null, null, true); } + @Test + void countAffectedRowsRelationalRewritesDeleteIntoGovernedCount() throws SQLException { + var rs = mock(ResultSet.class); + var metadata = mock(ResultSetMetaData.class); + when(metadata.getColumnCount()).thenReturn(1); + when(metadata.getColumnLabel(1)).thenReturn("count"); + when(metadata.getColumnType(1)).thenReturn(Types.BIGINT); + when(metadata.getColumnTypeName(1)).thenReturn("int8"); + when(rs.getMetaData()).thenReturn(metadata); + when(rs.getObject(1)).thenReturn(90L); + when(rs.getLong(1)).thenReturn(90L); + when(rs.next()).thenReturn(true, false); + when(statement.executeQuery()).thenReturn(rs); + when(engineCatalog.isEngineManaged(DbType.POSTGRESQL)).thenReturn(false); + + var request = new QueryExecutionRequest(datasourceId, + "DELETE FROM users WHERE active = false", QueryType.DELETE, null, null); + + var result = executor.countAffectedRows(request); + + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(90L); + verify(connection).prepareStatement( + "SELECT COUNT(*) FROM users WHERE active = false"); + verify(statement, never()).executeLargeUpdate(); + } + + @Test + void countAffectedRowsUnsupportedForJoinShapes() { + when(engineCatalog.isEngineManaged(DbType.POSTGRESQL)).thenReturn(false); + + var result = executor.countAffectedRows(new QueryExecutionRequest(datasourceId, + "DELETE FROM orders USING users WHERE orders.user_id = users.id", + QueryType.DELETE, null, null)); + + assertThat(result.supported()).isFalse(); + } + + @Test + void countAffectedRowsEngineManagedDispatchesToEngine() { + var mongoDescriptor = new DatasourceConnectionDescriptor(datasourceId, UUID.randomUUID(), + DbType.MONGODB, "h", 27017, "db", "u", "ENC", SslMode.DISABLE, 10, 2_000, + false, null, false, null, null, null, null, null, null, true); + when(lookupService.findById(datasourceId)).thenReturn(Optional.of(mongoDescriptor)); + var engine = mock(com.bablsoft.accessflow.core.api.QueryEngine.class); + when(engineCatalog.isEngineManaged(DbType.MONGODB)).thenReturn(true); + when(engineCatalog.engineFor(DbType.MONGODB)).thenReturn(engine); + var expected = com.bablsoft.accessflow.core.api.QueryAffectedRowsResult.of( + "mongodb", 42L, Duration.ZERO); + when(engine.countAffectedRows(org.mockito.ArgumentMatchers.any())).thenReturn(expected); + + var result = executor.countAffectedRows(new QueryExecutionRequest(datasourceId, + "db.users.deleteMany({})", QueryType.DELETE, null, null)); + + assertThat(result).isSameAs(expected); + var captor = org.mockito.ArgumentCaptor.forClass( + com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest.class); + verify(engine).countAffectedRows(captor.capture()); + assertThat(captor.getValue().descriptor()).isSameAs(mongoDescriptor); + assertThat(captor.getValue().effectiveTimeout()).isEqualTo(Duration.ofSeconds(30)); + } + private static ResultSet emptyResultSet() throws SQLException { var rs = mock(ResultSet.class); var metadata = mock(ResultSetMetaData.class); diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolPropertiesTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolPropertiesTest.java index 709d05ce..2017ce02 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolPropertiesTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/ProxyPoolPropertiesTest.java @@ -14,7 +14,7 @@ class ProxyPoolPropertiesTest { @Test void allNullArgumentsYieldDocumentedDefaults() { - var props = new ProxyPoolProperties(null, null, null, null, null, null); + var props = new ProxyPoolProperties(null, null, null, null, null, null, null); assertThat(props.connectionTimeout()).isEqualTo(Duration.ofSeconds(30)); assertThat(props.idleTimeout()).isEqualTo(Duration.ofMinutes(10)); @@ -60,7 +60,8 @@ void explicitValuesPassThroughUnchanged() { Duration.ofSeconds(5), Duration.ofSeconds(60), Duration.ofMinutes(15), Duration.ofSeconds(2), "custom-", new Execution(500, Duration.ofSeconds(7), 250, 100, 1_048_576L, 8, - Duration.ofSeconds(2))); + Duration.ofSeconds(2)), + Duration.ofSeconds(3)); assertThat(props.connectionTimeout()).isEqualTo(Duration.ofSeconds(5)); assertThat(props.idleTimeout()).isEqualTo(Duration.ofSeconds(60)); diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListenerTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListenerTest.java new file mode 100644 index 00000000..98af565e --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/QueryCostEstimateListenerTest.java @@ -0,0 +1,27 @@ +package com.bablsoft.accessflow.proxy.internal; + +import com.bablsoft.accessflow.core.events.QuerySubmittedEvent; +import com.bablsoft.accessflow.proxy.api.QueryCostEstimateService; +import org.junit.jupiter.api.Test; + +import java.util.Optional; +import java.util.UUID; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class QueryCostEstimateListenerTest { + + @Test + void onSubmittedTriggersEstimate() { + var service = mock(QueryCostEstimateService.class); + var queryRequestId = UUID.randomUUID(); + when(service.estimateSubmittedQuery(queryRequestId)).thenReturn(Optional.empty()); + + new QueryCostEstimateListener(service) + .onSubmitted(new QuerySubmittedEvent(queryRequestId)); + + verify(service).estimateSubmittedQuery(queryRequestId); + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounterTest.java b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounterTest.java new file mode 100644 index 00000000..b37f4868 --- /dev/null +++ b/backend/src/test/java/com/bablsoft/accessflow/proxy/internal/dryrun/AffectedRowCounterTest.java @@ -0,0 +1,65 @@ +package com.bablsoft.accessflow.proxy.internal.dryrun; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class AffectedRowCounterTest { + + @Test + void updateWithWhereBecomesCountOverSamePredicate() { + var count = AffectedRowCounter.toCountSql( + "UPDATE users SET active = false WHERE last_login < '2020-01-01'"); + assertThat(count).hasValue( + "SELECT COUNT(*) FROM users WHERE last_login < '2020-01-01'"); + } + + @Test + void updateWithoutWhereCountsWholeTable() { + var count = AffectedRowCounter.toCountSql("UPDATE users SET active = false"); + assertThat(count).hasValue("SELECT COUNT(*) FROM users"); + } + + @Test + void deleteWithWhereBecomesCount() { + var count = AffectedRowCounter.toCountSql( + "DELETE FROM payroll.salaries WHERE year < 2019"); + assertThat(count).hasValue("SELECT COUNT(*) FROM payroll.salaries WHERE year < 2019"); + } + + @Test + void deleteKeepsAlias() { + var count = AffectedRowCounter.toCountSql("DELETE FROM users u WHERE u.active = false"); + assertThat(count).hasValue("SELECT COUNT(*) FROM users u WHERE u.active = false"); + } + + @Test + void updateWithJoinIsUnsupported() { + assertThat(AffectedRowCounter.toCountSql( + "UPDATE orders o JOIN users u ON o.user_id = u.id SET o.flag = 1")).isEmpty(); + } + + @Test + void updateFromIsUnsupported() { + assertThat(AffectedRowCounter.toCountSql( + "UPDATE orders SET flag = 1 FROM users WHERE orders.user_id = users.id")).isEmpty(); + } + + @Test + void deleteUsingIsUnsupported() { + assertThat(AffectedRowCounter.toCountSql( + "DELETE FROM orders USING users WHERE orders.user_id = users.id")).isEmpty(); + } + + @Test + void selectAndInsertAreUnsupported() { + assertThat(AffectedRowCounter.toCountSql("SELECT * FROM users")).isEmpty(); + assertThat(AffectedRowCounter.toCountSql( + "INSERT INTO users (id) VALUES (1)")).isEmpty(); + } + + @Test + void unparseableSqlIsUnsupported() { + assertThat(AffectedRowCounter.toCountSql("DELETE FROM WHERE nope !!")).isEmpty(); + } +} diff --git a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachineTest.java b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachineTest.java index 9592a066..3e454095 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachineTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/QueryReviewStateMachineTest.java @@ -64,6 +64,7 @@ class QueryReviewStateMachineTest { @Mock RoutingDecisionService routingDecisionService; @Mock com.bablsoft.accessflow.ai.api.BehaviorAnomalyLookupService behaviorAnomalyLookupService; @Mock AccessGrantLookupService accessGrantLookupService; + @Mock com.bablsoft.accessflow.core.api.QueryEstimateLookupService queryEstimateLookupService; @Mock MessageSource messageSource; @Mock ApplicationEventPublisher eventPublisher; @InjectMocks QueryReviewStateMachine stateMachine; diff --git a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionCodecTest.java b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionCodecTest.java index c9be702a..d06b5cd4 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionCodecTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionCodecTest.java @@ -49,7 +49,9 @@ void roundTripsEveryLeafAndCombinator() { new ConditionNode.SourceIpMatches(List.of("203.0.113.0/24", "2001:db8::/32")), new ConditionNode.UserAgentMatches(List.of("*curl*", "*GitHubActions*")), new ConditionNode.TimeSinceLastApproval(ComparisonOperator.GT, 1440), - new ConditionNode.CiCdOrigin(true))); + new ConditionNode.CiCdOrigin(true), + new ConditionNode.EstimatedRows(ComparisonOperator.GT, 100_000L), + new ConditionNode.ScanTypeMatches(List.of("Seq*", "COLLSCAN")))); var json = codec.encode(tree); var decoded = codec.decode(json); @@ -57,6 +59,16 @@ void roundTripsEveryLeafAndCombinator() { assertThat(decoded).isEqualTo(tree); } + @Test + void estimatedRowsAndScanTypeUseSnakeCaseDiscriminators() { + var json = codec.encode(new ConditionNode.And(List.of( + new ConditionNode.EstimatedRows(ComparisonOperator.GTE, 500_000L), + new ConditionNode.ScanTypeMatches(List.of("Seq Scan"))))); + + assertThat(json).contains("\"type\":\"estimated_rows\"") + .contains("\"type\":\"scan_type\""); + } + @Test void wireShapeForClientContextLeaves() { assertThat(codec.encode(new ConditionNode.SourceIpMatches(List.of("10.0.0.0/8")))) diff --git a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluatorTest.java b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluatorTest.java index 38795ec7..b5930af9 100644 --- a/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluatorTest.java +++ b/backend/src/test/java/com/bablsoft/accessflow/workflow/internal/routing/RoutingConditionEvaluatorTest.java @@ -292,4 +292,47 @@ void ciCdOrigin() { assertThat(evaluator.matches(new ConditionNode.CiCdOrigin(false), clientContext("203.0.113.7", null, false, null))).isTrue(); } + + @Test + void estimatedRows() { + var ctx = estimateContext(2_400_000L, "Seq Scan"); + assertThat(evaluator.matches( + new ConditionNode.EstimatedRows(ComparisonOperator.GT, 100_000), ctx)).isTrue(); + assertThat(evaluator.matches( + new ConditionNode.EstimatedRows(ComparisonOperator.LT, 100_000), ctx)).isFalse(); + assertThat(evaluator.matches( + new ConditionNode.EstimatedRows(ComparisonOperator.EQ, 2_400_000), ctx)).isTrue(); + } + + @Test + void estimatedRowsFailsClosedWhenNoEstimate() { + var ctx = estimateContext(null, null); + assertThat(evaluator.matches( + new ConditionNode.EstimatedRows(ComparisonOperator.GTE, 0), ctx)).isFalse(); + } + + @Test + void scanType() { + var ctx = estimateContext(500L, "Seq Scan"); + assertThat(evaluator.matches( + new ConditionNode.ScanTypeMatches(List.of("Seq*")), ctx)).isTrue(); + assertThat(evaluator.matches( + new ConditionNode.ScanTypeMatches(List.of("seq scan")), ctx)).isTrue(); + assertThat(evaluator.matches( + new ConditionNode.ScanTypeMatches(List.of("Index*")), ctx)).isFalse(); + } + + @Test + void scanTypeFailsClosedWhenNoPlan() { + var ctx = estimateContext(500L, null); + assertThat(evaluator.matches( + new ConditionNode.ScanTypeMatches(List.of("*")), ctx)).isFalse(); + } + + private ConditionContext estimateContext(Long estimatedRows, String scanType) { + return new ConditionContext(QueryType.DELETE, Set.of("payroll.salaries"), RiskLevel.HIGH, 82, + "ANALYST", Set.of(groupId), LocalDateTime.of(2026, 6, 3, 14, 30), + false, false, false, "203.0.113.7", "Mozilla/5.0 (Macintosh)", false, 120, + false, estimatedRows, scanType); + } } diff --git a/connectors/couchbase/connector.json b/connectors/couchbase/connector.json index 158e41b2..f6a8f3de 100644 --- a/connectors/couchbase/connector.json +++ b/connectors/couchbase/connector.json @@ -13,8 +13,8 @@ "bundled": false, "driver": { "type": "url", - "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-couchbase-1.0.4-all.jar", - "fileName": "accessflow-engine-couchbase-1.0.4-all.jar", - "sha256": "8753f5bd9f1f962df5ecb621d10fc8bcec8a65b4338da1e1197f82e8dbc86cb1" + "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-couchbase-1.0.5-all.jar", + "fileName": "accessflow-engine-couchbase-1.0.5-all.jar", + "sha256": "bd78c1e8733e28fadf95b01cb82acf6e17be790316c7a5f526aa38899259a085" } } diff --git a/connectors/elasticsearch/connector.json b/connectors/elasticsearch/connector.json index 28a4e163..f879af59 100644 --- a/connectors/elasticsearch/connector.json +++ b/connectors/elasticsearch/connector.json @@ -13,8 +13,8 @@ "bundled": false, "driver": { "type": "url", - "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-elasticsearch-1.0.3-all.jar", - "fileName": "accessflow-engine-elasticsearch-1.0.3-all.jar", - "sha256": "416568fc8ba21388ae37bc8eeccff7fcd6b656cf29fe1fd8e26fb1f325648886" + "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-elasticsearch-1.0.4-all.jar", + "fileName": "accessflow-engine-elasticsearch-1.0.4-all.jar", + "sha256": "18e6889b6ca3b9777a23a896e52677dd4ed3fa1e2799a412bbea1566feb24520" } } diff --git a/connectors/mongodb/connector.json b/connectors/mongodb/connector.json index 40e4ca41..f6ff317a 100644 --- a/connectors/mongodb/connector.json +++ b/connectors/mongodb/connector.json @@ -13,8 +13,8 @@ "bundled": false, "driver": { "type": "url", - "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-mongodb-1.0.6-all.jar", - "fileName": "accessflow-engine-mongodb-1.0.6-all.jar", - "sha256": "af2109db37dc804c3e57dae724f9948c7f519ffdb4f47fe7839d7c35078c5fa0" + "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-mongodb-1.0.7-all.jar", + "fileName": "accessflow-engine-mongodb-1.0.7-all.jar", + "sha256": "56922d9da7a6f7ed958a928ac2878c0a5d6d862e4aa235023bb78dce4cecf85f" } } diff --git a/connectors/neo4j/connector.json b/connectors/neo4j/connector.json index 2022a227..1cd45d8a 100644 --- a/connectors/neo4j/connector.json +++ b/connectors/neo4j/connector.json @@ -13,8 +13,8 @@ "bundled": false, "driver": { "type": "url", - "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-neo4j-1.0.2-all.jar", - "fileName": "accessflow-engine-neo4j-1.0.2-all.jar", - "sha256": "9cf3f86e833e2556b2d07fd52fc0d3cb65cdec64c3e8bc61d1f039ae9afaed02" + "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-neo4j-1.0.3-all.jar", + "fileName": "accessflow-engine-neo4j-1.0.3-all.jar", + "sha256": "be6078ffff601e4e4d45a2c03177874345dafb1153da3521536411ad7e0f9870" } } diff --git a/connectors/opensearch/connector.json b/connectors/opensearch/connector.json index b80ebe2e..0a5f4697 100644 --- a/connectors/opensearch/connector.json +++ b/connectors/opensearch/connector.json @@ -13,8 +13,8 @@ "bundled": false, "driver": { "type": "url", - "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-elasticsearch-1.0.2-all.jar", - "fileName": "accessflow-engine-elasticsearch-1.0.2-all.jar", - "sha256": "660fc466f090ca3539f20126d45fac9553a0d557a539f44745b08c6da20982ca" + "url": "https://bablsoft.github.io/accessflow/engines/accessflow-engine-elasticsearch-1.0.4-all.jar", + "fileName": "accessflow-engine-elasticsearch-1.0.4-all.jar", + "sha256": "18e6889b6ca3b9777a23a896e52677dd4ed3fa1e2799a412bbea1566feb24520" } } diff --git a/docs/03-data-model.md b/docs/03-data-model.md index 5b32e0d5..fc6b2b9a 100644 --- a/docs/03-data-model.md +++ b/docs/03-data-model.md @@ -477,6 +477,8 @@ The condition is a polymorphic, `"type"`-discriminated tree (snake_case, no exte | `time_since_last_approval` (AF-446) | `operator` (`LT`/`LTE`/`GT`/`GTE`/`EQ`), `minutes` | minutes since the requester's last APPROVED/EXECUTED query **on the same datasource** satisfy the comparison. **Fails closed**: false when the requester has no prior approval there | | `cicd_origin` (AF-446) | `expected: bool` | whether the request came from a CI/CD pipeline (submitted via an API key or with the `X-AccessFlow-CI` header) equals `expected`. Deterministic — the flag defaults to `false` | | `anomaly_detected` (AF-383) | `expected: bool` | whether the submitter currently has an `OPEN` `behavior_anomaly` on the target datasource equals `expected`. The UBA detector is a periodic batch over **past** data, so this signal escalates the flagged user's **next** query — pair it with `ESCALATE`. Deterministic — false when the user has no open anomaly there | +| `estimated_rows` (AF-624) | `operator` (`LT`/`LTE`/`GT`/`GTE`/`EQ`), `value` | the query's pre-flight estimated row impact (the exact affected-row count for UPDATE/DELETE when available, else the EXPLAIN estimate) satisfies the comparison. Read live from `query_estimates` at routing time. **Fails closed**: false when no estimate signal exists (not yet computed, unsupported engine, or failed) | +| `scan_type` (AF-624) | `patterns: [string]` | the pre-flight plan's root operation (e.g. `Seq Scan`, `COLLSCAN`) matches any glob (`*` wildcard, case-insensitive). **Fails closed**: false when no plan was captured | On the AI-skipped path (`datasource.ai_analysis_enabled=false`) the risk-based operands (`risk_level`, `risk_score`) evaluate to **false** — there is no AI signal. Routing is **not** run on the AI-failure path. @@ -648,6 +650,7 @@ The central entity. Represents a single SQL submission through the platform. | `submission_reason` | ENUM `submission_reason`: `USER_SUBMITTED` (default) \| `AI_SUGGESTION` (AF-451) \| `EMERGENCY_ACCESS` (AF-385, Flyway V93). `AI_SUGGESTION` marks a draft created by applying an AI optimization suggestion in the editor; `EMERGENCY_ACCESS` marks a query that bypassed pre-approval through the break-glass path. Recorded in the `QUERY_SUBMITTED` audit metadata. NOT NULL DEFAULT `'USER_SUBMITTED'`. | | `justification` | TEXT nullable — requester's stated reason for the query | | `ai_analysis_id` | FK → `ai_analyses` nullable | +| `query_estimate_id` | FK → `query_estimates` nullable (AF-624, Flyway `V128`) — the query's persisted pre-flight cost estimate; bare-UUID back-pointer mirroring `ai_analysis_id` (the real NOT NULL FK lives on `query_estimates.query_request_id`) | | `execution_started_at` | TIMESTAMPTZ nullable | | `execution_completed_at` | TIMESTAMPTZ nullable | | `rows_affected` | BIGINT nullable | @@ -745,6 +748,31 @@ Stores the result of an AI analysis run for a query request. --- +## query_estimates + +Persisted pre-flight cost / blast-radius estimate for a submitted query (AF-624, Flyway `V127`/`V128`). Computed asynchronously right after submission by the proxy's `QueryCostEstimateService` — the AF-445 dry-run machinery (relational dialect `EXPLAIN`, engine `dryRun` overrides) plus, for UPDATE/DELETE, a governed non-mutating affected-row count (relational `SELECT COUNT(*)` rewrite; MongoDB `countDocuments`, SQL++ `SELECT COUNT(*)`, Cypher `MATCH … RETURN count(*)`, Elasticsearch `_count` via the engines' `countAffectedRows` SPI overrides). One row per query request; every computation path persists a row (success, unsupported-engine degrade, or `failed=true` sentinel), mirroring `ai_analyses`. Read by reviewers (`GET /queries/{id}` → `cost_estimate`), by the routing-policy engine (`estimated_rows` / `scan_type` conditions), and folded into the AI analyzer's prompt (`{{cost_estimate}}`). + +| Column | Type / Notes | +|--------|-------------| +| `id` | UUID PK | +| `query_request_id` | FK → `query_requests` NOT NULL **UNIQUE** ON DELETE CASCADE (one estimate per query; deleted with its query) | +| `engine_id` | VARCHAR(64) nullable — connector id (e.g. `postgresql`, `mongodb`); null on the failed-sentinel path | +| `query_type` | ENUM `query_type` nullable | +| `supported` | BOOLEAN NOT NULL DEFAULT false — false when the engine has no plan concept or the statement shape isn't explainable | +| `estimated_rows` | BIGINT nullable — the plan's root row estimate | +| `affected_row_count` | BIGINT nullable — exact governed count for UPDATE/DELETE; null when the shape can't be provably counted (joins, `USING`, MERGE, …) or the engine doesn't support counting | +| `scan_type` | VARCHAR(128) nullable — the plan's root operation (e.g. `Seq Scan`, `COLLSCAN`) | +| `estimated_cost` | DOUBLE PRECISION nullable — the plan's root cost when the engine exposes one | +| `plan` | JSONB nullable — the snake_case plan-node tree (same shape as the dry-run endpoint's `plan`) | +| `raw_plan` | TEXT nullable — the engine's raw plan output | +| `unsupported_reason` | VARCHAR(500) nullable — localized reason when `supported=false` | +| `failed` | BOOLEAN NOT NULL DEFAULT false — true when the computation hit an unexpected error (sentinel row) | +| `error_message` | VARCHAR(500) nullable — failure reason when `failed=true` | +| `duration_ms` | INTEGER nullable | +| `created_at` | TIMESTAMPTZ | + +--- + ## ai_config Per-organization AI provider configurations. Many rows per organization — admins create diff --git a/docs/04-api-spec.md b/docs/04-api-spec.md index 7d5d6c2f..e49fffc3 100644 --- a/docs/04-api-spec.md +++ b/docs/04-api-spec.md @@ -1313,6 +1313,22 @@ Each subsequent row contains the same fields as `QueryListItemView`. `ai_risk_le "failed": false, "error_message": null }, + "cost_estimate": { + "id": "uuid", + "engine_id": "postgresql", + "query_type": "UPDATE", + "supported": true, + "estimated_rows": 1, + "affected_row_count": 1, + "scan_type": "Index Scan", + "estimated_cost": 8.29, + "plan": { "operation": "Index Scan", "target": "orders", "estimated_rows": 1, "estimated_cost": 8.29, "detail": "(id = 123)", "children": [] }, + "raw_plan": "[...]", + "unsupported_reason": null, + "failed": false, + "error_message": null, + "duration_ms": 12 + }, "review_decisions": [ { "id": "uuid", @@ -1364,6 +1380,8 @@ Each subsequent row contains the same fields as `QueryListItemView`. `ai_risk_le `matched_policy` is the routing policy that decided this query's routing (AF-379); `null` when no policy matched and the query fell through to the datasource's review plan. `policy_name` is `null` when the matched policy was later deleted. The frontend renders a "matched policy" alert on the detail page when this object is present. See [docs/05-backend.md → "Policy-as-code routing engine"](05-backend.md#policy-as-code-routing-engine-af-379). +`cost_estimate` is the query's persisted pre-flight cost / blast-radius estimate (AF-624), computed automatically and asynchronously right after submission — `null` while it is still being computed (typically only during `PENDING_AI`). `supported=false` with a localized `unsupported_reason` marks engines/statement shapes with no plan concept; `failed=true` with `error_message` marks a computation error. `affected_row_count` is the exact governed row count for UPDATE/DELETE (relational `SELECT COUNT(*)` rewrite, or the engine's native non-mutating count) and is `null` when the shape cannot be provably counted; `estimated_rows`/`scan_type`/`estimated_cost` come from the plan root, and `plan` reuses the dry-run endpoint's plan-node shape — see [POST /queries/dry-run](#post-queriesdry-run--response-200) and [docs/05-backend.md → "Automatic pre-flight cost estimate"](05-backend.md#automatic-pre-flight-cost-estimate-af-624). The `query.estimate_complete` WebSocket event signals completion. + `approved_by_grant` is the provenance of a grant-covered auto-approval (#582); `null` for every query that was not auto-approved by a pre-approving JIT access grant. `grant_id` always reflects the query's persisted `approved_by_grant_id`; the approver fields (`approver_id`, `approver_email`, `approved_at` — the grant's final-stage approval) and `expires_at` are resolved from the grant row at read time and are `null` when the grant row or its approver is no longer resolvable. The frontend renders an "auto-approved under an access grant" alert on the detail page when this object is present. See [docs/05-backend.md → "Grant-covered query auto-approval"](05-backend.md#grant-covered-query-auto-approval-582). `linked_tickets` lists the tickets auto-created in an external ticketing system (ServiceNow / Jira, AF-453) for this query's workflow events, oldest first — empty array when none. `system` is `SERVICENOW` | `JIRA`; `status` / `resolution` reflect the external system's labels as last synced by the [ticketing inbound webhook](08-notifications.md#ticketing-inbound-webhooks--bi-directional-sync-af-453). @@ -4536,6 +4554,7 @@ Clients subscribe to real-time updates for their own queries and (for reviewers) | `review.decision_made` | Reviewer approved/rejected submitter's query | `query_id`, `decision`, `reviewer`, `comment` | | `query.executed` | Execution completed | `query_id`, `rows_affected`, `duration_ms` | | `ai.analysis_complete` | AI analysis finished | `query_id`, `risk_level`, `risk_score` | +| `query.estimate_complete` | The automatic pre-flight cost estimate finished (AF-624) — detail views refetch to render the cost-estimate panel | `query_id`, `supported` | | `notification.created` | A new in-app notification was persisted for the caller | `notification_id`, `event_type`, `query_id`, `created_at` | | `access_request.created` | New JIT access request needs a reviewer's decision | `access_request_id`, `requester_id` | | `access_request.status_changed` | Access request changed status (approved/rejected/expired/revoked/cancelled) — pushed to the requester | `access_request_id`, `old_status`, `new_status` | diff --git a/docs/05-backend.md b/docs/05-backend.md index 2e03e322..0ebef174 100644 --- a/docs/05-backend.md +++ b/docs/05-backend.md @@ -817,6 +817,16 @@ The result is a `SelectExecutionResult` mapped to `SampleRowsResponse` for `GET The statement-timeout cap reuses `ACCESSFLOW_PROXY_EXECUTION_STATEMENT_TIMEOUT`; there is no row cap (a dry-run returns no rows). The result is mapped to `QueryDryRunResponse` by the controller in the `security` module (which already depends on `proxy`, so it can host the `/queries/dry-run` endpoint and use `JwtClaims` without a module cycle — the same arrangement as the sample-rows endpoint). +### Automatic pre-flight cost estimate (AF-624) + +`proxy.api.QueryCostEstimateService` (`DefaultQueryCostEstimateService`) turns the AF-445 dry-run machinery into an **automatic, persisted, per-submission blast-radius estimate**: right after a query is submitted, the engine's non-committing plan (estimated rows, root scan type, cost, plan tree) plus — for UPDATE/DELETE — a governed, non-mutating **exact affected-row count** are computed and stored as the query's single `query_estimates` row (see [docs/03-data-model.md → query_estimates](03-data-model.md#query_estimates)). + +- **Two independent triggers, safe by idempotency.** `proxy.internal.QueryCostEstimateListener` consumes `QuerySubmittedEvent` (mirroring the AI module's `AiAnalysisListener`) and runs unconditionally — reviewers and routing want the estimate regardless of `ai_analysis_enabled`. Independently, `DefaultAiAnalyzerService.analyzeSubmittedQuery` calls `estimateSubmittedQuery(id)` itself before building its prompt, so the estimate is deterministically available for `{{cost_estimate}}` no matter which trigger wins the race. The service fast-paths on an existing row and `DefaultQueryEstimateService.persist` is insert-once, so the race is harmless. +- **Computation.** The submitter's row-security directives are resolved (`RowSecurityResolutionService`) so the plan and count reflect the *governed* statement, then `QueryExecutor.dryRun` runs the existing per-dialect EXPLAIN / engine `dryRun` path. For UPDATE/DELETE, `QueryExecutor.countAffectedRows` additionally computes the exact count: **relational** datasources rewrite the parsed single-table statement into `SELECT COUNT(*) FROM [WHERE …]` (`proxy.internal.dryrun.AffectedRowCounter` — join / `UPDATE … FROM` / `DELETE … USING` shapes degrade to null) and run it through the normal SELECT path (so `RowSecurityRewriter` applies); **engine-managed** datasources delegate to the `QueryEngine.countAffectedRows` SPI default method (overridden by MongoDB `countDocuments`, Couchbase SQL++ `SELECT COUNT(*)` splice, Neo4j `MATCH … RETURN count(*)`, Elasticsearch/OpenSearch `_count` — each applying its native row security and failing closed on uncountable shapes; Redis, Cassandra/ScyllaDB, DynamoDB, and the warehouse engines inherit the unsupported default). +- **Bounded.** Both calls run under the dedicated `accessflow.proxy.estimate-timeout` (`ACCESSFLOW_PROXY_ESTIMATE_TIMEOUT`, default `PT5S`) instead of the full execution statement timeout — an estimate is a best-effort signal, never worth a long lock. +- **Every path persists a row.** Success stores the full estimate; engines with no plan concept store `supported=false` + a localized `unsupported_reason` (a transactional `BEGIN…COMMIT` envelope short-circuits the same way); an unexpected error stores a `failed=true` sentinel with the message — mirroring the AI module's sentinel convention, so the frontend can always render a definitive state. Completion publishes `QueryEstimateCompletedEvent` (or `QueryEstimateFailedEvent`), which the realtime module fans out as the `query.estimate_complete` WebSocket event. +- **Consumers.** `GET /queries/{id}` embeds the row as `cost_estimate`; `QueryReviewStateMachine.buildContext` reads it live (fail-closed) for the `estimated_rows` / `scan_type` routing conditions; `DefaultAiAnalyzerService` renders it into the `{{cost_estimate}}` prompt placeholder. + ### Data classification & derivation (AF-447) `data_classification_tag` rows (see [docs/03-data-model.md](03-data-model.md)) tag tables/columns with @@ -978,7 +988,7 @@ Decision rules: Routing policies are ordered, attribute-based rules that decide how a submitted query is routed **before** the default review-plan logic runs. The engine is owned by the `workflow` module and evaluated inside the same `QueryReviewStateMachine` listener, **after** AI analysis (or the skip event) and **before** reviewer fan-out: -1. `RoutingPolicyEngine` loads the org's enabled policies (org-wide + this datasource) in ascending `priority` and evaluates each `condition` against the query context (query type, referenced tables, AI risk level / score, requester role + group memberships, time-of-day / day-of-week, WHERE / LIMIT presence, transactional flag, and the client context captured at submission — source IP / CIDR, user-agent, time-since-last-approval, CI/CD origin) via `RoutingConditionEvaluator`. +1. `RoutingPolicyEngine` loads the org's enabled policies (org-wide + this datasource) in ascending `priority` and evaluates each `condition` against the query context (query type, referenced tables, AI risk level / score, requester role + group memberships, time-of-day / day-of-week, WHERE / LIMIT presence, transactional flag, the pre-flight cost estimate — estimated/affected rows and root scan type, read live from `query_estimates` and fail-closed when absent (AF-624) — and the client context captured at submission — source IP / CIDR, user-agent, time-since-last-approval, CI/CD origin) via `RoutingConditionEvaluator`. 2. **First match wins.** The first enabled policy whose condition matches decides the action; evaluation stops there. On **no match** the grant-covered auto-approval fast-path (#582, see the [JIT section](#grant-covered-query-auto-approval-582)) is consulted next, and only then does the query fall through to the datasource's review plan exactly as before — so **any** matching policy (AUTO_REJECT, REQUIRE_APPROVALS, ESCALATE — including anomaly-driven ones) always wins over the grant fast-path. 3. The outcome (matched policy id, action, resolved `effective_min_approvals`, reason) is persisted as a single `routing_decision` row (`RoutingDecisionService`), and surfaced on `GET /queries/{id}` as `matched_policy`. @@ -1498,9 +1508,10 @@ Spring context refresh. ### Editable system prompt `SystemPromptRenderer` holds the built-in analyzer prompt (`DEFAULT_TEMPLATE`) and renders it -with four named placeholders substituted at call time: `{{db_type}}`, `{{schema_context}}`, -`{{sql}}` and `{{language}}`. `{{sql}}` is replaced last so SQL text that happens to contain -another token string is never re-substituted. +with named placeholders substituted at call time: `{{db_type}}`, `{{schema_context}}`, +`{{rag_context}}`, `{{cost_estimate}}` (AF-624 — the pre-flight estimate summary, falling back to +`(no cost estimate available)`), `{{sql}}` and `{{language}}`. `{{sql}}` is replaced last so SQL +text that happens to contain another token string is never re-substituted. Admins may override the prompt per `ai_config` row via the `system_prompt_template` column (`NULL`/blank ⇒ use `DEFAULT_TEMPLATE`). `DefaultAiConfigService` validates that a custom template @@ -1641,6 +1652,8 @@ Optimization suggestions: when the query would benefit from an index or a rewrit Database type: {db_type} Schema context: {schema_context} +Pre-flight cost estimate (from the database engine's own EXPLAIN / affected-row count — treat it as the authoritative blast radius and factor it into risk_score and risk_level): +{cost_estimate} SQL to analyze: {sql} ``` @@ -1649,7 +1662,7 @@ SQL to analyze: ### Response language -`AiAnalyzerStrategy.analyze(sql, dbType, schemaContext, language)` takes a BCP-47 code (`en`, `es`, `de`, `fr`, `zh-CN`, `ru`, `hy`). The renderer appends one line at the end of the user prompt: `Respond in: . Translate the free-form fields (summary, issues[].message, issues[].suggestion) into that language. Keep risk_level and issues[].category as their original English enum values.` +`AiAnalyzerStrategy.analyze(sql, dbType, schemaContext, costEstimateContext, language, aiConfigId)` takes a BCP-47 `language` code (`en`, `es`, `de`, `fr`, `zh-CN`, `ru`, `hy`). The renderer appends one line at the end of the user prompt: `Respond in: . Translate the free-form fields (summary, issues[].message, issues[].suggestion) into that language. Keep risk_level and issues[].category as their original English enum values.` `DefaultAiAnalyzerService` resolves the language per call by reading the org's `localization_config.ai_review_language` via `LocalizationConfigService.getOrDefault(organizationId)`. If the lookup fails or returns an unknown code the service silently falls back to English so prompt construction never blocks AI analysis. The `/admin/ai-config/test` smoke endpoint always passes `"en"` since it is a synthetic, language-agnostic call. @@ -2029,6 +2042,7 @@ Browsers cannot set a custom `Authorization` header on a WebSocket upgrade, so t | `query.status_changed` | `QueryStatusChangedEvent` (in `core/events/`) | submitter | | `query.executed` | `QueryExecutedEvent` (in `workflow/events/`) | submitter | | `ai.analysis_complete` | `AiAnalysisCompletedEvent` (in `core/events/`) | submitter | +| `query.estimate_complete` | `QueryEstimateCompletedEvent` / `QueryEstimateFailedEvent` (in `core/events/`, AF-624) | submitter | | `review.new_request` | `QueryReadyForReviewEvent` (in `core/events/`) | eligible reviewers | | `review.decision_made` | `ReviewDecisionMadeEvent` (in `workflow/events/`) | submitter | | `notification.created` | `UserNotificationCreatedEvent` (in `notifications/events/`) | the recipient user | diff --git a/docs/06-frontend.md b/docs/06-frontend.md index c798e0f8..fa1a94e1 100644 --- a/docs/06-frontend.md +++ b/docs/06-frontend.md @@ -240,6 +240,7 @@ Full detail view for any query: - When `status === 'TIMED_OUT'`, a warning callout above the SQL block names the review plan, the configured `approval_timeout_hours`, and how long ago the timeout fired. The metadata sidebar surfaces `plan` / `timeout.hours` for any query whose datasource has a review plan, regardless of status. Status-pill colour and label come from `statusColors.ts` (`TIMED_OUT` → warn-amber palette, label "TIMED OUT"). - When `ai_analysis.failed === true` (AF-249), a warning `Alert` at the top of the main column tells the reviewer that AI analysis didn't complete and that review is proceeding without an AI recommendation; the analyzer's reason is shown both in the banner detail and in a dedicated failure variant of the AI accordion. The `RiskPill` in the accordion header switches to a neutral grey "AI N/A" variant (`failed` prop on `RiskPill`). For `REVIEWER` / `ADMIN` callers a primary "Re-analyze" button (in both the banner and the accordion) calls `POST /queries/{id}/reanalyze`; the page invalidates its TanStack Query entries on success and picks up the new analysis via the existing `ai.analysis_complete` WebSocket event. The list page (`QueryListPage`) renders the same "AI N/A" pill in the risk column when `ai_failed=true` on the list row, so a CRITICAL-looking sentinel is never mistaken for a real risk verdict. - When the latest entry in `review_decisions[]` has `decision: REQUESTED_CHANGES` AND the query is still `PENDING_REVIEW` (AF-269), an info `Alert` at the top of the main column tells the submitter that the reviewer asked for changes — body interpolates `{{reviewer}}`, `{{when}}`, and `{{comment}}`. The reviewer decision panel itself requires a non-empty comment for both **Reject** (disabled until typed) and **Request changes** (already disabled); approving still allows an empty comment. The rejected stage of `ApprovalTimeline` carries the last `REJECTED` decision's comment (wrapped in `"…"` so the existing italic style in [ApprovalTimeline.tsx](../frontend/src/components/review/ApprovalTimeline.tsx) applies). +- **Cost-estimate card** (AF-624) — a "Cost estimate" `DetailCard` (`components/review/CostEstimatePanel.tsx`) renders the query's persisted pre-flight blast-radius estimate from `cost_estimate` on `GET /queries/{id}`: the exact affected-row count for UPDATE/DELETE ("Affected rows (exact)"), the plan's estimated rows / scan type / cost, and the execution-plan tree (reusing the editor's `PlanTree` + `utils/queryPlan.ts`), falling back to the raw plan text. State machine mirrors the AI card: while `status === 'PENDING_AI'` and `cost_estimate` is null it shows "Computing the cost estimate…"; a null estimate past that shows "No cost estimate is available for this query."; `supported=false` renders the localized `unsupported_reason` (still showing the exact count when one was computed); `failed=true` renders a warning with `error_message`. The `query.estimate_complete` WebSocket event invalidates `['queries','detail',id]` so the panel fills in without polling. - When `ai_analysis === null` and the query has already advanced out of `PENDING_AI` (AF-307), the AI step is rendered as **bypassed** rather than waiting. The card title becomes "AI analysis (skipped)" with a muted body — "AI analysis was skipped — this datasource has AI analysis disabled." — and the `ApprovalTimeline` shows a gray stage labeled "AI analysis skipped" (dot uses `--fg-muted`). The skipped state is derived on the frontend (`!ai_analysis && status !== 'PENDING_AI'`); the backend persists no `ai_analyses` row on the skip path. While the query is still in `PENDING_AI`, the original "Awaiting analysis…" fallback continues to render. ### DatasourceCreateWizardPage *(ADMIN)* @@ -651,6 +652,7 @@ useEffect(() => | `query.status_changed` | `['queries','detail',query_id]` and `['queries','list']` | | `query.executed` | `['queries','detail',query_id]` and `['queries','list']` | | `ai.analysis_complete` | `['queries','detail',query_id]` | +| `query.estimate_complete` | `['queries','detail',query_id]` | | `review.new_request` | `['reviews','pending']` | | `review.decision_made` | `['reviews','pending']` and `['queries','detail',query_id]` | | `notification.created` | `['notifications','list']` and `['notifications','unread-count']` | @@ -834,8 +836,10 @@ conditions, each optionally negated (NOT); there is no raw-JSON editor. API acce [frontend/src/api/routingPolicies.ts](../frontend/src/api/routingPolicies.ts); the form↔wire mapping helper is [frontend/src/pages/admin/routingPolicyForm.ts](../frontend/src/pages/admin/routingPolicyForm.ts); types (`RoutingPolicy`, `RoutingCondition`, `RoutingAction`, …) live in `src/types/api.ts`. -`QueryDetailPage` shows a **matched-policy** alert when `GET /queries/{id}` returns a non-null -`matched_policy`. +The builder covers the `estimated_rows` (comparison operator + row count) and `scan_type` +(glob tags, e.g. `Seq*`, `COLLSCAN`) pre-flight-estimate operands (AF-624) alongside the +original leaf set. `QueryDetailPage` shows a **matched-policy** alert when `GET /queries/{id}` +returns a non-null `matched_policy`. ### OAuth 2.0 sign-in diff --git a/docs/09-deployment.md b/docs/09-deployment.md index 9722a821..81334fcb 100644 --- a/docs/09-deployment.md +++ b/docs/09-deployment.md @@ -983,6 +983,7 @@ Deployment-wide tuning for the `ai` module's `BehaviorAnomalyDetectionJob`, whic | `ACCESSFLOW_PROXY_LEAK_DETECTION_THRESHOLD` | Optional | `0s` | HikariCP leak-detection threshold (`0s` disables) | | `ACCESSFLOW_PROXY_EXECUTION_MAX_ROWS` | Optional | `10000` | Hard cap on rows returned by a single query execution | | `ACCESSFLOW_PROXY_EXECUTION_STATEMENT_TIMEOUT` | Optional | `30s` | Statement-level timeout applied to customer-DB JDBC statements | +| `ACCESSFLOW_PROXY_ESTIMATE_TIMEOUT` | Optional | `PT5S` | Statement timeout for the automatic pre-flight cost estimate (AF-624) — the async EXPLAIN + affected-row COUNT run after each submission use this tighter bound instead of the full execution statement timeout | | `ACCESSFLOW_PROXY_EXECUTION_DEFAULT_FETCH_SIZE` | Optional | `1000` | Default JDBC fetch size | | `ACCESSFLOW_PROXY_EXECUTION_INSERT_BATCH_CHUNK_SIZE` | Optional | `1000` | Rows per JDBC `executeBatch()` flush when a `BEGIN…COMMIT` envelope's homogeneous single-row INSERTs are batched (AF-457) | | `ACCESSFLOW_PROXY_EXECUTION_MAX_RESULT_BYTES` | Optional | `52428800` | Per-result byte cap enforced while materializing SELECT rows (#49); when exceeded the result is truncated with `truncated_reason=BYTE_LIMIT`. Relational JDBC path only | diff --git a/docs/15-engine-sdk.md b/docs/15-engine-sdk.md index 62c5ba3d..de9191a3 100644 --- a/docs/15-engine-sdk.md +++ b/docs/15-engine-sdk.md @@ -55,6 +55,7 @@ A plugin implements **`com.bablsoft.accessflow.core.api.QueryEngine`**: | `execute(QueryEngineExecutionRequest)` | Run an approved query with the host-computed row cap / timeout → `QueryExecutionResult` | | `sampleTable(QueryEngineSampleRequest)` | Read a bounded, governance-applied sample of one table/collection (AF-443) — issue the engine's native "read all rows, capped at N" and funnel it through the same row-security + masking path as `execute` → `SelectExecutionResult`. Engines whose row-security model has no per-row meaning (key-value prefixes) must **fail closed** when a directive applies. | | `dryRun(QueryEngineDryRunRequest)` | **Default method** (AF-445). Return a non-committing execution plan + best-effort estimated rows → `QueryDryRunResult`, applying the request's row-security directives so the plan reflects the governed query. **Must never execute or mutate** — plan/estimate only (e.g. Mongo `explain` queryPlanner, Couchbase / Neo4j `EXPLAIN`, Elasticsearch `_validate/query?explain`). The default returns `QueryDryRunResult.unsupported(engineId())`, so an engine with no plan concept (Redis, Cassandra/ScyllaDB, DynamoDB) degrades gracefully **without overriding** — and, critically, without changing its compiled bytecode, so its connector pin needs no re-bump. Only engines that override `dryRun` re-pin. | +| `countAffectedRows(QueryEngineDryRunRequest)` | **Default method** (AF-624). Count the rows the request's UPDATE/DELETE would affect → `QueryAffectedRowsResult`, applying the request's row-security directives, **without executing or mutating** — the engine's native non-mutating count (MongoDB `countDocuments`, SQL++ `SELECT COUNT(*)` splice, Cypher `MATCH … RETURN count(*)`, Elasticsearch `_count`). Fail closed — return `QueryAffectedRowsResult.unsupported(engineId())` — for any statement shape whose count is not provably identical to the write's affected rows (joins, MERGE, LIMIT-capped DML, unspliceable row security). The default returns unsupported, so engines with no count concept (Redis, Cassandra/ScyllaDB, DynamoDB, the warehouses) degrade gracefully without overriding — and without a connector re-pin. Only engines that override re-pin. | | `testConnection(descriptor)` | Admin "Test connection" → `ConnectionTestResult` | | `introspectSchema(descriptor)` | Schema/collection sampling → `DatabaseSchemaView` (drives the ER diagram, autocomplete, AI schema context) | | `evictDatasource(id)` | Drop cached clients when a datasource's config changes or it is deactivated | diff --git a/e2e/tests/query-submit.spec.ts b/e2e/tests/query-submit.spec.ts index 26fb7624..fe3573e6 100644 --- a/e2e/tests/query-submit.spec.ts +++ b/e2e/tests/query-submit.spec.ts @@ -124,6 +124,14 @@ test.describe.serial('query submission from /editor', () => { // (statusColors.ts:22), so we assert on the visible text. await expect(page.getByRole('heading', { level: 1 }).getByText('Pending review')) .toBeVisible({ timeout: 15_000 }); + + // AF-624: the pre-flight cost estimate is computed asynchronously right after + // submission and lands on the detail page via the query.estimate_complete WS + // invalidation (or is already present on first load — EXPLAIN is fast). For + // SELECT 1 on the real Postgres the plan is supported, so the estimate card + // shows the "Estimated rows" stat instead of a pending/unavailable fallback. + await expect(page.getByText('Cost estimate')).toBeVisible(); + await expect(page.getByText('Estimated rows')).toBeVisible({ timeout: 20_000 }); }); // ── 2. Empty SQL — Submit is disabled until at least one char is typed ─────────── diff --git a/engines/couchbase/dependency-reduced-pom.xml b/engines/couchbase/dependency-reduced-pom.xml index 80475e81..a8432d74 100644 --- a/engines/couchbase/dependency-reduced-pom.xml +++ b/engines/couchbase/dependency-reduced-pom.xml @@ -4,7 +4,7 @@ com.bablsoft.accessflow accessflow-engine-couchbase AccessFlow Couchbase Query Engine - 1.0.4 + 1.0.5 On-demand Couchbase engine plugin for AccessFlow (issue AF-412). Implements the core.api.QueryEngine SPI for SQL++ (N1QL) and ships as a self-contained shaded JAR (bundling the Couchbase Java SDK with a relocated Reactor and Jackson) resolved through diff --git a/engines/couchbase/pom.xml b/engines/couchbase/pom.xml index 8b22d721..1bf368ab 100644 --- a/engines/couchbase/pom.xml +++ b/engines/couchbase/pom.xml @@ -12,7 +12,7 @@ (connectors/couchbase/connector.json) pins this version + the SHA-256 of the shaded JAR; bump BOTH together whenever the engine (or a core.api type it compiles against) changes. --> - 1.0.4 + 1.0.5 jar AccessFlow Couchbase Query Engine diff --git a/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriter.java b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriter.java new file mode 100644 index 00000000..675f3eb6 --- /dev/null +++ b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriter.java @@ -0,0 +1,134 @@ +package com.bablsoft.accessflow.engine.couchbase; + +import com.bablsoft.accessflow.engine.couchbase.SqlPlusPlusTokenizer.Kind; +import com.bablsoft.accessflow.engine.couchbase.SqlPlusPlusTokenizer.Token; + +import java.util.List; +import java.util.Set; + +/** + * Rewrites a parsed UPDATE / DELETE into the non-mutating {@code SELECT COUNT(*)} probe backing + * {@code countAffectedRows} (issue AF-624): the statement's target keyspace (original source text, + * so case and backticks survive) and its top-level WHERE clause are carried over verbatim — + * {@code SELECT COUNT(*) AS af_count FROM [AS ] [WHERE ]} — and the + * row-security applier then splices the request's directives into the count statement exactly as + * the execute path would (same alias, same qualifier). Only the shapes the WHERE-splice rewriter + * can provably reach are countable; anything else (CTE / subquery / JOIN / NEST / UNNEST / + * USE KEYS / set operation / multi-keyspace / MERGE) returns {@code null} so the caller degrades + * to {@code unsupported}. A {@code LIMIT} also returns {@code null}: it caps the mutation count, + * which a plain COUNT(*) would overstate. + */ +final class CouchbaseCountRewriter { + + static final String COUNT_ALIAS = "af_count"; + + /** Keywords that can follow the WHERE clause at depth 0 (= the clause boundary). */ + private static final Set WHERE_TAIL = Set.of( + "GROUP", "LETTING", "HAVING", "WINDOW", "ORDER", "LIMIT", "OFFSET", "RETURNING"); + + private CouchbaseCountRewriter() { + } + + /** The count statement, or {@code null} when the shape cannot be provably counted. */ + static String toCountStatement(CouchbaseStatement statement) { + if (statement.kind() != CouchbaseStatementKind.UPDATE + && statement.kind() != CouchbaseStatementKind.DELETE) { + return null; + } + if (statement.hasCte() || statement.hasSubquery() || statement.hasJoinLike() + || statement.hasUseKeys() || statement.hasSetOperation() || statement.hasLimit() + || statement.target() == null || statement.keyspaces().size() > 1) { + return null; + } + var target = targetSourceText(statement); + if (target == null) { + return null; + } + var count = new StringBuilder("SELECT COUNT(*) AS ").append(COUNT_ALIAS) + .append(" FROM ").append(target); + if (statement.targetAlias() != null) { + count.append(" AS `").append(statement.targetAlias().replace("`", "``")).append('`'); + } + var where = whereClauseText(statement); + if (where != null) { + count.append(" WHERE ").append(where); + } + return count.toString(); + } + + /** + * The original source text of the target keyspace ref (after FROM for a DELETE, after the + * UPDATE verb), including an optional {@code default:} namespace prefix — extracted by token + * offsets so backticks and case are preserved exactly as submitted. + */ + private static String targetSourceText(CouchbaseStatement statement) { + var tokens = statement.tokens(); + int refStart = refStartIndex(statement, tokens); + if (refStart < 0 || refStart >= tokens.size()) { + return null; + } + int i = refStart; + // Optional namespace prefix: default:. + if (tokens.get(i).isWord("DEFAULT") && i + 1 < tokens.size() + && tokens.get(i + 1).isSymbol(":")) { + i += 2; + } + int end = -1; + while (i < tokens.size()) { + var token = tokens.get(i); + if (token.kind() != Kind.WORD && token.kind() != Kind.QUOTED_IDENT) { + break; + } + end = token.end(); + if (i + 1 < tokens.size() && tokens.get(i + 1).isSymbol(".")) { + i += 2; + continue; + } + break; + } + if (end < 0) { + return null; + } + return statement.sql().substring(tokens.get(refStart).start(), end); + } + + private static int refStartIndex(CouchbaseStatement statement, List tokens) { + if (statement.kind() == CouchbaseStatementKind.UPDATE) { + // hasCte shapes were rejected above, so the UPDATE verb is the first token. + return 1; + } + for (int i = 0; i < tokens.size(); i++) { + var token = tokens.get(i); + if (token.depth() == 0 && token.isWord("FROM")) { + return i + 1; + } + } + return -1; + } + + /** The top-level WHERE expression text, or {@code null} when the statement has none. */ + private static String whereClauseText(CouchbaseStatement statement) { + var tokens = statement.tokens(); + for (int i = 0; i < tokens.size(); i++) { + var token = tokens.get(i); + if (token.depth() == 0 && token.isWord("WHERE")) { + var clause = statement.sql() + .substring(token.end(), clauseEndOffset(tokens, i + 1)).strip(); + return clause.isEmpty() ? null : clause; + } + } + return null; + } + + /** Offset of the first depth-0 tail keyword at/after {@code from}, else the statement end. */ + private static int clauseEndOffset(List tokens, int from) { + for (int i = from; i < tokens.size(); i++) { + var token = tokens.get(i); + if (token.depth() == 0 && token.kind() == Kind.WORD + && WHERE_TAIL.contains(token.value())) { + return token.start(); + } + } + return tokens.get(tokens.size() - 1).end(); + } +} diff --git a/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngine.java b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngine.java index 257af1a8..2ae1ff42 100644 --- a/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngine.java +++ b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngine.java @@ -3,6 +3,7 @@ import com.bablsoft.accessflow.core.api.ConnectionTestResult; import com.bablsoft.accessflow.core.api.DatabaseSchemaView; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngine; import com.bablsoft.accessflow.core.api.QueryEngineContext; @@ -86,6 +87,12 @@ public QueryDryRunResult dryRun(QueryEngineDryRunRequest request) { request.effectiveTimeout()); } + @Override + public QueryAffectedRowsResult countAffectedRows(QueryEngineDryRunRequest request) { + return initialized(executor).countAffectedRows(request.request(), request.descriptor(), + request.effectiveTimeout()); + } + @Override public ConnectionTestResult testConnection(DatasourceConnectionDescriptor descriptor) { return initialized(connectionProbe).test(descriptor); diff --git a/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutor.java b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutor.java index e843dde0..791c4513 100644 --- a/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutor.java +++ b/engines/couchbase/src/main/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutor.java @@ -1,12 +1,15 @@ package com.bablsoft.accessflow.engine.couchbase; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.InvalidSqlException; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; import com.bablsoft.accessflow.core.api.QueryType; import com.bablsoft.accessflow.core.api.SampleTableRequest; import com.bablsoft.accessflow.core.api.SelectExecutionResult; +import com.bablsoft.accessflow.core.api.UnrewritableRowSecurityException; import com.bablsoft.accessflow.core.api.UpdateExecutionResult; import com.bablsoft.accessflow.engine.couchbase.CouchbaseRowSecurityApplier.Applied; import com.couchbase.client.core.error.CouchbaseException; @@ -21,6 +24,7 @@ import java.time.Instant; import java.util.ArrayList; import java.util.List; +import java.util.Map; /** * Executes a SQL++ statement for a {@code COUCHBASE} datasource — the document-engine analogue of @@ -106,6 +110,61 @@ QueryDryRunResult dryRun(QueryExecutionRequest request, } } + /** + * Governed affected-row count (issue AF-624): the submitted UPDATE / DELETE is rewritten by + * {@link CouchbaseCountRewriter} into a {@code SELECT COUNT(*)} over the same keyspace and + * WHERE clause, the request's row-security directives are spliced into the count statement + * exactly as {@link #execute} would splice them, and the count runs read-only through the + * bucket's default-scope query context with the same scan consistency and timeout as a normal + * read. Shape reasons never throw — non-UPDATE/DELETE types, parse failures, and every shape + * the rewriter or the row-security splicer rejects degrade to {@code unsupported}; genuine + * server errors propagate through the usual exception translation, like {@link #dryRun}. + */ + QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request, + DatasourceConnectionDescriptor descriptor, + Duration timeout) { + var start = clock.instant(); + if (request.queryType() != QueryType.UPDATE && request.queryType() != QueryType.DELETE) { + return QueryAffectedRowsResult.unsupported(CouchbaseQueryEngine.ENGINE_ID); + } + Applied applied; + try { + var countSql = CouchbaseCountRewriter.toCountStatement( + parser.parseStatement(request.sql())); + if (countSql == null) { + return QueryAffectedRowsResult.unsupported(CouchbaseQueryEngine.ENGINE_ID); + } + applied = rowSecurityApplier.apply(parser.parseStatement(countSql), + request.rowSecurityPredicates()); + } catch (InvalidSqlException | UnrewritableRowSecurityException ex) { + return QueryAffectedRowsResult.unsupported(CouchbaseQueryEngine.ENGINE_ID); + } + var scope = clusterManager.defaultScope(descriptor); + try { + var result = scope.query(applied.sql(), options(applied, timeout).readonly(true)); + var count = extractCount(result.rowsAs(byte[].class)); + if (count == null) { + return QueryAffectedRowsResult.unsupported(CouchbaseQueryEngine.ENGINE_ID); + } + return QueryAffectedRowsResult.of(CouchbaseQueryEngine.ENGINE_ID, count, + durationSince(start)); + } catch (CouchbaseException ex) { + throw exceptionTranslator.translate(ex, timeout); + } + } + + /** The {@code af_count} value of the single COUNT(*) row, or {@code null} when malformed. */ + private static Long extractCount(List rows) { + if (rows.isEmpty()) { + return null; + } + if (CouchbaseJson.parseRow(rows.get(0)) instanceof Map row + && row.get(CouchbaseCountRewriter.COUNT_ALIAS) instanceof Number count) { + return count.longValue(); + } + return null; + } + SelectExecutionResult sampleTable(SampleTableRequest request, DatasourceConnectionDescriptor descriptor, int maxRows, Duration timeout) { diff --git a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriterTest.java b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriterTest.java new file mode 100644 index 00000000..85f45f5a --- /dev/null +++ b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseCountRewriterTest.java @@ -0,0 +1,94 @@ +package com.bablsoft.accessflow.engine.couchbase; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class CouchbaseCountRewriterTest { + + private final CouchbaseQueryParser parser = new CouchbaseQueryParser(TestMessages.keyEcho()); + + private String rewrite(String sql) { + return CouchbaseCountRewriter.toCountStatement(parser.parseStatement(sql)); + } + + // ---- countable shapes ---------------------------------------------------------------------- + + @Test + void deleteWithWhereBecomesCountOverSamePredicate() { + assertThat(rewrite("DELETE FROM users WHERE team = 'eng'")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users WHERE team = 'eng'"); + } + + @Test + void deleteWithoutWhereCountsTheWholeKeyspace() { + assertThat(rewrite("DELETE FROM users")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users"); + } + + @Test + void updateDropsSetAndReturningClauses() { + assertThat(rewrite( + "UPDATE users SET bonus = 1, level = 2 WHERE team = 'eng' RETURNING META(users).id")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users WHERE team = 'eng'"); + } + + @Test + void updateUnsetIsDroppedToo() { + assertThat(rewrite("UPDATE users UNSET bonus WHERE team = 'eng'")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users WHERE team = 'eng'"); + } + + @Test + void preservesTheTargetAliasSoWhereReferencesStillResolve() { + assertThat(rewrite("UPDATE users AS u SET u.bonus = 1 WHERE u.team = 'eng'")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users AS `u` WHERE u.team = 'eng'"); + assertThat(rewrite("DELETE FROM users u WHERE u.age < 18")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users AS `u` WHERE u.age < 18"); + } + + @Test + void preservesDottedAndBacktickedKeyspaceSourceText() { + assertThat(rewrite("DELETE FROM `Bucket-1`.app.`Users` WHERE x = 1")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM `Bucket-1`.app.`Users` WHERE x = 1"); + } + + @Test + void preservesTheDefaultNamespacePrefix() { + assertThat(rewrite("DELETE FROM default:users WHERE a = 1")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM default:users WHERE a = 1"); + } + + @Test + void deleteWithComplexWhereExpressionIsCarriedVerbatim() { + assertThat(rewrite("DELETE FROM users WHERE (age > 21 AND team = 'eng') OR vip = true")) + .isEqualTo("SELECT COUNT(*) AS af_count FROM users " + + "WHERE (age > 21 AND team = 'eng') OR vip = true"); + } + + // ---- uncountable shapes -------------------------------------------------------------------- + + @Test + void nonMutatingAndNonRewritableKindsAreNull() { + assertThat(rewrite("SELECT * FROM users")).isNull(); + assertThat(rewrite("INSERT INTO users (KEY, VALUE) VALUES ('k', {'a': 1})")).isNull(); + assertThat(rewrite("UPSERT INTO users (KEY, VALUE) VALUES ('k', {'a': 1})")).isNull(); + assertThat(rewrite("CREATE INDEX idx ON users(age)")).isNull(); + assertThat(rewrite("MERGE INTO users AS t USING staged AS s ON t.id = s.id " + + "WHEN MATCHED THEN UPDATE SET t.a = s.a")).isNull(); + } + + @Test + void failClosedShapesAreNull() { + // The same shapes the row-security splicer refuses to rewrite. + assertThat(rewrite("WITH x AS (SELECT 1) DELETE FROM users WHERE a = 1")).isNull(); + assertThat(rewrite("DELETE FROM users WHERE uid IN (SELECT RAW id FROM admins)")).isNull(); + assertThat(rewrite("UPDATE users SET a = (SELECT RAW b FROM c)[0] WHERE x = 1")).isNull(); + assertThat(rewrite("DELETE FROM users USE KEYS ['k1']")).isNull(); + } + + @Test + void limitIsNullBecauseItCapsTheMutationCount() { + assertThat(rewrite("DELETE FROM users WHERE team = 'eng' LIMIT 5")).isNull(); + } +} diff --git a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineIntegrationTest.java b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineIntegrationTest.java index 557c6dcb..e483ffd9 100644 --- a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineIntegrationTest.java +++ b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineIntegrationTest.java @@ -5,6 +5,7 @@ import com.bablsoft.accessflow.core.api.DatasourceConnectionTestException; import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.MaskingStrategy; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngineContext; import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; @@ -351,6 +352,50 @@ void dryRunAppliesRowSecurity() { assertThat(result.appliedRowSecurityPolicyIds()).containsExactly(directive.policyId()); } + @Test + void countAffectedRowsCountsDeleteMatchesWithoutMutating() { + var result = countAffected("DELETE FROM _default WHERE team = 'eng'", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("couchbase"); + assertThat(result.affectedRows()).isEqualTo(2); + assertThat(result.duration()).isNotNull(); + // The count is a read-only probe — all three documents survive. + assertThat(((SelectExecutionResult) run("SELECT * FROM _default")).rowCount()).isEqualTo(3); + } + + @Test + void countAffectedRowsCountsUpdateWithoutWhereAsWholeCollection() { + var result = countAffected("UPDATE _default SET bonus = 1", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(3); + } + + @Test + void countAffectedRowsAppliesRowSecurity() { + var directive = new RowSecurityDirective(UUID.randomUUID(), "_default", "team", + RowSecurityOperator.EQUALS, List.of("eng")); + // Cy (sales, 80) is shielded by the eng-only predicate; only Bo (eng, 90) matches. + var result = countAffected("DELETE FROM _default WHERE salary < 95", List.of(directive)); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(1); + } + + @Test + void countAffectedRowsFailsClosedOnUnsupportedShapes() { + assertThat(countAffected("DELETE FROM _default USE KEYS ['p1']", List.of()).supported()) + .isFalse(); + assertThat(countAffected("SELECT * FROM _default", List.of()).supported()).isFalse(); + } + + private static QueryAffectedRowsResult countAffected(String query, + List rls) { + var request = new QueryExecutionRequest(descriptor.id(), query, + engine.parse(query).type(), null, null, List.of(), List.of(), rls, false, + List.of(query)); + return engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(30))); + } + private static QueryDryRunResult dryRun(String query, List rls) { var type = engine.parse(query).type(); var request = new QueryExecutionRequest(descriptor.id(), query, type, null, null, diff --git a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineTest.java b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineTest.java index 604bfd53..30aa055b 100644 --- a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineTest.java +++ b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryEngineTest.java @@ -1,9 +1,16 @@ package com.bablsoft.accessflow.engine.couchbase; +import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.QueryEngineContext; +import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.SslMode; import org.junit.jupiter.api.Test; import java.time.Clock; +import java.time.Duration; import java.util.Map; import java.util.UUID; @@ -23,11 +30,23 @@ void usingTheEngineBeforeInitializeFails() { assertThatThrownBy(() -> engine.parse("SELECT 1")) .isInstanceOf(IllegalStateException.class) .hasMessageContaining("before initialize"); + assertThatThrownBy(() -> engine.countAffectedRows(dryRunRequest())) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("before initialize"); // Eviction and shutdown before initialize are safe no-ops. engine.evictDatasource(UUID.randomUUID()); engine.shutdown(); } + private static QueryEngineDryRunRequest dryRunRequest() { + var request = new QueryExecutionRequest(UUID.randomUUID(), "DELETE FROM users", + QueryType.DELETE, null, null); + var descriptor = new DatasourceConnectionDescriptor(UUID.randomUUID(), UUID.randomUUID(), + DbType.COUCHBASE, "127.0.0.1", 1, "bucket", "user", "", SslMode.DISABLE, 10, + 1000, true, null, false, null, "couchbase", null, null, null, null, true); + return new QueryEngineDryRunRequest(request, descriptor, Duration.ofSeconds(5)); + } + @Test void initializeWiresTheParser() { var engine = new CouchbaseQueryEngine(); diff --git a/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutorTest.java b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutorTest.java new file mode 100644 index 00000000..5c498cb7 --- /dev/null +++ b/engines/couchbase/src/test/java/com/bablsoft/accessflow/engine/couchbase/CouchbaseQueryExecutorTest.java @@ -0,0 +1,70 @@ +package com.bablsoft.accessflow.engine.couchbase; + +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryType; +import com.couchbase.client.java.query.QueryScanConsistency; +import org.junit.jupiter.api.Test; + +import java.time.Clock; +import java.time.Duration; +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Offline unit coverage of {@link CouchbaseQueryExecutor#countAffectedRows}: every degrade-to- + * unsupported branch returns before any cluster access, so the executor is wired with a + * {@code null} cluster manager — reaching it would NPE and fail the test. The live counting + * paths (real COUNT(*), row security, non-mutation) are covered by + * {@link CouchbaseQueryEngineIntegrationTest}. + */ +class CouchbaseQueryExecutorTest { + + private final CouchbaseQueryExecutor executor = new CouchbaseQueryExecutor( + null, new CouchbaseQueryParser(TestMessages.keyEcho()), + new CouchbaseRowSecurityApplier(TestMessages.keyEcho()), new CouchbaseResultMapper(), + new CouchbaseExceptionTranslator(TestMessages.keyEcho()), + QueryScanConsistency.REQUEST_PLUS, Clock.systemUTC()); + + private QueryAffectedRowsResult count(String sql, QueryType type) { + var request = new QueryExecutionRequest(UUID.randomUUID(), sql, type, null, null, + List.of(), List.of(), List.of(), false, List.of(sql)); + return executor.countAffectedRows(request, null, Duration.ofSeconds(5)); + } + + @Test + void nonMutatingQueryTypesAreUnsupported() { + for (var type : List.of(QueryType.SELECT, QueryType.INSERT, QueryType.DDL)) { + var result = count("SELECT * FROM users", type); + assertThat(result.supported()).as(type.name()).isFalse(); + assertThat(result.engineId()).isEqualTo("couchbase"); + assertThat(result.affectedRows()).isNull(); + } + } + + @Test + void unparseableSqlIsUnsupportedNotThrown() { + assertThat(count("GRANT ALL ON users", QueryType.DELETE).supported()).isFalse(); + assertThat(count("DELETE FROM users WHERE (a = 1", QueryType.DELETE) + .supported()).isFalse(); + } + + @Test + void mergeClassifiedAsUpdateIsUnsupported() { + assertThat(count("MERGE INTO users AS t USING staged AS s ON t.id = s.id " + + "WHEN MATCHED THEN UPDATE SET t.a = s.a", QueryType.UPDATE) + .supported()).isFalse(); + } + + @Test + void failClosedShapesAreUnsupported() { + for (var sql : List.of( + "DELETE FROM users USE KEYS ['k1']", + "DELETE FROM users WHERE uid IN (SELECT RAW id FROM admins)", + "DELETE FROM users WHERE team = 'eng' LIMIT 5")) { + assertThat(count(sql, QueryType.DELETE).supported()).as(sql).isFalse(); + } + } +} diff --git a/engines/elasticsearch/pom.xml b/engines/elasticsearch/pom.xml index 715a1f3a..82c8821a 100644 --- a/engines/elasticsearch/pom.xml +++ b/engines/elasticsearch/pom.xml @@ -13,7 +13,7 @@ this one JAR) pin this version + the SHA-256 of the shaded JAR; bump them together whenever the engine (or a core.api type it compiles against) changes. --> - 1.0.3 + 1.0.4 jar AccessFlow Elasticsearch Query Engine diff --git a/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngine.java b/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngine.java index ca67a5ac..0b551010 100644 --- a/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngine.java +++ b/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngine.java @@ -3,6 +3,7 @@ import com.bablsoft.accessflow.core.api.ConnectionTestResult; import com.bablsoft.accessflow.core.api.DatabaseSchemaView; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngine; import com.bablsoft.accessflow.core.api.QueryEngineContext; @@ -93,6 +94,13 @@ public QueryDryRunResult dryRun(QueryEngineDryRunRequest request) { request.effectiveTimeout()); } + @Override + public QueryAffectedRowsResult countAffectedRows(QueryEngineDryRunRequest request) { + // Shared with OpenSearchQueryEngine: engineId() stamps the right provider on the result. + return initialized(executor).countAffectedRows(engineId(), request.request(), + request.descriptor(), request.effectiveTimeout()); + } + @Override public ConnectionTestResult testConnection(DatasourceConnectionDescriptor descriptor) { return initialized(connectionProbe).test(descriptor); diff --git a/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutor.java b/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutor.java index 0d749f60..886cfcfa 100644 --- a/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutor.java +++ b/engines/elasticsearch/src/main/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutor.java @@ -1,12 +1,15 @@ package com.bablsoft.accessflow.engine.elasticsearch; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.InvalidSqlException; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; import com.bablsoft.accessflow.core.api.QueryType; import com.bablsoft.accessflow.core.api.SampleTableRequest; import com.bablsoft.accessflow.core.api.SelectExecutionResult; +import com.bablsoft.accessflow.core.api.UnrewritableRowSecurityException; import com.bablsoft.accessflow.core.api.UpdateExecutionResult; import tools.jackson.databind.JsonNode; import tools.jackson.databind.node.ObjectNode; @@ -120,6 +123,38 @@ QueryDryRunResult dryRun(String engineId, QueryExecutionRequest request, } } + QueryAffectedRowsResult countAffectedRows(String engineId, QueryExecutionRequest request, + DatasourceConnectionDescriptor descriptor, + Duration timeout) { + var start = clock.instant(); + EsCommand cmd; + try { + var command = parser.parseCommand(request.sql()); + // Only the by-query mutations have a countable pre-image; everything else degrades. + if (command.operation() != EsOperation.UPDATE_BY_QUERY + && command.operation() != EsOperation.DELETE_BY_QUERY) { + return QueryAffectedRowsResult.unsupported(engineId); + } + cmd = rowSecurityApplier.apply(command, request.rowSecurityPredicates()).command(); + } catch (InvalidSqlException | UnrewritableRowSecurityException ex) { + // A preflight count never throws for shape reasons — fail closed to "unsupported". + return QueryAffectedRowsResult.unsupported(engineId, ex.getMessage()); + } + var transport = clientManager.transport(descriptor); + try { + var body = EsJson.object(); + body.set("query", cmd.query() != null ? cmd.query() : EsJson.matchAll()); + // Non-mutating _count of the governed (bool.filter-wrapped) query — the same endpoint + // the COUNT read uses; _count does not accept a ?timeout= query parameter. + var response = EsJson.parse(transport.perform("POST", path(cmd.index(), "_count"), + Map.of(), EsJson.write(body), JSON)); + return QueryAffectedRowsResult.of(engineId, longField(response, "count"), + durationSince(start)); + } catch (SearchTransportException ex) { + throw exceptionTranslator.translate(ex, timeout); + } + } + private QueryExecutionResult search(SearchTransport transport, EsCommand cmd, QueryExecutionRequest request, int maxRows, Duration timeout, Instant start, Set policyIds) { diff --git a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineIntegrationTest.java b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineIntegrationTest.java index c8653dba..3dbe9a23 100644 --- a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineIntegrationTest.java +++ b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineIntegrationTest.java @@ -6,6 +6,7 @@ import com.bablsoft.accessflow.core.api.DatasourceConnectionTestException; import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.MaskingStrategy; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngineContext; import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; @@ -116,6 +117,47 @@ void dryRunIndexOperationIsUnsupported() { assertThat(result.supported()).isFalse(); } + private static QueryAffectedRowsResult countAffected(String sql, + List rls) { + var type = engine.parse(sql).type(); + var request = new QueryExecutionRequest(DS_ID, sql, type, null, null, List.of(), List.of(), + rls, false, null); + return engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(30))); + } + + @Test + void countAffectedRowsCountsDeleteByQueryWithoutMutating() { + seed("logs-preflight"); + var result = countAffected( + "{\"delete_by_query\":\"logs-preflight\",\"query\":{\"term\":{\"tenant\":\"acme\"}}}", + List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("elasticsearch"); + assertThat(result.affectedRows()).isEqualTo(2L); + // Non-mutating preflight: all three seeded documents survive. + var count = (SelectExecutionResult) exec("{\"count\":\"logs-preflight\"}", QueryType.SELECT); + assertThat(count.rows().get(0).get(0)).isEqualTo(3L); + } + + @Test + void countAffectedRowsAppliesRowSecurityToTheGovernedCount() { + seed("logs-preflight"); + var directive = new RowSecurityDirective(UUID.randomUUID(), "logs-preflight", "tenant", + RowSecurityOperator.EQUALS, List.of("acme")); + var result = countAffected("{\"update_by_query\":\"logs-preflight\"}", List.of(directive)); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(2L); + } + + @Test + void countAffectedRowsIsUnsupportedForReads() { + seed("logs-preflight"); + var result = countAffected("{\"search\":\"logs-preflight\"}", List.of()); + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("elasticsearch"); + } + @Test void dryRunAppliesRowSecurity() { seed("logs-dryrun"); diff --git a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineTest.java b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineTest.java index c4e105d6..458cbe89 100644 --- a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineTest.java +++ b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/ElasticsearchQueryEngineTest.java @@ -1,10 +1,18 @@ package com.bablsoft.accessflow.engine.elasticsearch; +import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.QueryEngineContext; +import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.SslMode; import org.junit.jupiter.api.Test; import java.time.Clock; +import java.time.Duration; import java.time.ZoneOffset; +import java.util.List; import java.util.Map; import java.util.UUID; @@ -38,4 +46,19 @@ void parsesAfterInitialize() { Clock.systemDefaultZone().withZone(ZoneOffset.UTC))); assertThat(engine.parse("{\"search\":\"logs\"}").referencedTables()).containsExactly("logs"); } + + @Test + void countAffectedRowsForAReadEnvelopeIsUnsupportedWithTheElasticsearchEngineId() { + engine.initialize(new QueryEngineContext(TestMessages.keyEcho(), c -> c, Map.of(), + Clock.systemDefaultZone().withZone(ZoneOffset.UTC))); + var descriptor = new DatasourceConnectionDescriptor(UUID.randomUUID(), UUID.randomUUID(), + DbType.ELASTICSEARCH, "localhost", 9200, null, "", "", SslMode.DISABLE, 10, 10000, + false, null, false, null, "elasticsearch", null, null, null, null, true, null, null); + var request = new QueryExecutionRequest(UUID.randomUUID(), "{\"search\":\"logs\"}", + QueryType.SELECT, null, null, List.of(), List.of(), List.of(), false, null); + var result = engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(30))); + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("elasticsearch"); + } } diff --git a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutorTest.java b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutorTest.java new file mode 100644 index 00000000..7dc7465e --- /dev/null +++ b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/EsQueryExecutorTest.java @@ -0,0 +1,194 @@ +package com.bablsoft.accessflow.engine.elasticsearch; + +import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.DbType; +import com.bablsoft.accessflow.core.api.QueryExecutionFailedException; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryExecutionTimeoutException; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.RowSecurityDirective; +import com.bablsoft.accessflow.core.api.RowSecurityOperator; +import com.bablsoft.accessflow.core.api.SslMode; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.time.Clock; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +/** + * Unit tests for the executor's governed affected-row preflight (AF-624): the by-query mutations + * are counted via a non-mutating {@code _count} of the row-security-wrapped query; every other + * operation and every shape/parse failure degrades to {@code unsupported} without throwing. + * Transport-level tests mock the {@link SearchClientManager} / {@link SearchTransport} seam, the + * convention of the engine's other Spring-free unit tests. + */ +class EsQueryExecutorTest { + + private static final UUID DS_ID = UUID.randomUUID(); + private static final Duration TIMEOUT = Duration.ofSeconds(30); + + private final SearchClientManager clientManager = mock(SearchClientManager.class); + private final SearchTransport transport = mock(SearchTransport.class); + private final EsQueryExecutor executor = new EsQueryExecutor(clientManager, + new EsQueryParser(TestMessages.keyEcho()), + new EsRowSecurityApplier(TestMessages.keyEcho()), new EsResultMapper(), + new EsExceptionTranslator(TestMessages.keyEcho()), Clock.systemUTC()); + private final DatasourceConnectionDescriptor descriptor = new DatasourceConnectionDescriptor( + DS_ID, UUID.randomUUID(), DbType.ELASTICSEARCH, "localhost", 9200, null, "", "", + SslMode.DISABLE, 10, 10000, false, null, false, null, "elasticsearch", null, null, + null, null, true, null, null); + + private static QueryExecutionRequest request(String sql, QueryType type, + List rls) { + return new QueryExecutionRequest(DS_ID, sql, type, null, null, List.of(), List.of(), rls, + false, null); + } + + private void givenTransportReturns(String response) { + when(clientManager.transport(descriptor)).thenReturn(transport); + when(transport.perform(anyString(), anyString(), anyMap(), any(), anyString())) + .thenReturn(response); + } + + @Test + void countsDeleteByQueryViaNonMutatingCount() { + givenTransportReturns("{\"count\":2}"); + var result = executor.countAffectedRows("elasticsearch", + request("{\"delete_by_query\":\"logs\",\"query\":{\"term\":{\"tenant\":\"acme\"}}}", + QueryType.DELETE, List.of()), + descriptor, TIMEOUT); + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("elasticsearch"); + assertThat(result.affectedRows()).isEqualTo(2L); + assertThat(result.duration()).isNotNull(); + var body = ArgumentCaptor.forClass(String.class); + verify(transport).perform(eq("POST"), eq("/logs/_count"), eq(Map.of()), body.capture(), + eq("application/json")); + assertThat(EsJson.parse(body.getValue()).path("query").path("term").path("tenant").asString()) + .isEqualTo("acme"); + } + + @Test + void countsUpdateByQueryDefaultingToMatchAllWhenQueryAbsent() { + givenTransportReturns("{\"count\":3}"); + var result = executor.countAffectedRows("elasticsearch", + request("{\"update_by_query\":\"logs\"}", QueryType.UPDATE, List.of()), + descriptor, TIMEOUT); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(3L); + var body = ArgumentCaptor.forClass(String.class); + verify(transport).perform(eq("POST"), eq("/logs/_count"), eq(Map.of()), body.capture(), + eq("application/json")); + assertThat(EsJson.parse(body.getValue()).path("query").has("match_all")).isTrue(); + } + + @Test + void wrapsRowSecurityDirectivesInBoolFilterAroundTheUserQuery() { + givenTransportReturns("{\"count\":1}"); + var directive = new RowSecurityDirective(UUID.randomUUID(), "logs", "tenant", + RowSecurityOperator.EQUALS, List.of("acme")); + var result = executor.countAffectedRows("elasticsearch", + request("{\"delete_by_query\":\"logs\",\"query\":{\"term\":{\"level\":\"debug\"}}}", + QueryType.DELETE, List.of(directive)), + descriptor, TIMEOUT); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(1L); + var body = ArgumentCaptor.forClass(String.class); + verify(transport).perform(eq("POST"), eq("/logs/_count"), eq(Map.of()), body.capture(), + eq("application/json")); + var query = EsJson.parse(body.getValue()).path("query"); + assertThat(query.path("bool").path("must").get(0).path("term").path("level").asString()) + .isEqualTo("debug"); + assertThat(query.path("bool").path("filter").get(0).path("term").path("tenant").asString()) + .isEqualTo("acme"); + } + + @Test + void emptyDirectiveValuesFailClosedToAMatchNothingCount() { + givenTransportReturns("{\"count\":0}"); + var directive = new RowSecurityDirective(UUID.randomUUID(), "logs", "tenant", + RowSecurityOperator.IN, List.of()); + var result = executor.countAffectedRows("elasticsearch", + request("{\"delete_by_query\":\"logs\"}", QueryType.DELETE, List.of(directive)), + descriptor, TIMEOUT); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isZero(); + var body = ArgumentCaptor.forClass(String.class); + verify(transport).perform(eq("POST"), eq("/logs/_count"), eq(Map.of()), body.capture(), + eq("application/json")); + assertThat(body.getValue()).contains("must_not"); + } + + @Test + void nonMutatingOperationsAreUnsupportedWithoutTouchingTheTransport() { + for (var sql : List.of( + "{\"search\":\"logs\"}", + "{\"count\":\"logs\"}", + "{\"index\":\"logs\",\"document\":{\"a\":1}}", + "{\"bulk\":\"logs\",\"operations\":[{\"document\":{\"a\":1}}]}", + "{\"create_index\":\"logs\"}", + "{\"delete_index\":\"logs\"}")) { + var result = executor.countAffectedRows("elasticsearch", + request(sql, QueryType.SELECT, List.of()), descriptor, TIMEOUT); + assertThat(result.supported()).as(sql).isFalse(); + assertThat(result.engineId()).isEqualTo("elasticsearch"); + } + verifyNoInteractions(clientManager); + } + + @Test + void parseFailureIsUnsupportedNotThrown() { + var result = executor.countAffectedRows("elasticsearch", + request("not json at all", QueryType.DELETE, List.of()), descriptor, TIMEOUT); + assertThat(result.supported()).isFalse(); + assertThat(result.unsupportedReason()).isNotBlank(); + verifyNoInteractions(clientManager); + } + + @Test + void stampsTheOpenSearchEngineIdWhenAsked() { + givenTransportReturns("{\"count\":5}"); + var result = executor.countAffectedRows("opensearch", + request("{\"delete_by_query\":\"logs\"}", QueryType.DELETE, List.of()), + descriptor, TIMEOUT); + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("opensearch"); + assertThat(result.affectedRows()).isEqualTo(5L); + } + + @Test + void transportFailurePropagatesAsTranslatedExecutionException() { + when(clientManager.transport(descriptor)).thenReturn(transport); + when(transport.perform(anyString(), anyString(), anyMap(), any(), anyString())) + .thenThrow(new SearchTransportException(500, "boom", false, null)); + assertThatThrownBy(() -> executor.countAffectedRows("elasticsearch", + request("{\"delete_by_query\":\"logs\"}", QueryType.DELETE, List.of()), + descriptor, TIMEOUT)) + .isInstanceOf(QueryExecutionFailedException.class); + } + + @Test + void transportTimeoutPropagatesAsTimeoutException() { + when(clientManager.transport(descriptor)).thenReturn(transport); + when(transport.perform(anyString(), anyString(), anyMap(), any(), anyString())) + .thenThrow(new SearchTransportException(0, null, true, null)); + assertThatThrownBy(() -> executor.countAffectedRows("elasticsearch", + request("{\"delete_by_query\":\"logs\"}", QueryType.DELETE, List.of()), + descriptor, TIMEOUT)) + .isInstanceOf(QueryExecutionTimeoutException.class); + } +} diff --git a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineIntegrationTest.java b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineIntegrationTest.java index 1650b8b3..77d43342 100644 --- a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineIntegrationTest.java +++ b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineIntegrationTest.java @@ -5,6 +5,7 @@ import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.MaskingStrategy; import com.bablsoft.accessflow.core.api.QueryEngineContext; +import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; import com.bablsoft.accessflow.core.api.QueryEngineExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; @@ -107,8 +108,17 @@ void drivesSearchRowSecurityMaskingAndIntrospectionEndToEnd() { var user = (Map) filtered.rows().get(0).get(userCol); assertThat((String) user.get("email")).contains("***@"); - var delete = (UpdateExecutionResult) exec( - "{\"delete_by_query\":\"events\",\"query\":{\"term\":{\"tenant\":\"globex\"}}}", + // AF-624 preflight: the governed _count of the delete's pre-image, without mutating. + var deleteSql = "{\"delete_by_query\":\"events\",\"query\":{\"term\":{\"tenant\":\"globex\"}}}"; + var preflight = engine.countAffectedRows(new QueryEngineDryRunRequest( + new QueryExecutionRequest(DS_ID, deleteSql, QueryType.DELETE, null, null, List.of(), + List.of(), List.of(), false, null), + descriptor, Duration.ofSeconds(30))); + assertThat(preflight.supported()).isTrue(); + assertThat(preflight.engineId()).isEqualTo("opensearch"); + assertThat(preflight.affectedRows()).isEqualTo(1L); + + var delete = (UpdateExecutionResult) exec(deleteSql, QueryType.DELETE, List.of(), List.of()); assertThat(delete.rowsAffected()).isEqualTo(1); diff --git a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineTest.java b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineTest.java index 38c62c34..7f6ccfc3 100644 --- a/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineTest.java +++ b/engines/elasticsearch/src/test/java/com/bablsoft/accessflow/engine/elasticsearch/OpenSearchQueryEngineTest.java @@ -1,11 +1,20 @@ package com.bablsoft.accessflow.engine.elasticsearch; +import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.QueryEngineContext; +import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; +import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.SslMode; import org.junit.jupiter.api.Test; import java.time.Clock; +import java.time.Duration; import java.time.ZoneOffset; +import java.util.List; import java.util.Map; +import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; @@ -25,4 +34,19 @@ void sharesTheElasticsearchParseLogic() { assertThat(engine.parse("{\"count\":\"events\"}").referencedTables()) .containsExactly("events"); } + + @Test + void countAffectedRowsInheritsTheSharedLogicButStampsTheOpenSearchEngineId() { + engine.initialize(new QueryEngineContext(TestMessages.keyEcho(), c -> c, Map.of(), + Clock.systemDefaultZone().withZone(ZoneOffset.UTC))); + var descriptor = new DatasourceConnectionDescriptor(UUID.randomUUID(), UUID.randomUUID(), + DbType.OPENSEARCH, "localhost", 9200, null, "", "", SslMode.DISABLE, 10, 10000, + false, null, false, null, "opensearch", null, null, null, null, true, null, null); + var request = new QueryExecutionRequest(UUID.randomUUID(), "{\"search\":\"events\"}", + QueryType.SELECT, null, null, List.of(), List.of(), List.of(), false, null); + var result = engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(30))); + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("opensearch"); + } } diff --git a/engines/mongodb/dependency-reduced-pom.xml b/engines/mongodb/dependency-reduced-pom.xml index c2cbfcea..98589675 100644 --- a/engines/mongodb/dependency-reduced-pom.xml +++ b/engines/mongodb/dependency-reduced-pom.xml @@ -4,7 +4,7 @@ com.bablsoft.accessflow accessflow-engine-mongodb AccessFlow MongoDB Query Engine - 1.0.6 + 1.0.7 On-demand MongoDB engine plugin for AccessFlow (issue AF-414). Implements the core.api.QueryEngine SPI and ships as a self-contained shaded JAR (bundling mongodb-driver-sync and a relocated Jackson) resolved through the connector catalog and diff --git a/engines/mongodb/pom.xml b/engines/mongodb/pom.xml index c17ce1cf..a9e78ae1 100644 --- a/engines/mongodb/pom.xml +++ b/engines/mongodb/pom.xml @@ -12,7 +12,7 @@ (connectors/mongodb/connector.json) pins this version + the SHA-256 of the shaded JAR; bump BOTH together whenever the engine (or a core.api type it compiles against) changes. --> - 1.0.6 + 1.0.7 jar AccessFlow MongoDB Query Engine diff --git a/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngine.java b/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngine.java index 8824cf1f..8c403800 100644 --- a/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngine.java +++ b/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngine.java @@ -3,6 +3,7 @@ import com.bablsoft.accessflow.core.api.ConnectionTestResult; import com.bablsoft.accessflow.core.api.DatabaseSchemaView; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngine; import com.bablsoft.accessflow.core.api.QueryEngineContext; @@ -85,6 +86,12 @@ public QueryDryRunResult dryRun(QueryEngineDryRunRequest request) { request.effectiveTimeout()); } + @Override + public QueryAffectedRowsResult countAffectedRows(QueryEngineDryRunRequest request) { + return initialized(executor).countAffectedRows(request.request(), request.descriptor(), + request.effectiveTimeout()); + } + @Override public ConnectionTestResult testConnection(DatasourceConnectionDescriptor descriptor) { return initialized(connectionProbe).test(descriptor); diff --git a/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryExecutor.java b/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryExecutor.java index 69db46a1..4e7cba03 100644 --- a/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryExecutor.java +++ b/engines/mongodb/src/main/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryExecutor.java @@ -1,17 +1,22 @@ package com.bablsoft.accessflow.engine.mongodb; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.InvalidSqlException; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; +import com.bablsoft.accessflow.core.api.QueryType; import com.bablsoft.accessflow.core.api.SampleTableRequest; import com.bablsoft.accessflow.core.api.SelectExecutionResult; +import com.bablsoft.accessflow.core.api.UnrewritableRowSecurityException; import com.bablsoft.accessflow.core.api.UpdateExecutionResult; import com.mongodb.MongoException; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.MongoIterable; +import com.mongodb.client.model.CountOptions; import com.mongodb.client.model.IndexOptions; import org.bson.Document; @@ -121,6 +126,45 @@ QueryDryRunResult dryRun(QueryExecutionRequest request, } } + QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request, + DatasourceConnectionDescriptor descriptor, + Duration timeout) { + if (request.queryType() != QueryType.UPDATE && request.queryType() != QueryType.DELETE) { + return QueryAffectedRowsResult.unsupported(MongoQueryEngine.ENGINE_ID); + } + var start = clock.instant(); + MongoCommand command; + try { + var applied = rowSecurityApplier.apply(parser.parseCommand(request.sql()), + request.rowSecurityPredicates()); + command = applied.command(); + } catch (InvalidSqlException | UnrewritableRowSecurityException ex) { + // Fail closed: a shape we cannot parse or securely filter is never counted unfiltered. + return QueryAffectedRowsResult.unsupported(MongoQueryEngine.ENGINE_ID, ex.getMessage()); + } + var operation = command.operation(); + if (operation.queryType() != QueryType.UPDATE && operation.queryType() != QueryType.DELETE) { + return QueryAffectedRowsResult.unsupported(MongoQueryEngine.ENGINE_ID); + } + var database = clientManager.database(descriptor, true); + try { + var collection = database.getCollection(command.collection()); + var options = new CountOptions().maxTime(Math.max(1, timeout.toMillis()), + TimeUnit.MILLISECONDS); + long matched = collection.countDocuments(filterOrEmpty(command.filter()), options); + // *_ONE (and findOneAndUpdate/replaceOne) statements touch at most one document. + long affected = switch (operation) { + case UPDATE_ONE, REPLACE_ONE, FIND_ONE_AND_UPDATE, DELETE_ONE -> + Math.min(matched, 1); + default -> matched; + }; + return QueryAffectedRowsResult.of(MongoQueryEngine.ENGINE_ID, affected, + durationSince(start)); + } catch (MongoException ex) { + throw exceptionTranslator.translate(ex, timeout); + } + } + /** The explain sub-command document for an explainable operation, or {@code null} otherwise. */ private static Document explainTarget(MongoCommand command) { var coll = command.collection(); diff --git a/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineIntegrationTest.java b/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineIntegrationTest.java index 351f40b1..6dd3d796 100644 --- a/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineIntegrationTest.java +++ b/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineIntegrationTest.java @@ -6,6 +6,7 @@ import com.bablsoft.accessflow.core.api.DatasourceConnectionTestException; import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.MaskingStrategy; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngineContext; import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; @@ -305,6 +306,58 @@ void dryRunAppliesRowSecurity() { assertThat(result.appliedRowSecurityPolicyIds()).containsExactly(directive.policyId()); } + @Test + void countAffectedRowsForDeleteManyCountsWithoutMutating() { + var result = countAffectedRows("db.people.deleteMany({ team: 'eng' })", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.engineId()).isEqualTo("mongodb"); + assertThat(result.affectedRows()).isEqualTo(2L); + // No mutation occurred: all three documents survive. + assertThat(((SelectExecutionResult) run("db.people.find({})")).rowCount()).isEqualTo(3); + } + + @Test + void countAffectedRowsForUpdateManyCountsMatchingDocuments() { + var result = countAffectedRows( + "db.people.updateMany({ salary: { $gte: 90 } }, { $set: { bonus: 1 } })", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(2L); + assertThat(((SelectExecutionResult) run("db.people.find({ bonus: 1 })")).rowCount()) + .isZero(); + } + + @Test + void countAffectedRowsCapsSingleDocumentOperationsAtOne() { + var result = countAffectedRows("db.people.deleteOne({ team: 'eng' })", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(1L); + } + + @Test + void countAffectedRowsAppliesRowSecurity() { + var directive = new RowSecurityDirective(UUID.randomUUID(), "people", "team", + RowSecurityOperator.EQUALS, List.of("sales")); + var result = countAffectedRows("db.people.deleteMany({})", List.of(directive)); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(1L); + } + + @Test + void countAffectedRowsIsUnsupportedForReads() { + var result = countAffectedRows("db.people.find({})", List.of()); + assertThat(result.supported()).isFalse(); + assertThat(result.affectedRows()).isNull(); + } + + private QueryAffectedRowsResult countAffectedRows(String query, + List rls) { + var type = engine.parse(query).type(); + var request = new QueryExecutionRequest(descriptor.id(), query, type, null, null, + List.of(), List.of(), rls, false, List.of(query)); + return engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(10))); + } + private QueryDryRunResult dryRun(String query, List rls) { var type = engine.parse(query).type(); var request = new QueryExecutionRequest(descriptor.id(), query, type, null, null, diff --git a/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineTest.java b/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineTest.java index 3d22ffc4..1d180acd 100644 --- a/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineTest.java +++ b/engines/mongodb/src/test/java/com/bablsoft/accessflow/engine/mongodb/MongoQueryEngineTest.java @@ -1,10 +1,17 @@ package com.bablsoft.accessflow.engine.mongodb; +import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.QueryEngineContext; +import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; +import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryType; +import com.bablsoft.accessflow.core.api.SslMode; import org.junit.jupiter.api.Test; import java.time.Clock; +import java.time.Duration; +import java.util.List; import java.util.Map; import java.util.UUID; @@ -20,6 +27,18 @@ private static MongoQueryEngine initialized() { return engine; } + private static DatasourceConnectionDescriptor descriptor() { + return new DatasourceConnectionDescriptor(UUID.randomUUID(), UUID.randomUUID(), + DbType.MONGODB, "127.0.0.1", 27017, "db", null, "", SslMode.DISABLE, 10, 1000, + true, null, false, null, null, null, null, null, null, true); + } + + private static QueryEngineDryRunRequest dryRunRequest(String query, QueryType type) { + var request = new QueryExecutionRequest(UUID.randomUUID(), query, type, null, null, + List.of(), List.of(), List.of(), false, List.of(query)); + return new QueryEngineDryRunRequest(request, descriptor(), Duration.ofSeconds(5)); + } + @Test void engineIdIsTheConnectorId() { assertThat(new MongoQueryEngine().engineId()).isEqualTo("mongodb"); @@ -40,6 +59,48 @@ void methodsThrowBeforeInitialize() { .hasMessageContaining("initialize"); } + @Test + void countAffectedRowsIsUnsupportedForNonWriteQueryTypes() { + var result = initialized().countAffectedRows( + dryRunRequest("db.users.find({})", QueryType.SELECT)); + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("mongodb"); + assertThat(result.affectedRows()).isNull(); + } + + @Test + void countAffectedRowsIsUnsupportedForInsertQueryTypes() { + var result = initialized().countAffectedRows( + dryRunRequest("db.users.insertOne({ a: 1 })", QueryType.INSERT)); + assertThat(result.supported()).isFalse(); + } + + @Test + void countAffectedRowsIsUnsupportedForUnparseableQueries() { + var result = initialized().countAffectedRows( + dryRunRequest("this is not a mongo query", QueryType.DELETE)); + assertThat(result.supported()).isFalse(); + assertThat(result.engineId()).isEqualTo("mongodb"); + assertThat(result.unsupportedReason()).isNotBlank(); + } + + @Test + void countAffectedRowsIsUnsupportedWhenParsedOperationIsNotUpdateOrDelete() { + // Declared DELETE but the statement parses as a read — fail closed, no count. + var result = initialized().countAffectedRows( + dryRunRequest("db.users.find({})", QueryType.DELETE)); + assertThat(result.supported()).isFalse(); + } + + @Test + void countAffectedRowsThrowsBeforeInitialize() { + var engine = new MongoQueryEngine(); + assertThatThrownBy(() -> engine.countAffectedRows( + dryRunRequest("db.users.deleteMany({})", QueryType.DELETE))) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("initialize"); + } + @Test void evictAndShutdownAreSafeBeforeInitialize() { var engine = new MongoQueryEngine(); diff --git a/engines/neo4j/pom.xml b/engines/neo4j/pom.xml index e472761b..86dc8d19 100644 --- a/engines/neo4j/pom.xml +++ b/engines/neo4j/pom.xml @@ -12,7 +12,7 @@ (connectors/neo4j/connector.json) pins this version + the SHA-256 of the shaded JAR; bump them together whenever the engine (or a core.api type it compiles against) changes. --> - 1.0.2 + 1.0.3 jar AccessFlow Neo4j Query Engine diff --git a/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriver.java b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriver.java new file mode 100644 index 00000000..c52bd18b --- /dev/null +++ b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriver.java @@ -0,0 +1,72 @@ +package com.bablsoft.accessflow.engine.neo4j; + +import com.bablsoft.accessflow.engine.neo4j.CypherTokenizer.Kind; + +import java.util.Set; + +/** + * Derives the non-mutating affected-row count query for an UPDATE / DELETE Cypher statement + * (issue AF-624): the statement's read prefix — its {@code MATCH} / {@code OPTIONAL MATCH} / + * {@code WHERE} clauses before the first write clause — with {@code RETURN count(*)} appended in + * place of the write clauses (and any trailing {@code RETURN}). Only the provably-simple shape is + * derived: {@code MATCH} / {@code OPTIONAL MATCH} / {@code WHERE} clauses, then {@code SET} / + * {@code DELETE} / {@code DETACH DELETE} / {@code REMOVE} segments, then an optional trailing + * {@code RETURN} (with {@code ORDER BY} / {@code SKIP} / {@code LIMIT}). Anything else — {@code MERGE}, {@code CREATE}-then-{@code SET}, {@code FOREACH}, + * {@code CALL} subqueries, {@code WITH} / {@code UNWIND} pipelines, reads resuming after the write + * — yields {@code null} so the caller degrades to an unsupported count rather than risk a wrong + * one. {@code count(*)} of the matched rows approximates the affected entities (the same rows the + * write clauses operate on), which is the accepted AF-624 semantics. + */ +final class CypherAffectedCountDeriver { + + /** Depth-0 keywords that begin / end a clause (mirrors the parser's clause boundaries). */ + private static final Set CLAUSE_KEYWORDS = Set.of( + "MATCH", "OPTIONAL", "WHERE", "RETURN", "WITH", "CREATE", "MERGE", "SET", "DELETE", + "DETACH", "REMOVE", "UNWIND", "CALL", "FOREACH", "ORDER", "SKIP", "LIMIT", "UNION", + "USE", "ON", "YIELD"); + /** Clause keywords permitted in the read prefix ahead of the first write clause. */ + private static final Set PREFIX_READ = Set.of("MATCH", "OPTIONAL", "WHERE"); + /** Clause keywords that start the write region of an UPDATE / DELETE statement. */ + private static final Set WRITE_STARTERS = Set.of("SET", "DELETE", "DETACH", "REMOVE"); + /** Clause keywords permitted after the write region (further writes + a trailing RETURN). */ + private static final Set AFTER_WRITE = Set.of( + "SET", "DELETE", "DETACH", "REMOVE", "RETURN", "ORDER", "SKIP", "LIMIT"); + + private CypherAffectedCountDeriver() { + } + + /** + * The count query for {@code statement}, or {@code null} when the statement is not an + * UPDATE / DELETE or its shape makes the read-prefix extraction ambiguous. + */ + static String derive(CypherStatement statement) { + if (statement.kind() != CypherStatementKind.UPDATE + && statement.kind() != CypherStatementKind.DELETE) { + return null; + } + int writeStart = -1; + boolean sawMatch = false; + for (var token : statement.tokens()) { + if (token.depth() != 0 || token.kind() != Kind.WORD + || !CLAUSE_KEYWORDS.contains(token.value())) { + continue; + } + var value = token.value(); + if (writeStart < 0) { + if (WRITE_STARTERS.contains(value)) { + writeStart = token.start(); + } else if (value.equals("MATCH")) { + sawMatch = true; + } else if (!PREFIX_READ.contains(value)) { + return null; // WITH / UNWIND / CALL / CREATE / MERGE / FOREACH / … before the write + } + } else if (!AFTER_WRITE.contains(value)) { + return null; // a read pipeline resumes after the write — ambiguous + } + } + if (writeStart < 0 || !sawMatch) { + return null; // nothing to count: no write clause, or no MATCH selecting the rows + } + return statement.cypher().substring(0, writeStart).strip() + " RETURN count(*) AS affected"; + } +} diff --git a/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngine.java b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngine.java index 2b1e718a..1d53d04c 100644 --- a/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngine.java +++ b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngine.java @@ -3,6 +3,7 @@ import com.bablsoft.accessflow.core.api.ConnectionTestResult; import com.bablsoft.accessflow.core.api.DatabaseSchemaView; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngine; import com.bablsoft.accessflow.core.api.QueryEngineContext; @@ -83,6 +84,12 @@ public QueryDryRunResult dryRun(QueryEngineDryRunRequest request) { request.effectiveTimeout()); } + @Override + public QueryAffectedRowsResult countAffectedRows(QueryEngineDryRunRequest request) { + return initialized(executor).countAffectedRows(request.request(), request.descriptor(), + request.effectiveTimeout()); + } + @Override public ConnectionTestResult testConnection(DatasourceConnectionDescriptor descriptor) { return initialized(connectionProbe).test(descriptor); diff --git a/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryExecutor.java b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryExecutor.java index 725840ac..2a1f0800 100644 --- a/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryExecutor.java +++ b/engines/neo4j/src/main/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryExecutor.java @@ -1,12 +1,14 @@ package com.bablsoft.accessflow.engine.neo4j; import com.bablsoft.accessflow.core.api.DatasourceConnectionDescriptor; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryExecutionRequest; import com.bablsoft.accessflow.core.api.QueryExecutionResult; import com.bablsoft.accessflow.core.api.QueryType; import com.bablsoft.accessflow.core.api.SampleTableRequest; import com.bablsoft.accessflow.core.api.SelectExecutionResult; +import com.bablsoft.accessflow.core.api.UnrewritableRowSecurityException; import com.bablsoft.accessflow.core.api.UpdateExecutionResult; import org.neo4j.driver.Query; import org.neo4j.driver.Result; @@ -108,6 +110,42 @@ QueryDryRunResult dryRun(QueryExecutionRequest request, } } + /** + * Governed affected-row count for an UPDATE / DELETE (issue AF-624): derives the statement's + * non-mutating {@code MATCH … RETURN count(*)} read prefix via {@link CypherAffectedCountDeriver}, + * applies the request's row-security directives to it exactly as {@link #execute} does (any shape + * the applier rejects degrades to unsupported), and runs it in a read transaction bounded by the + * host-computed timeout. Genuine server / connection errors translate like every other path. + */ + QueryAffectedRowsResult countAffectedRows(QueryExecutionRequest request, + DatasourceConnectionDescriptor descriptor, + Duration timeout) { + var start = clock.instant(); + var statement = parser.parseStatement(request.sql()); + var countCypher = CypherAffectedCountDeriver.derive(statement); + if (countCypher == null) { + return QueryAffectedRowsResult.unsupported(Neo4jQueryEngine.ENGINE_ID); + } + Neo4jRowSecurityApplier.Applied applied; + try { + applied = rowSecurityApplier.apply(parser.parseStatement(countCypher), + request.rowSecurityPredicates()); + } catch (UnrewritableRowSecurityException ex) { + // Fail closed: a shape the execute path would reject must never yield a count. + return QueryAffectedRowsResult.unsupported(Neo4jQueryEngine.ENGINE_ID); + } + var driver = driverManager.driver(descriptor); + var txConfig = TransactionConfig.builder().withTimeout(timeout).build(); + var query = new Query(applied.cypher(), applied.parameters()); + try (Session session = driver.session(Neo4jConnectionProbe.sessionConfig(descriptor))) { + long count = session.executeRead(tx -> tx.run(query).single().get(0).asLong(), txConfig); + return QueryAffectedRowsResult.of(Neo4jQueryEngine.ENGINE_ID, count, + durationSince(start)); + } catch (Neo4jException ex) { + throw exceptionTranslator.translate(ex, timeout); + } + } + SelectExecutionResult sampleTable(SampleTableRequest request, DatasourceConnectionDescriptor descriptor, int maxRows, Duration timeout) { diff --git a/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriverTest.java b/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriverTest.java new file mode 100644 index 00000000..d98613ff --- /dev/null +++ b/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/CypherAffectedCountDeriverTest.java @@ -0,0 +1,89 @@ +package com.bablsoft.accessflow.engine.neo4j; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class CypherAffectedCountDeriverTest { + + private final CypherQueryParser parser = new CypherQueryParser(TestMessages.keyEcho()); + + private String derive(String cypher) { + return CypherAffectedCountDeriver.derive(parser.parseStatement(cypher)); + } + + @Test + void derivesCountForMatchWhereSet() { + assertThat(derive("MATCH (n:Person) WHERE n.age > 30 SET n.flag = true")) + .isEqualTo("MATCH (n:Person) WHERE n.age > 30 RETURN count(*) AS affected"); + } + + @Test + void derivesCountForDetachDelete() { + assertThat(derive("MATCH (u:User) DETACH DELETE u")) + .isEqualTo("MATCH (u:User) RETURN count(*) AS affected"); + } + + @Test + void derivesCountForRemove() { + assertThat(derive("MATCH (u:User) WHERE u.legacy = true REMOVE u.flag")) + .isEqualTo("MATCH (u:User) WHERE u.legacy = true RETURN count(*) AS affected"); + } + + @Test + void dropsTrailingReturnWithOrderAndLimit() { + assertThat(derive("MATCH (u:User) SET u.x = 1 RETURN u ORDER BY u.name LIMIT 5")) + .isEqualTo("MATCH (u:User) RETURN count(*) AS affected"); + } + + @Test + void keepsOptionalMatchInThePrefix() { + assertThat(derive("MATCH (u:User) OPTIONAL MATCH (u)-[:OWNS]->(a:Account) SET u.n = 1")) + .isEqualTo("MATCH (u:User) OPTIONAL MATCH (u)-[:OWNS]->(a:Account)" + + " RETURN count(*) AS affected"); + } + + @Test + void writeKeywordInsideStringOrBracketsIsIgnored() { + assertThat(derive("MATCH (u:User) WHERE u.note = 'please DELETE me' SET u.x = 1")) + .isEqualTo("MATCH (u:User) WHERE u.note = 'please DELETE me'" + + " RETURN count(*) AS affected"); + } + + @Test + void selectIsNotDerived() { + assertThat(derive("MATCH (u:User) RETURN u")).isNull(); + } + + @Test + void insertShapesAreNotDerived() { + assertThat(derive("CREATE (:User {id: 1})")).isNull(); + assertThat(derive("MERGE (u:User {id: 1}) SET u.seen = true")).isNull(); + assertThat(derive("CREATE (u:User) SET u.x = 1")).isNull(); + } + + @Test + void ddlIsNotDerived() { + assertThat(derive("DROP INDEX user_id IF EXISTS")).isNull(); + } + + @Test + void withPipelineBeforeWriteFailsClosed() { + assertThat(derive("MATCH (u:User) WITH u LIMIT 1 SET u.x = 1")).isNull(); + } + + @Test + void unwindBeforeWriteFailsClosed() { + assertThat(derive("UNWIND [1, 2] AS x MATCH (u:User {id: x}) SET u.x = 1")).isNull(); + } + + @Test + void readResumingAfterWriteFailsClosed() { + assertThat(derive("MATCH (u:User) SET u.x = 1 WITH u MATCH (v:User) SET v.y = 2")).isNull(); + } + + @Test + void writeWithoutMatchFailsClosed() { + assertThat(derive("SET u.x = 1")).isNull(); + } +} diff --git a/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngineIntegrationTest.java b/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngineIntegrationTest.java index ba304ee3..d8ca185c 100644 --- a/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngineIntegrationTest.java +++ b/engines/neo4j/src/test/java/com/bablsoft/accessflow/engine/neo4j/Neo4jQueryEngineIntegrationTest.java @@ -6,6 +6,7 @@ import com.bablsoft.accessflow.core.api.DbType; import com.bablsoft.accessflow.core.api.InvalidSqlException; import com.bablsoft.accessflow.core.api.MaskingStrategy; +import com.bablsoft.accessflow.core.api.QueryAffectedRowsResult; import com.bablsoft.accessflow.core.api.QueryDryRunResult; import com.bablsoft.accessflow.core.api.QueryEngineContext; import com.bablsoft.accessflow.core.api.QueryEngineDryRunRequest; @@ -275,6 +276,51 @@ void dryRunAppliesRowSecurity() { assertThat(result.appliedRowSecurityPolicyIds()).containsExactly(directive.policyId()); } + @Test + void countAffectedRowsForUpdateCountsMatchedNodesWithoutMutating() { + var result = countAffected("MATCH (u:User) WHERE u.id > 10 SET u.flag = true", List.of()); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(2); + assertThat(result.engineId()).isEqualTo("neo4j"); + var flagged = (SelectExecutionResult) run("MATCH (u:User) WHERE u.flag = true RETURN u"); + assertThat(flagged.rowCount()).isZero(); + } + + @Test + void countAffectedRowsAppliesRowSecurity() { + var directive = new RowSecurityDirective(UUID.randomUUID(), "User", "region", + RowSecurityOperator.EQUALS, List.of("EU")); + var result = countAffected("MATCH (u:User) DETACH DELETE u", List.of(directive)); + assertThat(result.supported()).isTrue(); + assertThat(result.affectedRows()).isEqualTo(2); + assertThat(((SelectExecutionResult) run("MATCH (u:User) RETURN u")).rowCount()).isEqualTo(3); + } + + @Test + void countAffectedRowsFailsClosedOnMerge() { + var result = countAffected("MERGE (u:User {id: 10}) SET u.seen = true", List.of()); + assertThat(result.supported()).isFalse(); + assertThat(result.affectedRows()).isNull(); + } + + @Test + void countAffectedRowsFailsClosedOnUnrewritableRowSecurityShape() { + run("MATCH (u:User {id: 10}) CREATE (u)-[:PLACED]->(:Order {id: 1})"); + var directive = new RowSecurityDirective(UUID.randomUUID(), "Order", "id", + RowSecurityOperator.EQUALS, List.of(1)); + var result = countAffected("MATCH (u:User)-[:PLACED]->(:Order) DELETE u", List.of(directive)); + assertThat(result.supported()).isFalse(); + } + + private static QueryAffectedRowsResult countAffected(String query, + List rls) { + var type = engine.parse(query).type(); + var request = new QueryExecutionRequest(descriptor.id(), query, type, null, null, + List.of(), List.of(), rls, false, List.of(query)); + return engine.countAffectedRows(new QueryEngineDryRunRequest(request, descriptor, + Duration.ofSeconds(30))); + } + private static QueryDryRunResult dryRun(String query, List rls) { var type = engine.parse(query).type(); var request = new QueryExecutionRequest(descriptor.id(), query, type, null, null, diff --git a/frontend/src/components/review/CostEstimatePanel.test.tsx b/frontend/src/components/review/CostEstimatePanel.test.tsx new file mode 100644 index 00000000..6dca0eae --- /dev/null +++ b/frontend/src/components/review/CostEstimatePanel.test.tsx @@ -0,0 +1,107 @@ +import { describe, expect, it } from 'vitest'; +import { render, screen } from '@testing-library/react'; +import '@/i18n'; +import type { CostEstimateDetail } from '@/types/api'; +import { CostEstimatePanel } from './CostEstimatePanel'; + +function estimate(overrides: Partial = {}): CostEstimateDetail { + return { + id: 'e1', + engine_id: 'postgresql', + query_type: 'DELETE', + supported: true, + estimated_rows: 2400000, + affected_row_count: null, + scan_type: 'Seq Scan', + estimated_cost: 44543.5, + plan: { operation: 'Seq Scan', target: 'users', estimated_rows: 2400000, children: [] }, + raw_plan: null, + unsupported_reason: null, + failed: false, + error_message: null, + duration_ms: 12, + ...overrides, + }; +} + +describe('CostEstimatePanel', () => { + it('shows the pending text while the query is still in PENDING_AI', () => { + render(); + expect(screen.getByText(/computing the cost estimate/i)).toBeInTheDocument(); + }); + + it('shows the unavailable fallback when no estimate exists past PENDING_AI', () => { + render(); + expect(screen.getByText(/no cost estimate is available/i)).toBeInTheDocument(); + }); + + it('renders the failure state with the error message', () => { + render( + , + ); + expect(screen.getByText(/cost estimation failed/i)).toBeInTheDocument(); + expect(screen.getByText(/timeout/)).toBeInTheDocument(); + }); + + it('renders the unsupported reason for engines without a plan', () => { + render( + , + ); + expect(screen.getByText(/not supported for the redis engine/i)).toBeInTheDocument(); + }); + + it('renders estimated rows, scan type, cost, and the plan tree', () => { + render(); + expect(screen.getAllByText('Seq Scan').length).toBeGreaterThan(0); + expect(screen.getByText('2,400,000')).toBeInTheDocument(); + expect(screen.getByText('44,543.5')).toBeInTheDocument(); + }); + + it('renders the exact affected-row count for writes when present', () => { + render( + , + ); + expect(screen.getByText(/affected rows \(exact\)/i)).toBeInTheDocument(); + expect(screen.getByText('90')).toBeInTheDocument(); + }); + + it('shows the affected-row count even when the plan itself is unsupported', () => { + render( + , + ); + expect(screen.getByText('7')).toBeInTheDocument(); + }); + + it('falls back to the raw plan when no structured plan is present', () => { + render( + , + ); + expect(screen.getByText('Seq Scan on users')).toBeInTheDocument(); + }); +}); diff --git a/frontend/src/components/review/CostEstimatePanel.tsx b/frontend/src/components/review/CostEstimatePanel.tsx new file mode 100644 index 00000000..d84f1aee --- /dev/null +++ b/frontend/src/components/review/CostEstimatePanel.tsx @@ -0,0 +1,148 @@ +import { InfoCircleOutlined, WarningOutlined } from '@ant-design/icons'; +import { useTranslation } from 'react-i18next'; +import type { CostEstimateDetail, QueryStatus } from '@/types/api'; +import { fmtNum } from '@/utils/dateFormat'; +import { formatCost } from '@/utils/queryPlan'; +import { PlanTree } from '@/components/editor/PlanTree'; + +interface CostEstimatePanelProps { + estimate: CostEstimateDetail | null; + status: QueryStatus; +} + +/** + * Persisted pre-flight cost / blast-radius panel on the query detail page (AF-624): the estimate + * computed automatically at submission — estimated rows, exact affected-row count for writes, scan + * type, cost, and the execution-plan tree — with pending / unavailable / failed fallbacks. + */ +export function CostEstimatePanel({ estimate, status }: CostEstimatePanelProps) { + const { t } = useTranslation(); + + if (!estimate) { + return ( +

+ {status === 'PENDING_AI' + ? t('cost_estimate_panel.pending') + : t('cost_estimate_panel.unavailable')} +
+ ); + } + + if (estimate.failed) { + return ( +
+ + + {t('cost_estimate_panel.failed')} + {estimate.error_message ? ` — ${estimate.error_message}` : ''} + +
+ ); + } + + const hasAffectedCount = + estimate.affected_row_count !== null && estimate.affected_row_count !== undefined; + + if (!estimate.supported && !hasAffectedCount) { + return ( +
+ + {estimate.unsupported_reason ?? t('cost_estimate_panel.unavailable')} +
+ ); + } + + const cost = formatCost(estimate.estimated_cost); + return ( +
+
+ {hasAffectedCount && ( +
+ {t('cost_estimate_panel.affected_rows_label')} + {fmtNum(estimate.affected_row_count as number)} +
+ )} +
+ {t('cost_estimate_panel.est_rows_label')} + + {estimate.estimated_rows === null || estimate.estimated_rows === undefined + ? '—' + : fmtNum(estimate.estimated_rows)} + +
+ {estimate.scan_type && ( +
+ {t('cost_estimate_panel.scan_type_label')} + {estimate.scan_type} +
+ )} + {cost && ( +
+ {t('cost_estimate_panel.cost_label')} + {cost} +
+ )} +
+ {estimate.plan ? ( + + ) : estimate.raw_plan ? ( + 
+          {estimate.raw_plan}
+
+ ) : null} +
+ ); +} diff --git a/frontend/src/locales/de.json b/frontend/src/locales/de.json index 3649cf7f..5f0f0900 100644 --- a/frontend/src/locales/de.json +++ b/frontend/src/locales/de.json @@ -523,6 +523,15 @@ "no_plan": "Die Engine hat keinen Plan für diese Abfrage zurückgegeben.", "rerun_button": "Erneut ausführen" }, + "cost_estimate_panel": { + "pending": "Kostenschätzung wird berechnet…", + "unavailable": "Für diese Abfrage ist keine Kostenschätzung verfügbar.", + "failed": "Kostenschätzung fehlgeschlagen", + "affected_rows_label": "Betroffene Zeilen (exakt)", + "est_rows_label": "Geschätzte Zeilen", + "scan_type_label": "Scan-Typ", + "cost_label": "Geschätzte Kosten" + }, "ai_panel": { "title": "KI-Analyse", "disabled_badge": "deaktiviert", @@ -577,6 +586,7 @@ "card_justification": "Begründung", "no_justification": "Der Antragsteller hat keine Begründung angegeben.", "card_ai": "KI-Analyse", + "card_cost_estimate": "Kostenschätzung", "ai_awaiting": "Warten auf Analyse…", "ai_failed_banner_title": "KI-Analyse fehlgeschlagen", "ai_failed_banner_detail": "Die Prüfung läuft ohne KI-Empfehlung weiter. Grund: {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Kein gültiger CIDR-Block: {{value}}", "user_agent_placeholder": "z. B. *curl*, *GitHubActions*", "minutes_suffix": "Min.", - "cicd_present_label": "Ist CI/CD-Herkunft" + "cicd_present_label": "Ist CI/CD-Herkunft", + "scan_type_placeholder": "z. B. Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "Quell-IP / CIDR", "user_agent": "User-Agent", "time_since_last_approval": "Zeit seit letzter Genehmigung", - "cicd_origin": "CI/CD-Herkunft" + "cicd_origin": "CI/CD-Herkunft", + "estimated_rows": "Geschätzte Zeilen", + "scan_type": "Scan-Typ" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/en.json b/frontend/src/locales/en.json index 2fa618b4..746ec2c8 100644 --- a/frontend/src/locales/en.json +++ b/frontend/src/locales/en.json @@ -523,6 +523,15 @@ "no_plan": "The engine returned no plan for this query.", "rerun_button": "Run again" }, + "cost_estimate_panel": { + "pending": "Computing the cost estimate…", + "unavailable": "No cost estimate is available for this query.", + "failed": "Cost estimation failed", + "affected_rows_label": "Affected rows (exact)", + "est_rows_label": "Estimated rows", + "scan_type_label": "Scan type", + "cost_label": "Estimated cost" + }, "ai_panel": { "title": "AI analysis", "disabled_badge": "disabled", @@ -577,6 +586,7 @@ "card_justification": "Justification", "no_justification": "Requester did not provide a justification.", "card_ai": "AI analysis", + "card_cost_estimate": "Cost estimate", "ai_awaiting": "Awaiting analysis…", "ai_failed_banner_title": "AI analysis failed", "ai_failed_banner_detail": "Review is proceeding without an AI recommendation. Reason: {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Not a valid CIDR block: {{value}}", "user_agent_placeholder": "e.g. *curl*, *GitHubActions*", "minutes_suffix": "min", - "cicd_present_label": "Is CI/CD origin" + "cicd_present_label": "Is CI/CD origin", + "scan_type_placeholder": "e.g. Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2394,7 +2405,9 @@ "source_ip": "Source IP / CIDR", "user_agent": "User-agent", "time_since_last_approval": "Time since last approval", - "cicd_origin": "CI/CD origin" + "cicd_origin": "CI/CD origin", + "estimated_rows": "Estimated rows", + "scan_type": "Scan type" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/es.json b/frontend/src/locales/es.json index 22ebc1a3..9ace6fc9 100644 --- a/frontend/src/locales/es.json +++ b/frontend/src/locales/es.json @@ -523,6 +523,15 @@ "no_plan": "El motor no devolvió ningún plan para esta consulta.", "rerun_button": "Ejecutar de nuevo" }, + "cost_estimate_panel": { + "pending": "Calculando la estimación de coste…", + "unavailable": "No hay estimación de coste disponible para esta consulta.", + "failed": "La estimación de coste falló", + "affected_rows_label": "Filas afectadas (exacto)", + "est_rows_label": "Filas estimadas", + "scan_type_label": "Tipo de escaneo", + "cost_label": "Coste estimado" + }, "ai_panel": { "title": "Análisis de IA", "disabled_badge": "deshabilitado", @@ -577,6 +586,7 @@ "card_justification": "Justificación", "no_justification": "El solicitante no proporcionó una justificación.", "card_ai": "Análisis de IA", + "card_cost_estimate": "Estimación de coste", "ai_awaiting": "Esperando análisis…", "ai_failed_banner_title": "El análisis de IA falló", "ai_failed_banner_detail": "La revisión continúa sin una recomendación de IA. Motivo: {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Bloque CIDR no válido: {{value}}", "user_agent_placeholder": "p. ej. *curl*, *GitHubActions*", "minutes_suffix": "min", - "cicd_present_label": "Es origen CI/CD" + "cicd_present_label": "Es origen CI/CD", + "scan_type_placeholder": "p. ej. Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "IP de origen / CIDR", "user_agent": "Agente de usuario", "time_since_last_approval": "Tiempo desde la última aprobación", - "cicd_origin": "Origen CI/CD" + "cicd_origin": "Origen CI/CD", + "estimated_rows": "Filas estimadas", + "scan_type": "Tipo de escaneo" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/fr.json b/frontend/src/locales/fr.json index 85774f7e..0ac01835 100644 --- a/frontend/src/locales/fr.json +++ b/frontend/src/locales/fr.json @@ -523,6 +523,15 @@ "no_plan": "Le moteur n'a renvoyé aucun plan pour cette requête.", "rerun_button": "Relancer" }, + "cost_estimate_panel": { + "pending": "Calcul de l’estimation du coût…", + "unavailable": "Aucune estimation de coût n’est disponible pour cette requête.", + "failed": "Échec de l’estimation du coût", + "affected_rows_label": "Lignes affectées (exact)", + "est_rows_label": "Lignes estimées", + "scan_type_label": "Type de scan", + "cost_label": "Coût estimé" + }, "ai_panel": { "title": "Analyse IA", "disabled_badge": "désactivée", @@ -577,6 +586,7 @@ "card_justification": "Justification", "no_justification": "Le demandeur n'a pas fourni de justification.", "card_ai": "Analyse IA", + "card_cost_estimate": "Estimation du coût", "ai_awaiting": "En attente d'analyse…", "ai_failed_banner_title": "Échec de l'analyse IA", "ai_failed_banner_detail": "La revue se poursuit sans recommandation de l'IA. Raison : {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Bloc CIDR non valide : {{value}}", "user_agent_placeholder": "p. ex. *curl*, *GitHubActions*", "minutes_suffix": "min", - "cicd_present_label": "Provenance CI/CD" + "cicd_present_label": "Provenance CI/CD", + "scan_type_placeholder": "p. ex. Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "IP source / CIDR", "user_agent": "Agent utilisateur", "time_since_last_approval": "Temps depuis la dernière approbation", - "cicd_origin": "Provenance CI/CD" + "cicd_origin": "Provenance CI/CD", + "estimated_rows": "Lignes estimées", + "scan_type": "Type de scan" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/hy.json b/frontend/src/locales/hy.json index e185f508..4181bd04 100644 --- a/frontend/src/locales/hy.json +++ b/frontend/src/locales/hy.json @@ -523,6 +523,15 @@ "no_plan": "Շարժիչը այս հարցման համար պլան չվերադարձրեց։", "rerun_button": "Կրկին գործարկել" }, + "cost_estimate_panel": { + "pending": "Հաշվարկվում է արժեքի գնահատումը…", + "unavailable": "Այս հարցման համար արժեքի գնահատում հասանելի չէ։", + "failed": "Արժեքի գնահատումը ձախողվեց", + "affected_rows_label": "Ազդված տողեր (ճշգրիտ)", + "est_rows_label": "Գնահատված տողեր", + "scan_type_label": "Սկանավորման տեսակ", + "cost_label": "Գնահատված արժեք" + }, "ai_panel": { "title": "AI վերլուծություն", "disabled_badge": "անջատված", @@ -577,6 +586,7 @@ "card_justification": "Հիմնավորում", "no_justification": "Հայցողը հիմնավորում չի տրամադրել։", "card_ai": "AI վերլուծություն", + "card_cost_estimate": "Արժեքի գնահատում", "ai_awaiting": "Սպասում է վերլուծության…", "ai_failed_banner_title": "AI վերլուծությունը ձախողվեց", "ai_failed_banner_detail": "Վերանայումը շարունակվում է առանց AI առաջարկության։ Պատճառ՝ {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Անվավեր CIDR բլոկ՝ {{value}}", "user_agent_placeholder": "օր. *curl*, *GitHubActions*", "minutes_suffix": "ր", - "cicd_present_label": "CI/CD ծագում" + "cicd_present_label": "CI/CD ծագում", + "scan_type_placeholder": "օր.՝ Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "Աղբյուրի IP / CIDR", "user_agent": "User-Agent", "time_since_last_approval": "Ժամանակ վերջին հաստատումից", - "cicd_origin": "CI/CD ծագում" + "cicd_origin": "CI/CD ծագում", + "estimated_rows": "Գնահատված տողեր", + "scan_type": "Սկանավորման տեսակ" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/ru.json b/frontend/src/locales/ru.json index 1d21caa5..27bad69b 100644 --- a/frontend/src/locales/ru.json +++ b/frontend/src/locales/ru.json @@ -523,6 +523,15 @@ "no_plan": "Движок не вернул план для этого запроса.", "rerun_button": "Запустить снова" }, + "cost_estimate_panel": { + "pending": "Вычисляется оценка стоимости…", + "unavailable": "Оценка стоимости для этого запроса недоступна.", + "failed": "Оценка стоимости не удалась", + "affected_rows_label": "Затронутые строки (точно)", + "est_rows_label": "Оценка строк", + "scan_type_label": "Тип сканирования", + "cost_label": "Оценочная стоимость" + }, "ai_panel": { "title": "ИИ-анализ", "disabled_badge": "отключено", @@ -577,6 +586,7 @@ "card_justification": "Обоснование", "no_justification": "Заявитель не указал обоснование.", "card_ai": "ИИ-анализ", + "card_cost_estimate": "Оценка стоимости", "ai_awaiting": "Ожидание анализа…", "ai_failed_banner_title": "ИИ-анализ не удался", "ai_failed_banner_detail": "Проверка продолжается без рекомендации ИИ. Причина: {{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "Недопустимый блок CIDR: {{value}}", "user_agent_placeholder": "напр. *curl*, *GitHubActions*", "minutes_suffix": "мин", - "cicd_present_label": "Источник CI/CD" + "cicd_present_label": "Источник CI/CD", + "scan_type_placeholder": "напр. Seq Scan, COLLSCAN, Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "IP-адрес источника / CIDR", "user_agent": "User-Agent", "time_since_last_approval": "Время с последнего одобрения", - "cicd_origin": "Источник CI/CD" + "cicd_origin": "Источник CI/CD", + "estimated_rows": "Оценка строк", + "scan_type": "Тип сканирования" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/locales/zh-CN.json b/frontend/src/locales/zh-CN.json index a916a5f2..f3dbbdc0 100644 --- a/frontend/src/locales/zh-CN.json +++ b/frontend/src/locales/zh-CN.json @@ -523,6 +523,15 @@ "no_plan": "引擎未返回此查询的计划。", "rerun_button": "重新运行" }, + "cost_estimate_panel": { + "pending": "正在计算成本预估…", + "unavailable": "此查询没有可用的成本预估。", + "failed": "成本预估失败", + "affected_rows_label": "受影响行数（精确）", + "est_rows_label": "预估行数", + "scan_type_label": "扫描类型", + "cost_label": "预估成本" + }, "ai_panel": { "title": "AI 分析", "disabled_badge": "已禁用", @@ -577,6 +586,7 @@ "card_justification": "理由", "no_justification": "申请人未提供说明。", "card_ai": "AI 分析", + "card_cost_estimate": "成本预估", "ai_awaiting": "正在等待分析…", "ai_failed_banner_title": "AI 分析失败", "ai_failed_banner_detail": "审核将继续，但没有 AI 建议。原因：{{reason}}", @@ -1952,7 +1962,8 @@ "cidr_invalid": "无效的 CIDR 网段：{{value}}", "user_agent_placeholder": "例如 *curl*、*GitHubActions*", "minutes_suffix": "分钟", - "cicd_present_label": "来自 CI/CD" + "cicd_present_label": "来自 CI/CD", + "scan_type_placeholder": "例如 Seq Scan、COLLSCAN、Index*" }, "langfuse": { "title": "Langfuse", @@ -2365,7 +2376,9 @@ "source_ip": "来源 IP / CIDR", "user_agent": "用户代理", "time_since_last_approval": "距上次批准时长", - "cicd_origin": "CI/CD 来源" + "cicd_origin": "CI/CD 来源", + "estimated_rows": "预估行数", + "scan_type": "扫描类型" }, "row_security_operator": { "EQUALS": "=", diff --git a/frontend/src/pages/admin/RoutingPoliciesPage.tsx b/frontend/src/pages/admin/RoutingPoliciesPage.tsx index bcf518e2..5521f636 100644 --- a/frontend/src/pages/admin/RoutingPoliciesPage.tsx +++ b/frontend/src/pages/admin/RoutingPoliciesPage.tsx @@ -815,6 +815,31 @@ function ConditionValueEditor({ name, operand, groups, roleOptions }: ConditionV ); + case 'estimated_rows': + return ( +
+ + + + ); default: return null; } diff --git a/frontend/src/pages/admin/routingPolicyForm.test.ts b/frontend/src/pages/admin/routingPolicyForm.test.ts index ccfe8f99..ca64ad44 100644 --- a/frontend/src/pages/admin/routingPolicyForm.test.ts +++ b/frontend/src/pages/admin/routingPolicyForm.test.ts @@ -74,6 +74,8 @@ describe('routingPolicyForm', () => { tsla_minutes: 1440, }, { operand: 'cicd_origin', negate: false, bool_value: true }, + { operand: 'estimated_rows', negate: false, est_operator: 'GT', est_value: 100000 }, + { operand: 'scan_type', negate: true, scan_patterns: ['Seq*'] }, ]; const condition = rowsToCondition('ALL', rows); const parsed = conditionToForm(condition); @@ -188,4 +190,27 @@ describe('routingPolicyForm', () => { ]); expect(conditionSummary(t, condition)).toContain('22:00'); }); + + it('defaultRow and conditionSummary cover estimated_rows and scan_type (AF-624)', () => { + expect(defaultRow('estimated_rows')).toEqual({ + operand: 'estimated_rows', + negate: false, + est_operator: 'GT', + est_value: 100000, + }); + expect(defaultRow('scan_type')).toEqual({ + operand: 'scan_type', + negate: false, + scan_patterns: [], + }); + const summary = conditionSummary( + t, + rowsToCondition('ALL', [ + { operand: 'estimated_rows', negate: false, est_operator: 'GT', est_value: 100000 }, + { operand: 'scan_type', negate: false, scan_patterns: ['Seq Scan', 'COLLSCAN'] }, + ]), + ); + expect(summary).toContain('100000'); + expect(summary).toContain('Seq Scan, COLLSCAN'); + }); }); diff --git a/frontend/src/pages/admin/routingPolicyForm.ts b/frontend/src/pages/admin/routingPolicyForm.ts index 7b2c8a5f..568eba05 100644 --- a/frontend/src/pages/admin/routingPolicyForm.ts +++ b/frontend/src/pages/admin/routingPolicyForm.ts @@ -44,6 +44,9 @@ export interface RoutingConditionRow { ua_patterns?: string[]; tsla_operator?: ComparisonOperator; tsla_minutes?: number; + est_operator?: ComparisonOperator; + est_value?: number; + scan_patterns?: string[]; } export interface RoutingPolicyFormValues { @@ -91,6 +94,10 @@ export function defaultRow(operand: RoutingConditionOperand): RoutingConditionRo return { ...base, tsla_operator: 'GT', tsla_minutes: 1440 }; case 'cicd_origin': return { ...base, bool_value: true }; + case 'estimated_rows': + return { ...base, est_operator: 'GT', est_value: 100000 }; + case 'scan_type': + return { ...base, scan_patterns: [] }; default: return base; } @@ -166,6 +173,14 @@ function rowToLeaf(row: RoutingConditionRow): RoutingCondition { }; case 'cicd_origin': return { type: 'cicd_origin', expected: row.bool_value ?? false }; + case 'estimated_rows': + return { + type: 'estimated_rows', + operator: row.est_operator ?? 'GT', + value: row.est_value ?? 0, + }; + case 'scan_type': + return { type: 'scan_type', patterns: row.scan_patterns ?? [] }; } } @@ -227,6 +242,15 @@ function leafToRow(node: RoutingCondition, negate: boolean): RoutingConditionRow }; case 'cicd_origin': return { operand: 'cicd_origin', negate, bool_value: node.expected }; + case 'estimated_rows': + return { + operand: 'estimated_rows', + negate, + est_operator: node.operator, + est_value: node.value, + }; + case 'scan_type': + return { operand: 'scan_type', negate, scan_patterns: node.patterns }; default: // Nested and/or/not (other than not-of-leaf) cannot be represented by the flat builder. return null; @@ -343,6 +367,12 @@ function rowSummary(t: TFunction, row: RoutingConditionRow): string { value = `${comparisonOperatorLabel(t, row.tsla_operator ?? 'GT')} ${row.tsla_minutes ?? 0} ` + t('admin.routing_policies.minutes_suffix'); break; + case 'estimated_rows': + value = `${comparisonOperatorLabel(t, row.est_operator ?? 'GT')} ${row.est_value ?? 0}`; + break; + case 'scan_type': + value = (row.scan_patterns ?? []).join(', '); + break; } return `${prefix}${label}: ${value}`; } diff --git a/frontend/src/pages/queries/QueryDetailPage.test.tsx b/frontend/src/pages/queries/QueryDetailPage.test.tsx index d8b33605..0d89525d 100644 --- a/frontend/src/pages/queries/QueryDetailPage.test.tsx +++ b/frontend/src/pages/queries/QueryDetailPage.test.tsx @@ -116,6 +116,7 @@ function failedQuery(): QueryDetail { failed: true, error_message: 'provider unavailable', }, + cost_estimate: null, rows_affected: null, duration_ms: null, error_message: null, diff --git a/frontend/src/pages/queries/QueryDetailPage.tsx b/frontend/src/pages/queries/QueryDetailPage.tsx index 26b4bd29..c1146d61 100644 --- a/frontend/src/pages/queries/QueryDetailPage.tsx +++ b/frontend/src/pages/queries/QueryDetailPage.tsx @@ -26,6 +26,7 @@ import { QueryTypePill } from '@/components/common/QueryTypePill'; import { SqlBlock } from '@/components/common/SqlBlock'; import { DetailCard } from '@/components/common/DetailCard'; import { ApprovalTimeline, type TimelineStage } from '@/components/review/ApprovalTimeline'; +import { CostEstimatePanel } from '@/components/review/CostEstimatePanel'; import { IssueCard } from '@/components/editor/IssueCard'; import { OptimizationCard } from '@/components/editor/OptimizationCard'; import { QueryCollaboration } from '@/components/editor/QueryCollaboration'; @@ -590,6 +591,13 @@ export function QueryDetailPage() { )} + } + > + + + {query.status === 'EXECUTED' && ( }>
): QueryDetail { status: 'PENDING_AI' as QueryStatus, justification: '', ai_analysis: null, + cost_estimate: null, rows_affected: null, duration_ms: null, error_message: null, diff --git a/frontend/src/realtime/websocketManager.ts b/frontend/src/realtime/websocketManager.ts index 29f10f16..18f5a6fe 100644 --- a/frontend/src/realtime/websocketManager.ts +++ b/frontend/src/realtime/websocketManager.ts @@ -189,6 +189,7 @@ class WebSocketManager { this.queryClient.invalidateQueries({ queryKey: ['queries', 'list'] }); break; case 'ai.analysis_complete': + case 'query.estimate_complete': if (queryId) this.queryClient.invalidateQueries({ queryKey: ['queries', 'detail', queryId] }); break; case 'review.new_request': diff --git a/frontend/src/types/api.ts b/frontend/src/types/api.ts index 77cb09ae..bccb99a6 100644 --- a/frontend/src/types/api.ts +++ b/frontend/src/types/api.ts @@ -75,7 +75,9 @@ export type RoutingConditionOperand = | 'source_ip' | 'user_agent' | 'time_since_last_approval' - | 'cicd_origin'; + | 'cicd_origin' + | 'estimated_rows' + | 'scan_type'; export type Weekday = | 'MONDAY' | 'TUESDAY' @@ -1145,6 +1147,24 @@ export interface AiAnalysisDetail { error_message: string | null; } +/** Persisted pre-flight cost / blast-radius estimate on a submitted query (AF-624). */ +export interface CostEstimateDetail { + id: string; + engine_id: string | null; + query_type: QueryType | null; + supported: boolean; + estimated_rows: number | null; + affected_row_count: number | null; + scan_type: string | null; + estimated_cost: number | null; + plan: QueryPlanNode | null; + raw_plan: string | null; + unsupported_reason: string | null; + failed: boolean; + error_message: string | null; + duration_ms: number | null; +} + export interface ReviewDecisionDetail { id: string; reviewer: UserRef; @@ -1164,6 +1184,7 @@ export interface QueryDetail { status: QueryStatus; justification: string; ai_analysis: AiAnalysisDetail | null; + cost_estimate: CostEstimateDetail | null; rows_affected: number | null; duration_ms: number | null; error_message: string | null; @@ -1231,7 +1252,9 @@ export type RoutingCondition = | { type: 'source_ip'; cidrs: string[] } | { type: 'user_agent'; patterns: string[] } | { type: 'time_since_last_approval'; operator: ComparisonOperator; minutes: number } - | { type: 'cicd_origin'; expected: boolean }; + | { type: 'cicd_origin'; expected: boolean } + | { type: 'estimated_rows'; operator: ComparisonOperator; value: number } + | { type: 'scan_type'; patterns: string[] }; export interface RoutingPolicy { id: string; diff --git a/frontend/src/types/ws.ts b/frontend/src/types/ws.ts index 732cce1a..6562d7ff 100644 --- a/frontend/src/types/ws.ts +++ b/frontend/src/types/ws.ts @@ -12,6 +12,7 @@ export type WsEventName = | 'review.decision_made' | 'query.executed' | 'ai.analysis_complete' + | 'query.estimate_complete' | 'notification.created' | 'anomaly.detected' | 'collab.joined' @@ -63,6 +64,10 @@ export interface WsEventPayloadMap { risk_level: RiskLevel | null; risk_score: number | null; }; + 'query.estimate_complete': { + query_id: string; + supported: boolean; + }; 'notification.created': { notification_id: string; event_type: UserNotificationEventType; @@ -152,6 +157,7 @@ export const WS_EVENT_NAMES: ReadonlyArray = [ 'review.decision_made', 'query.executed', 'ai.analysis_complete', + 'query.estimate_complete', 'notification.created', 'anomaly.detected', 'collab.joined', diff --git a/frontend/src/utils/enumLabels.ts b/frontend/src/utils/enumLabels.ts index 9838b3d9..404098ad 100644 --- a/frontend/src/utils/enumLabels.ts +++ b/frontend/src/utils/enumLabels.ts @@ -313,6 +313,8 @@ export const CONDITION_OPERANDS: readonly RoutingConditionOperand[] = [ 'user_agent', 'time_since_last_approval', 'cicd_origin', + 'estimated_rows', + 'scan_type', ] as const; export const routingActionLabel = (t: TFunction, v: RoutingAction): string => diff --git a/website/docs/index.html b/website/docs/index.html index bc6040a6..7a5d977c 100644 --- a/website/docs/index.html +++ b/website/docs/index.html @@ -1534,7 +1534,7 @@

Routing policies

  1. Open /admin/routing-policies (the Routing policies entry in the Security nav group, next to Review plans) and click Add policy.
  2. Name the policy and optionally scope it to one datasource — leave the datasource blank for an org-wide rule. Set its priority (unique per organisation; lower runs first) and the enabled toggle.
    3. -
  3. Build the condition with the guided builder: pick match ALL (AND) or match ANY (OR), then add leaf conditions — each can be negated (NOT). Operands include query type, referenced tables (glob, e.g. payroll.*), AI risk level, AI risk score (with a comparison operator), requester role, requester group, time-of-day window, day-of-week, presence of a WHERE clause, presence of a LIMIT clause, and the transactional (BEGIN…COMMIT) flag.
    4. +
  4. Build the condition with the guided builder: pick match ALL (AND) or match ANY (OR), then add leaf conditions — each can be negated (NOT). Operands include query type, referenced tables (glob, e.g. payroll.*), AI risk level, AI risk score (with a comparison operator), requester role, requester group, time-of-day window, day-of-week, presence of a WHERE clause, presence of a LIMIT clause, the transactional (BEGIN…COMMIT) flag, and the pre-flight cost estimate — estimated rows (comparison against the engine's own EXPLAIN estimate, or the exact affected-row count for UPDATE/DELETE) and scan type (glob match on the plan's root operation, e.g. Seq*) — so a policy can route a 10-million-row sequential-scan DELETE differently from a 10-row indexed one.
  5. Choose the action. Auto-approve (skip human review), Auto-reject (block the query), Require approvals (force human review with an absolute minimum number of approvers), or Escalate (force human review, adding a delta on top of the review plan's minimum). The approver count applies only to the last two actions.
  6. Reorder policies any time with the per-row up/down controls — the order is the evaluation order.
@@ -1542,7 +1542,9 @@

Routing policies

How it routes. Time-of-day and day-of-week conditions are evaluated in the server's local timezone (overnight windows wrap around midnight). On datasources with AI analysis disabled, risk-based conditions never match (there's no AI signal); routing does not - run when AI analysis fails — the query goes to a human instead. Every automated decision is + run when AI analysis fails — the query goes to a human instead. The cost-estimate conditions + likewise never match while no estimate exists (engine without a plan concept, or the estimate + failed) — they fail closed rather than auto-approving blind. Every automated decision is recorded in the audit log (QUERY_APPROVED / QUERY_REJECTED with source: "ROUTING_POLICY"), and the query detail page shows which policy matched. Routing policies are managed via the ADMIN-only diff --git a/website/index.html b/website/index.html index e424eac6..ba64182b 100644 --- a/website/index.html +++ b/website/index.html @@ -328,7 +328,7 @@

The middle ground between blanket access and a DBA ticket queue.

Full query proxy — SQL and NoSQL

-

Every query is inspected before it ever reaches the database — parsed, sorted by type (read, write, or schema change), checked against the tables a user is allowed to touch, and routed through your review policy. The same governance covers NoSQL and the cloud data warehouses too: MongoDB, Couchbase, Redis, Cassandra, Elasticsearch, DynamoDB, Neo4j, Snowflake, BigQuery, and Databricks all flow through identical checks, with dangerous or server-side commands blocked up front. Layer on row-level security so users only see the rows they're cleared for, dynamic masking that hides sensitive columns per person, and data-classification tags (PII, PCI, PHI, GDPR) that auto-apply masking and raise a query's risk score. A dry-run preview shows a query's impact before review without changing any data, and every executed query is saved as an exact snapshot you can replay in a test environment.

+

Every query is inspected before it ever reaches the database — parsed, sorted by type (read, write, or schema change), checked against the tables a user is allowed to touch, and routed through your review policy. The same governance covers NoSQL and the cloud data warehouses too: MongoDB, Couchbase, Redis, Cassandra, Elasticsearch, DynamoDB, Neo4j, Snowflake, BigQuery, and Databricks all flow through identical checks, with dangerous or server-side commands blocked up front. Layer on row-level security so users only see the rows they're cleared for, dynamic masking that hides sensitive columns per person, and data-classification tags (PII, PCI, PHI, GDPR) that auto-apply masking and raise a query's risk score. A dry-run preview shows a query's impact before review without changing any data, and every submitted query automatically gets a pre-flight cost estimate — the database's own EXPLAIN plan plus an exact affected-row count for updates and deletes — shown to reviewers, folded into the AI risk analysis, and usable in routing policies (“auto-escalate any full-scan DELETE over 100k rows”). Every executed query is saved as an exact snapshot you can replay in a test environment.

PostgreSQL · MySQL · MariaDB · Oracle · MSSQL · ClickHouse · MongoDB · Couchbase · Redis · Cassandra · ScyllaDB · Elasticsearch · OpenSearch · DynamoDB · Neo4j · Snowflake · BigQuery · Databricks · connector catalog · custom drivers · read-replica load balancing · result caching · dynamic masking · row-level security · data classification