diff --git a/docs/README.md b/docs/README.md index 4ad8dfe..f7d9fb2 100644 --- a/docs/README.md +++ b/docs/README.md @@ -12,6 +12,7 @@ | Contribute | [CONTRIBUTING.md](../CONTRIBUTING.md) and [Testing isolation](guides/testing-isolation.md) | | Know what is next, and what nobody has checked | [ROADMAP.md](../ROADMAP.md) | | Deploy the Databricks side | [Deploying Databricks](guides/deploying-databricks.md) | +| Attribute tenant job cost in currency | [Billing cost attribution](guides/billing-cost-attribution.md) | | Cut a release | [Releasing](guides/releasing.md) | | Assess it for security or compliance | [Threat model](security/threat-model.md), [SOC 2 mapping](compliance/soc2-mapping.md), [Data handling](compliance/data-handling.md) | diff --git a/docs/compatibility.md b/docs/compatibility.md index 7904330..f2ecc64 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -8,7 +8,8 @@ surrounding prose sounds. Last updated 2026-09-01 for the 1.0.1 release. The core Databricks statement and Jobs paths were re-run against the active development workspace, and the bundle was validated there. The proxy is verified against a real application database; the Bicep template compiles against the public -schema. The real currency path remains blocked on a metastore-admin grant. +schema. The currency path is implemented and locally verified, but its system-table read remains +blocked on a metastore-admin grant in the verification workspace. A `Documented` row is not promoted to `Verified` by a stable version number; the matrix records the work that has been done, not a claim about future work. @@ -63,6 +64,7 @@ request, package-consumer, publication and cleanup commands are in the | Capability | Status | Date | Evidence | |---|---|---|---| | Elapsed-time cost attribution (per-tenant, per-kind DBU) | **Documented** | 2026-08-29 | `OperationCostAttribution` against a real Postgres; the aggregation runs in `EXTRACT(EPOCH FROM (CompletedAt - ClaimedAt))` and is bounded by `ClaimedAt` to `CompletedAt` to exclude in-flight work. ADR 0012. | +| Billing cost attribution (per-tenant DBU and effective list-price currency) | **Documented** | 2026-09-01 | Local tests cover indexed distinct tenant-owned run selection in PostgreSQL, the 500-run/one-query budget, bound workspace/run/window parameters, report/price overlap proration, corrections, malformed rows, polling deadline and best-effort cancellation. The live system-table read is not verified because the development identity lacks the grant. [Runbook](guides/billing-cost-attribution.md), ADR 0012. | | Bicep reference deployment template compiles | **Documented** | 2026-08-29 | `az bicep build --file infra/azure-container-apps/main.bicep` against the public schema. No deploy has been run. ADR 0014. | | OpenTelemetry export via the sample's opt-in pipeline | **Documented** | 2026-08-29 | The sample's `Program.cs` subscribes to `LakeWright.Multitenancy` when `Lakewright:OpenTelemetry:Enabled=true`. Vendor-specific wiring is the adopter's. ADR 0013. | | Entra ID token accepted as a Databricks bearer token (user principal) | **Verified** | 2026-09-01 | `LiveDatabricksTests`; previously [spike 01](planning/spike-01-statement-execution.md) | @@ -121,10 +123,10 @@ request, package-consumer, publication and cleanup commands are in the ## Known gaps -- No per-tenant cost attribution **in currency**. The elapsed-time proxy in - `OperationCostAttribution` ships, labelled `CostSource.Proxy`; a real billing read remains - blocked on a metastore-admin grant on `system.billing.usage` — see T5 in - [the threat model](security/threat-model.md) and ADR 0012. +- Billing attribution is not live-verified in the development workspace. The implementation ships + as an opt-in, but the development identity still lacks access to `system.billing.usage` and + `system.billing.list_prices`; the proxy remains the default. See the + [billing runbook](guides/billing-cost-attribution.md), T5 in the threat model and ADR 0012. - No reference deployment *executed*. The Bicep template in `infra/azure-container-apps/` compiles; the workflow in `.github/workflows/deploy-azure.yml` is in place; no one has run a deploy with it, so the ingress half of encryption in transit and the managed identity path in diff --git a/docs/decisions/0012-elapsed-time-cost-attribution.md b/docs/decisions/0012-elapsed-time-cost-attribution.md index 2967519..78a4ebd 100644 --- a/docs/decisions/0012-elapsed-time-cost-attribution.md +++ b/docs/decisions/0012-elapsed-time-cost-attribution.md @@ -1,4 +1,4 @@ -# ADR 0012: Cost attribution as an interface, with the elapsed-time proxy as the first implementation +# ADR 0012: Cost attribution from an elapsed-time proxy or correlated billing usage Status: accepted Date: 2026-08-29 @@ -15,19 +15,49 @@ That leaves two questions for the codebase. What does the cost API look like, an **The first implementation, `OperationCostAttribution` in `LakeWright.Multitenancy`, reads the application database.** It sums `EXTRACT(EPOCH FROM (CompletedAt - ClaimedAt))` per kind, weights by the configured DBU/hour, and labels the result `CostSource.Proxy`. Only terminal-state operations are counted, because an open-ended duration is not a cost number. The aggregation runs in Postgres rather than the application: a 100k-row operations table pulled across the wire to sum in memory is the kind of "small data" that is small until it isn't. -**The interface exists so a real billing-table read replaces this implementation without changing the call sites.** The `CostSource` discriminator tells the caller which one ran, and `WarehouseSku` is null on a billing read because the rate comes from billing data rather than configuration. +**A billing implementation is split at the existing dependency boundary.** +`BillingCostAttribution` in `LakeWright.Multitenancy` selects tenant-owned job run ids and their +operation kinds from PostgreSQL. `IBillingUsageReader` in `LakeWright.Core` is the typed seam; +`DatabricksBillingUsageReader` in `LakeWright.Databricks` reads only those run ids from +`system.billing.usage`. `LakeWright.AspNetCore` composes both. Multitenancy therefore does not +reference the Databricks integration, and PostgreSQL is never assumed to be visible to Databricks +SQL. + +**The billing query is fixed SQL with bound values.** It filters both `workspace_id` and +`usage_metadata.job_run_id`, plus timestamp and `usage_date` bounds. PostgreSQL projects distinct +run ownership through a tenant/time partial covering index and stops at 501 rows. At most 500 run +ids enter one bound system-table query; a larger report returns `REPORT_TOO_LARGE` and HTTP 422 +rather than multiplying account-wide scans. Correlation happens in application code. A billing +row for an id not selected from the tenant's operations is rejected rather than ignored. + +**Currency is explicit and additive.** `TenantCostSummary` and `CostByKind` retain their original +constructors and gain init-only `EstimatedListCost` collections. Each `CurrencyAmount` keeps its +currency code beside its amount; unlike currencies are never added. Cost is calculated from +`system.billing.list_prices.pricing.effective_list.default`. A usage row is split at report and +price-validity boundaries; its quantity is prorated by each overlap and that same quantity feeds +both DBUs and cost. This is estimated effective list-price cost, not an adopter's negotiated invoice +amount. The query sums every billing record, including negative correction rows, before the +application aggregates by operation kind. Kinds are ordered by DBUs and then name, never by adding +unlike currency amounts. **Configuration belongs to the adopter, not the library.** A product running more than one warehouse SKU picks one as the proxy, or moves to a billing read. `CostAttributionOptions` carries the SKU and DBU/hour rate; the section is bound by `AddLakeWrightCostAttribution` rather than read by `AddLakeWright`, so a contributor working on the application without a workspace never hits a validation failure on a value they have no opinion about. **The endpoint is opt-in via `MapLakeWrightCost`, behind the Viewer policy, with a 31-day window cap.** A customer-facing usage page that lets a tenant ask for a multi-year range is asking Postgres to sum a range nobody actually wants. +**Billing statement work has an overall deadline.** `PollingTimeoutSeconds` defaults to 120. +Caller cancellation, deadline expiry, and poll transport failures all trigger a five-second +best-effort statement cancellation; cancellation failure never replaces the original result. + ## Consequences **The threat model updates from "partly mitigated" to "mitigated with a proxy."** The proxy is documented as such: the per-tenant report is "elapsed compute time on the configured warehouse SKU" and not currency, and the dollar number an operator would want is hours late rather than actionable in flight. The ceiling in `OperationWorker:MaxInFlightPerTenant` remains the control that acts in time; this one reports afterwards. -**`system.billing.usage` is now a real extension point rather than a sentence in the threat model.** A product that gets the grant wires an `ICostAttribution` against the API, returns `CostSource.Billing`, and the rest of the codebase reads it the same way. No second interface, no second endpoint, no migration. +**`system.billing.usage` is now a shipping opt-in rather than only an extension point.** A product +with the system-table grants calls `AddLakeWrightBillingCostAttribution` after the base Databricks +registration. Workspaces without those grants keep the proxy. Billing records can arrive hours +after a run, so the billing report is eventually consistent and an absent row is not replaced with +proxy data. **A test pins the property that the library's own instruments never carry a tenant id.** Per-tenant totals come from `operations` and `audit_events` rather than from the metrics, which is the property the cardinality-bomb rule was written to protect. A future change that adds a `tenant` or `tenantid` tag to a metric call site fails the build with the offending line. **`CostAttribution` joins the v0.2 milestone rather than v0.1.** v0.1 was the eight-week milestone whose definition of done did not name this. The published version is `0.1.2-preview.1`; the next published version with cost attribution is `0.2.0-preview.1`, which carries a breaking-change note for the `Operation` rows the implementation reads (none, today, but documented for the next maintainer). - diff --git a/docs/guides/billing-cost-attribution.md b/docs/guides/billing-cost-attribution.md new file mode 100644 index 0000000..eca3caf --- /dev/null +++ b/docs/guides/billing-cost-attribution.md @@ -0,0 +1,78 @@ +# Billing cost attribution + +The default `OperationCostAttribution` is deliberately a DBU proxy. Use the billing reader only +when the application identity can read the Databricks billing system tables and the product needs +currency-denominated effective list cost. + +## Prerequisites + +- The Databricks identity registered by `AddLakeWrightDatabricks` can use + `system.billing` and select from `system.billing.usage` and + `system.billing.list_prices`. +- `DatabricksBilling:WorkspaceId` is the workspace whose Lakeflow run ids are stored in this + application's `operations.ExternalId` column. +- Operation processing uses the LakeWright worker, so `ExternalId` is the Lakeflow + `job_run_id`. Statement ids and job ids are different identifiers and are not accepted. + +System billing rows are account-wide and typically arrive after the workload finishes. The reader +always filters `workspace_id`, job run id, timestamp and `usage_date`; a report can therefore be +empty while recent records are still being delivered. + +## Registration + +```csharp +builder.Services.AddLakeWright(builder.Configuration); +builder.Services.AddLakeWrightDatabricks(builder.Configuration); +builder.Services.AddLakeWrightBillingCostAttribution(builder.Configuration); +``` + +```json +{ + "Databricks": { + "WorkspaceUrl": "https://adb-....azuredatabricks.net", + "WarehouseId": "..." + }, + "DatabricksBilling": { + "WorkspaceId": "...", + "PollIntervalMilliseconds": 250, + "PollingTimeoutSeconds": 120 + } +} +``` + +`AddLakeWrightBillingCostAttribution` replaces the proxy registration. Keep +`AddLakeWrightCostAttribution` instead in environments without the grants. + +The response retains the existing DBU fields. `EstimatedListCost` is a collection of +`CurrencyAmount` values on both the summary and each operation-kind row. The amount uses +`pricing.effective_list.default`; it does not include a private negotiated discount and must not be +presented as an invoice total. `ElapsedSeconds` is zero for billing rows because the billing table +reports quantities and usage intervals, not the operation wall-clock value used by the proxy. + +Usage rows that cross either report-window or price-validity boundaries are prorated by their +overlap. The same prorated quantity feeds both DBUs and effective list cost, and a price change +inside one usage row contributes one segment at each effective price. The endpoint never adds +amounts in unlike currencies when ordering operation kinds; it orders by DBUs, then kind. + +One report is limited to 500 distinct tenant-owned job runs and issues one billing-system query. +HTTP 422 with code `REPORT_TOO_LARGE` means the caller must narrow the window. This prevents a +high-volume tenant from turning one request into repeated scans of the account-wide billing table. +A statement that remains pending past `PollingTimeoutSeconds` is cancelled best-effort and returns +the transient code `POLL_TIMEOUT`. + +## Live verification + +Run this only in the non-production workspace whose id is configured above. + +1. Start one known LakeWright operation and wait for its `ExternalId` job run to finish. +2. Wait for the corresponding `system.billing.usage` record to arrive. Databricks documents a + typical delay of up to 12 hours for original records. +3. Query the cost endpoint for a window that contains the run. Confirm `Source` is `Billing`, the + run contributes once to `Operations`, and `EstimatedListCost` carries the expected currency. +4. Compare the returned DBUs and effective list cost with a direct, read-only query over the same + workspace id, job run id and window. +5. Remove the system-table grants from a disposable verification principal and confirm the endpoint + returns HTTP 502 with `PERMISSION_DENIED`; restore the grant afterwards. + +No test data needs to be written to the system tables. Do not use a production workspace merely to +obtain a billing row. diff --git a/docs/security/threat-model.md b/docs/security/threat-model.md index 7fea8b9..117e80c 100644 --- a/docs/security/threat-model.md +++ b/docs/security/threat-model.md @@ -84,12 +84,11 @@ caps how many operations one tenant can have running at once, across every worke ceiling on the compute a runaway loop can buy before anyone notices. Warehouse auto-stop bounds the idle half. Evidence: `OperationClaimTests.A_tenant_at_its_ceiling_is_skipped_rather_than_failed`. -The reporting half is `ICostAttribution` (ADR 0012), with the elapsed-time proxy as its first -implementation. The proxy sums `EXTRACT(EPOCH FROM (CompletedAt - ClaimedAt))` per kind and -weights it by the configured warehouse SKU's DBU/hour. The result is `CostSource.Proxy` and -is documented as a proxy rather than a currency read. The interface exists so a real -billing-table read replaces this implementation without changing the call sites; the -`CostSource` discriminator tells the caller which one ran. +The reporting half is `ICostAttribution` (ADR 0012). The default elapsed-time proxy sums +`EXTRACT(EPOCH FROM (CompletedAt - ClaimedAt))` per kind and weights it by the configured +warehouse SKU's DBU/hour. The opt-in billing implementation instead returns DBUs and effective +list-price currency amounts from the Databricks system tables; the `CostSource` discriminator +tells the caller which one ran. Nothing here caps the *cost of one query*. A single operation against a large warehouse is bounded only by the run timeout — which cancels the run rather than merely abandoning it, so the timeout @@ -102,16 +101,23 @@ The tenant does not reach the compute in a form billing can see. `TenantScopedJo `lakewright_tenant_id` as a *job parameter*, and Databricks attributes usage in `system.billing.usage` by `custom_tags`, which come from the job or cluster definition, not from per-run parameters. Tagging per run is not available on `RunNow`, and a job per tenant does not -scale. So attribution has to go the other way: join `system.billing.usage` to a run id, and join -that run id to `operations.ExternalId`, which this project already stores. The tenant identity -lives in our database, not in theirs, and that is the correct place for it. - -Reading those tables needs a grant this project does not have. Querying `system.billing.usage` as +scale. So attribution goes the other way. PostgreSQL first selects `operations.ExternalId` for the +resolved tenant. The Databricks query receives only those job run ids as bound parameters and +also filters the configured `workspace_id`; application code joins the returned rows to operation +kinds. PostgreSQL is never named in Databricks SQL. A provider response containing any run id not +in the tenant-owned set fails the report instead of being attributed. The PostgreSQL selection is +index-backed and stops after 501 distinct runs; more than 500 returns HTTP 422, so one request can +issue at most one account-wide billing query. Pending statements also have an overall deadline and +are cancelled best-effort on deadline, transport failure, or caller cancellation. + +Reading those tables needs a grant the default development identity does not have. Querying `system.billing.usage` as the workspace identity returns `INSUFFICIENT_PERMISSIONS: User does not have USE SCHEMA on Schema -'system.billing'`. It is a metastore-admin grant, so the prerequisite for cost attribution in -currency is an administrative decision rather than a feature. The proxy is the alternative until -that grant is made; a product with the grant registers its own `ICostAttribution`, returns -`CostSource.Billing`, and the rest of the system reads it the same way. +'system.billing'`. It is an administrative decision, so the proxy remains the default. A product +with access to both `system.billing.usage` and `system.billing.list_prices` opts in with +`AddLakeWrightBillingCostAttribution`. The code path is covered locally with fixed-query, +window and price-boundary proration, malformed-row, correction, ownership, query-budget, polling +and cancellation tests; the system-table grants and live response shape remain workspace +verification steps. ### T6. Denial of service against the operation queue diff --git a/samples/Signalboard/Program.cs b/samples/Signalboard/Program.cs index 431b850..5693bd9 100644 --- a/samples/Signalboard/Program.cs +++ b/samples/Signalboard/Program.cs @@ -29,8 +29,8 @@ // Cost attribution: opt-in via configuration, like the Databricks clients. The sample ships the // elapsed-time proxy because it costs nothing to run and makes the cost endpoint answerable -// without a live workspace. A product wired to its own billing table replaces this registration -// with its own ICostAttribution. +// without a live workspace. A product with the system-table grants replaces this call with +// AddLakeWrightBillingCostAttribution; see docs/guides/billing-cost-attribution.md. builder.Services.AddLakeWrightCostAttribution(builder.Configuration); // OpenTelemetry: opt-in via configuration. The library publishes plain System.Diagnostics diff --git a/src/LakeWright.AspNetCore/CostEndpoints.cs b/src/LakeWright.AspNetCore/CostEndpoints.cs index 38bebeb..d779911 100644 --- a/src/LakeWright.AspNetCore/CostEndpoints.cs +++ b/src/LakeWright.AspNetCore/CostEndpoints.cs @@ -76,7 +76,38 @@ private static async Task GetCostAsync( }); } - var summary = await cost.ResolveAsync(tenant, effectiveFrom, effectiveUntil, cancellationToken); - return Results.Ok(summary); + try + { + var summary = await cost.ResolveAsync( + tenant, + effectiveFrom, + effectiveUntil, + cancellationToken); + return Results.Ok(summary); + } + catch (BillingUsageException exception) + { + if (exception.Code == "REPORT_TOO_LARGE") + { + return Results.Problem( + statusCode: StatusCodes.Status422UnprocessableEntity, + title: "Billing report is too large.", + detail: $"Narrow the window to at most {BillingUsageLimits.MaxJobRunsPerReport} distinct job runs.", + extensions: new Dictionary + { + ["code"] = exception.Code, + ["maxJobRuns"] = BillingUsageLimits.MaxJobRunsPerReport + }); + } + + return Results.Problem( + statusCode: StatusCodes.Status502BadGateway, + title: "Billing usage is unavailable.", + extensions: new Dictionary + { + ["code"] = exception.Code, + ["transient"] = exception.IsTransient + }); + } } } diff --git a/src/LakeWright.AspNetCore/LakeWrightServiceCollectionExtensions.cs b/src/LakeWright.AspNetCore/LakeWrightServiceCollectionExtensions.cs index 6cb6747..b04ccb6 100644 --- a/src/LakeWright.AspNetCore/LakeWrightServiceCollectionExtensions.cs +++ b/src/LakeWright.AspNetCore/LakeWrightServiceCollectionExtensions.cs @@ -1,4 +1,5 @@ using Azure.Core; +using LakeWright.Core.Cost; using LakeWright.Core.Jobs; using LakeWright.Core.Tenancy; using LakeWright.Databricks; @@ -160,6 +161,34 @@ public static IServiceCollection AddLakeWrightCostAttribution( return services; } + /// + /// Replaces proxy attribution with priced usage from Databricks billing system tables. + /// + /// + /// Call after and . The + /// workspace identity needs USE and SELECT access to + /// system.billing.usage and system.billing.list_prices. The configured workspace + /// id is always included in the billing query; job-run ids are first selected from the + /// tenant-filtered PostgreSQL operations table and correlated in application code. + /// + public static IServiceCollection AddLakeWrightBillingCostAttribution( + this IServiceCollection services, + IConfiguration configuration) + { + ArgumentNullException.ThrowIfNull(services); + ArgumentNullException.ThrowIfNull(configuration); + + services.AddOptions() + .Bind(configuration.GetSection(BillingUsageOptions.SectionName)) + .ValidateDataAnnotations() + .ValidateOnStart(); + + services.TryAddSingletonTimeProvider(); + services.AddScoped(); + services.AddScoped(); + return services; + } + private static void TryAddSingletonTimeProvider(this IServiceCollection services) { if (!services.Any(d => d.ServiceType == typeof(TimeProvider))) diff --git a/src/LakeWright.Core/Cost/IBillingUsageReader.cs b/src/LakeWright.Core/Cost/IBillingUsageReader.cs new file mode 100644 index 0000000..4aff0eb --- /dev/null +++ b/src/LakeWright.Core/Cost/IBillingUsageReader.cs @@ -0,0 +1,47 @@ +using LakeWright.Core.Tenancy; + +namespace LakeWright.Core.Cost; + +/// Reads priced Databricks billing rows for tenant-owned Lakeflow job runs. +/// +/// The tenant context and the run identifiers are both mandatory. Implementations read a global +/// system table, but the caller first resolves ownership in its transactional store and supplies +/// only identifiers owned by this tenant. This keeps the system-table escape explicit without +/// teaching the transactional layer about a Databricks client. +/// +public interface IBillingUsageReader +{ + Task> ReadAsync( + TenantContext tenant, + DateTimeOffset from, + DateTimeOffset until, + IReadOnlyCollection jobRunIds, + CancellationToken cancellationToken); +} + +/// Resource limits shared by billing attribution implementations. +public static class BillingUsageLimits +{ + /// + /// Maximum distinct Databricks job runs in one report. The billing reader performs one + /// account-wide system-table query, so callers must narrow the window when this limit is hit. + /// + public const int MaxJobRunsPerReport = 500; +} + +/// Priced usage attributed to one Lakeflow job run. +/// The Databricks job run identifier. +/// Net DBU quantity, including correction records. +/// Cost at the effective list price for one currency. +public sealed record BillingRunUsage( + long JobRunId, + decimal DbusConsumed, + CurrencyAmount EstimatedListCost); + +/// A safe, provider-neutral failure from a billing usage read. +public sealed class BillingUsageException(string code, bool isTransient) : Exception( + $"The billing usage read failed with code {code}.") +{ + public string Code { get; } = code; + public bool IsTransient { get; } = isTransient; +} diff --git a/src/LakeWright.Core/Cost/ICostAttribution.cs b/src/LakeWright.Core/Cost/ICostAttribution.cs index 0a43064..5ebec76 100644 --- a/src/LakeWright.Core/Cost/ICostAttribution.cs +++ b/src/LakeWright.Core/Cost/ICostAttribution.cs @@ -15,9 +15,8 @@ public enum CostSource Proxy = 0, /// - /// A direct read of system.billing.usage joined to operations.ExternalId. - /// Requires the grant this library does not hold, and is an extension point an adopter - /// implements against their own workspace. + /// A read of system.billing.usage for tenant-owned job run ids selected from + /// operations.ExternalId, correlated in application code. /// Billing = 1 } @@ -42,7 +41,20 @@ public sealed record TenantCostSummary( CostSource Source, string? WarehouseSku, decimal DbusConsumed, - IReadOnlyList ByKind); + IReadOnlyList ByKind) +{ + /// + /// Cost at Databricks' effective list price, grouped by the currency in + /// system.billing.list_prices. Empty for proxy attribution. + /// + /// + /// This is deliberately a collection rather than one amount: billing data can span price + /// rows in different currencies, and adding unlike currencies would produce a plausible but + /// meaningless number. It is an init-only property so the original positional constructor + /// remains source-compatible. + /// + public IReadOnlyList EstimatedListCost { get; init; } = []; +} /// /// One row of a . @@ -50,15 +62,25 @@ public sealed record TenantCostSummary( /// The operation's Kind string. /// Number of operations of this kind in the window. /// -/// Total wall-clock seconds these operations held compute, summed across the window. A proxy -/// number when ; the Databricks-reported duration otherwise. +/// Total wall-clock seconds these operations held compute, summed across the window. Populated by +/// and zero for a billing read, whose usage records do not report +/// the operation wall-clock duration. /// /// DBU attributed to this kind. public sealed record CostByKind( string Kind, int Operations, double ElapsedSeconds, - decimal DbusConsumed); + decimal DbusConsumed) +{ + /// Effective list-price cost for this operation kind, grouped by currency. + public IReadOnlyList EstimatedListCost { get; init; } = []; +} + +/// A monetary amount whose currency must travel with the value. +/// The ISO-style code reported by the Databricks price table. +/// The amount in . +public sealed record CurrencyAmount(string CurrencyCode, decimal Amount); /// /// Reports a tenant's Databricks compute consumption. @@ -71,10 +93,9 @@ public sealed record CostByKind( /// The shipped implementation in LakeWright.Multitenancy.Cost /// (OperationCostAttribution) weights operations.ClaimedAt to CompletedAt /// by the warehouse SKU's DBU/hour rate, which is what the threat model calls available without -/// a grant. reads system.billing.usage directly, which -/// is what a chargeback system wants but requires a metastore-admin grant this library does -/// not hold. An adopter wires whichever they have; the seam is the same interface so the rest -/// of the codebase does not know which one ran. +/// a grant. reads system.billing.usage and +/// system.billing.list_prices for tenant-owned job runs. It requires system-table grants, +/// so the proxy remains the default and an adopter opts into billing registration explicitly. /// public interface ICostAttribution { diff --git a/src/LakeWright.Databricks/BillingUsageOptions.cs b/src/LakeWright.Databricks/BillingUsageOptions.cs new file mode 100644 index 0000000..3906c9b --- /dev/null +++ b/src/LakeWright.Databricks/BillingUsageOptions.cs @@ -0,0 +1,23 @@ +using System.ComponentModel.DataAnnotations; + +namespace LakeWright.Databricks; + +/// Configuration for privileged reads of Databricks billing system tables. +public sealed class BillingUsageOptions +{ + public const string SectionName = "DatabricksBilling"; + + /// + /// Workspace whose usage may be correlated with this application's operation records. + /// + [Required] + public string WorkspaceId { get; set; } = string.Empty; + + /// Delay between Statement Execution polls after the initial wait expires. + [Range(50, 10_000)] + public int PollIntervalMilliseconds { get; set; } = 250; + + /// Overall deadline for a billing statement, including all polls. + [Range(1, 900)] + public int PollingTimeoutSeconds { get; set; } = 120; +} diff --git a/src/LakeWright.Databricks/DatabricksBillingUsageReader.cs b/src/LakeWright.Databricks/DatabricksBillingUsageReader.cs new file mode 100644 index 0000000..be7c552 --- /dev/null +++ b/src/LakeWright.Databricks/DatabricksBillingUsageReader.cs @@ -0,0 +1,284 @@ +using System.Globalization; +using LakeWright.Core.Cost; +using LakeWright.Core.Tenancy; +using Microsoft.Azure.Databricks.Client; +using Microsoft.Azure.Databricks.Client.Models; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Options; + +namespace LakeWright.Databricks; + +/// +/// Reads job-run usage and effective list-price cost from Databricks billing system tables. +/// +/// +/// The query is fixed text and every value is a Statement Execution parameter. Run identifiers +/// are capped and passed as one bound comma-delimited value; split turns it into an array in +/// Databricks SQL. The system table is account-wide, so workspace_id is an additional +/// mandatory bound filter rather than an assumption that run ids are globally unique. +/// +public sealed class DatabricksBillingUsageReader : IBillingUsageReader +{ + private const NumberStyles DecimalStyles = + NumberStyles.AllowLeadingSign | NumberStyles.AllowDecimalPoint | NumberStyles.AllowExponent; + + private const string BillingSql = + """ + WITH PricedUsage AS ( + SELECT CAST(u.usage_metadata.job_run_id AS STRING) AS JobRunId, + u.usage_unit AS UsageUnit, + u.usage_quantity + * CAST(timestampdiff( + MICROSECOND, + greatest(u.usage_start_time, :from, p.price_start_time), + least(u.usage_end_time, :until, coalesce(p.price_end_time, u.usage_end_time))) + AS DECIMAL(38, 12)) + / NULLIF(CAST(timestampdiff( + MICROSECOND, u.usage_start_time, u.usage_end_time) + AS DECIMAL(38, 12)), 0) AS WindowQuantity, + p.currency_code AS CurrencyCode, + p.pricing.effective_list.default AS EffectiveListPrice + FROM system.billing.usage u + JOIN system.billing.list_prices p + ON p.account_id = u.account_id + AND p.cloud = u.cloud + AND p.sku_name = u.sku_name + AND p.usage_unit = u.usage_unit + AND u.usage_end_time > p.price_start_time + AND (p.price_end_time IS NULL OR u.usage_start_time < p.price_end_time) + AND :until > p.price_start_time + AND (p.price_end_time IS NULL OR :from < p.price_end_time) + WHERE u.workspace_id = :workspace_id + AND u.usage_metadata.job_run_id IS NOT NULL + AND array_contains(split(:job_run_ids, ','), CAST(u.usage_metadata.job_run_id AS STRING)) + AND u.usage_start_time < :until + AND u.usage_end_time > :from + AND u.usage_date >= :from_date + AND u.usage_date <= :until_date + ) + SELECT JobRunId, + COALESCE(SUM(CASE WHEN UsageUnit = 'DBU' THEN WindowQuantity ELSE 0 END), 0) AS DbusConsumed, + CurrencyCode, + COALESCE(SUM(WindowQuantity * EffectiveListPrice), 0) AS EstimatedListCost + FROM PricedUsage + GROUP BY JobRunId, CurrencyCode + """; + + private readonly IDatabricksStatementSession _session; + private readonly DatabricksOptions _databricks; + private readonly BillingUsageOptions _billing; + private readonly TimeProvider _timeProvider; + + public DatabricksBillingUsageReader( + DatabricksClient client, + IOptions databricks, + IOptions billing, + ILogger logger, + TimeProvider timeProvider) + : this( + new DatabricksStatementSession(client, logger), + databricks.Value, + billing.Value, + timeProvider) + { + } + + internal DatabricksBillingUsageReader( + IDatabricksStatementSession session, + DatabricksOptions databricks, + BillingUsageOptions billing, + TimeProvider timeProvider) + { + _session = session; + _databricks = databricks; + _billing = billing; + _timeProvider = timeProvider; + } + + public async Task> ReadAsync( + TenantContext tenant, + DateTimeOffset from, + DateTimeOffset until, + IReadOnlyCollection jobRunIds, + CancellationToken cancellationToken) + { + ArgumentNullException.ThrowIfNull(tenant); + ArgumentNullException.ThrowIfNull(jobRunIds); + if (from >= until) + { + throw new ArgumentException("from must be earlier than until.", nameof(from)); + } + + if (jobRunIds.Count == 0) + { + return []; + } + + if (jobRunIds.Any(id => id <= 0)) + { + throw new ArgumentException("Job run ids must be positive.", nameof(jobRunIds)); + } + + var uniqueRunIds = jobRunIds.Distinct().Order().ToArray(); + if (uniqueRunIds.Length > BillingUsageLimits.MaxJobRunsPerReport) + { + throw new BillingUsageException("REPORT_TOO_LARGE", isTransient: false); + } + + var request = CreateRequest(from, until, uniqueRunIds); + var outcome = await _session.ExecuteAsync(request, tenant.TenantId, cancellationToken); + outcome = await WaitForCompletionAsync(tenant, outcome, cancellationToken); + return Parse(outcome); + } + + private SqlStatement CreateRequest( + DateTimeOffset from, + DateTimeOffset until, + IReadOnlyCollection jobRunIds) => new() + { + WarehouseId = _databricks.WarehouseId, + Catalog = "system", + Schema = "billing", + Statement = BillingSql, + Parameters = + [ + Parameter(StatementParameter.String("workspace_id", _billing.WorkspaceId)), + Parameter(StatementParameter.String( + "job_run_ids", + string.Join(',', jobRunIds.Select(id => id.ToString(CultureInfo.InvariantCulture))))), + Parameter(StatementParameter.Timestamp("from", from)), + Parameter(StatementParameter.Timestamp("until", until)), + Parameter(StatementParameter.Date("from_date", DateOnly.FromDateTime(from.UtcDateTime))), + Parameter(StatementParameter.Date("until_date", DateOnly.FromDateTime(until.UtcDateTime))) + ], + Disposition = SqlStatementDisposition.INLINE, + Format = StatementFormat.JSON_ARRAY, + RowLimit = 10_000, + WaitTimeout = _databricks.WaitTimeout, + OnWaitTimeout = SqlStatementOnWaitTimeout.CONTINUE + }; + + private async Task WaitForCompletionAsync( + TenantContext tenant, + StatementOutcome outcome, + CancellationToken cancellationToken) + { + string? activeStatementId = null; + var deadline = _timeProvider.GetUtcNow().AddSeconds(_billing.PollingTimeoutSeconds); + try + { + while (outcome is StatementOutcome.Pending pending) + { + activeStatementId = pending.StatementId; + var remaining = deadline - _timeProvider.GetUtcNow(); + if (remaining <= TimeSpan.Zero) + { + throw new BillingUsageException("POLL_TIMEOUT", isTransient: true); + } + + await Task.Delay( + TimeSpan.FromMilliseconds(_billing.PollIntervalMilliseconds) < remaining + ? TimeSpan.FromMilliseconds(_billing.PollIntervalMilliseconds) + : remaining, + _timeProvider, + cancellationToken); + outcome = await _session.GetAsync( + tenant.TenantId, + pending.StatementId, + cancellationToken); + } + + return outcome; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + await CancelBestEffortAsync(activeStatementId); + throw; + } + catch + { + await CancelBestEffortAsync(activeStatementId); + throw; + } + } + + private async Task CancelBestEffortAsync(string? statementId) + { + if (statementId is null) + { + return; + } + + using var cancelTimeout = new CancellationTokenSource(TimeSpan.FromSeconds(5)); + try + { + await _session.CancelAsync(statementId, cancelTimeout.Token); + } + catch (Exception exception) when (exception is not OutOfMemoryException) + { + // Cancellation must never replace the original timeout, transport error, or caller cancellation. + } + } + + private static List Parse(StatementOutcome outcome) + { + if (outcome is StatementOutcome.Failure failure) + { + throw new BillingUsageException(failure.ErrorCode, failure.IsTransient); + } + + if (outcome is not StatementOutcome.Success success) + { + throw new BillingUsageException("INLINE_RESULT_REQUIRED", isTransient: false); + } + + var indexes = success.ColumnNames + .Select((name, index) => (name, index)) + .ToDictionary(x => x.name, x => x.index, StringComparer.OrdinalIgnoreCase); + var required = new[] { "JobRunId", "DbusConsumed", "CurrencyCode", "EstimatedListCost" }; + if (required.Any(column => !indexes.ContainsKey(column))) + { + throw new BillingUsageException("INVALID_SCHEMA", isTransient: false); + } + + var rows = new List(success.Rows.Count); + foreach (var row in success.Rows) + { + if (!TryRead(row, indexes["JobRunId"], out var runText) + || !long.TryParse(runText, NumberStyles.None, CultureInfo.InvariantCulture, out var runId) + || runId <= 0 + || !TryRead(row, indexes["DbusConsumed"], out var dbuText) + || !decimal.TryParse(dbuText, DecimalStyles, CultureInfo.InvariantCulture, out var dbus) + || !TryRead(row, indexes["CurrencyCode"], out var currency) + || string.IsNullOrWhiteSpace(currency) + || !TryRead(row, indexes["EstimatedListCost"], out var costText) + || !decimal.TryParse(costText, DecimalStyles, CultureInfo.InvariantCulture, out var cost)) + { + throw new BillingUsageException("INVALID_ROW", isTransient: false); + } + + rows.Add(new BillingRunUsage( + runId, + dbus, + new CurrencyAmount(currency.ToUpperInvariant(), cost))); + } + + return rows; + } + + private static bool TryRead( + IReadOnlyList row, + int index, + [System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out string? value) + { + value = index < row.Count ? row[index] : null; + return value is not null; + } + + private static SqlStatementParameter Parameter(StatementParameter parameter) => new() + { + Name = parameter.Name, + Value = parameter.Value, + Type = parameter.Type + }; +} diff --git a/src/LakeWright.Databricks/DatabricksStatementExecutor.cs b/src/LakeWright.Databricks/DatabricksStatementExecutor.cs index 07687d0..136c8d8 100644 --- a/src/LakeWright.Databricks/DatabricksStatementExecutor.cs +++ b/src/LakeWright.Databricks/DatabricksStatementExecutor.cs @@ -1,4 +1,3 @@ -using Microsoft.Azure.Databricks.Client; using Microsoft.Azure.Databricks.Client.Models; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; @@ -9,35 +8,25 @@ namespace LakeWright.Databricks; /// Runs statements through Microsoft.Azure.Databricks.Client, translating its two failure /// modes into . /// -public sealed partial class DatabricksStatementExecutor : IStatementExecutor +public sealed class DatabricksStatementExecutor : IStatementExecutor { - // These log identifiers and codes, never free text. The client's exception message is the raw - // HTTP response body, and a Databricks error message can quote the value that caused it — a - // rejected parameter, a malformed literal — so both are tenant data. The threat model says - // logging never includes parameter values; an earlier version of these templates broke that. - // The detail still reaches the caller on StatementOutcome.Failure, where it is the caller's - // decision what to do with it. - [LoggerMessage(Level = LogLevel.Warning, Message = "Databricks rejected a statement request for tenant {TenantId} (HTTP {StatusCode})")] - private partial void LogRequestRejected(Core.Tenancy.TenantId? tenantId, int statusCode); - - [LoggerMessage(Level = LogLevel.Warning, Message = "Statement {StatementId} for tenant {TenantId} ended {State}: {ErrorCode}")] - private partial void LogStatementFailed(string? statementId, Core.Tenancy.TenantId? tenantId, StatementExecutionState? state, string errorCode); - - [LoggerMessage(Level = LogLevel.Warning, Message = "Statement {StatementId} returned unrecognised state {State}; treating as pending")] - private partial void LogUnrecognisedState(string? statementId, StatementExecutionState? state); - - private readonly DatabricksClient _client; + private readonly IDatabricksStatementSession _session; private readonly DatabricksOptions _options; - private readonly ILogger _logger; public DatabricksStatementExecutor( - DatabricksClient client, + Microsoft.Azure.Databricks.Client.DatabricksClient client, IOptions options, ILogger logger) + : this(new DatabricksStatementSession(client, logger), options.Value) { - _client = client; - _options = options.Value; - _logger = logger; + } + + internal DatabricksStatementExecutor( + IDatabricksStatementSession session, + DatabricksOptions options) + { + _session = session; + _options = options; } public async Task ExecuteAsync( @@ -82,19 +71,7 @@ public async Task ExecuteAsync( OnWaitTimeout = SqlStatementOnWaitTimeout.CONTINUE }; - StatementExecution response; - try - { - response = await _client.SQL.StatementExecution.Execute(request, cancellationToken); - } - catch (ClientApiException ex) - { - // The request itself was rejected: unknown warehouse, bad auth, malformed body. - LogRequestRejected(statement.Tenant.TenantId, (int)ex.StatusCode); - return new StatementOutcome.Failure("REQUEST_REJECTED", ex.Message, null, IsTransient: false); - } - - return Translate(response, statement.Tenant.TenantId); + return await _session.ExecuteAsync(request, statement.Tenant.TenantId, cancellationToken); } public async Task GetAsync( @@ -104,15 +81,7 @@ public async Task GetAsync( { ArgumentNullException.ThrowIfNull(tenant); - try - { - var response = await _client.SQL.StatementExecution.Get(statementId, cancellationToken); - return Translate(response, tenant.TenantId); - } - catch (ClientApiException ex) - { - return new StatementOutcome.Failure("REQUEST_REJECTED", ex.Message, statementId, IsTransient: false); - } + return await _session.GetAsync(tenant.TenantId, statementId, cancellationToken); } public Task CancelAsync( @@ -121,77 +90,6 @@ public Task CancelAsync( CancellationToken cancellationToken) { ArgumentNullException.ThrowIfNull(tenant); - return _client.SQL.StatementExecution.Cancel(statementId, cancellationToken); - } - - private StatementOutcome Translate(StatementExecution response, Core.Tenancy.TenantId? tenantId) - { - var state = response.Status?.State; - - switch (state) - { - case StatementExecutionState.SUCCEEDED: - return Succeeded(response); - - case StatementExecutionState.PENDING: - case StatementExecutionState.RUNNING: - return new StatementOutcome.Pending(response.StatementId); - - case StatementExecutionState.FAILED: - case StatementExecutionState.CLOSED: - case StatementExecutionState.CANCELED: - var error = response.Status?.Error; - var code = error?.ErrorCode.ToString() ?? state.ToString() ?? "UNKNOWN"; - LogStatementFailed(response.StatementId, tenantId, state, code); - return new StatementOutcome.Failure( - code, - error?.Message ?? $"Statement ended in state {state}.", - response.StatementId, - IsTransient: IsTransient(error?.ErrorCode)); - - default: - // Databricks reserves the right to add states. Treating an unrecognised one as a - // hard failure would turn a platform addition into an outage, so it is reported - // as still running and the caller polls. See ADR 0005. - LogUnrecognisedState(response.StatementId, state); - return new StatementOutcome.Pending(response.StatementId); - } - } - - private static StatementOutcome Succeeded(StatementExecution response) - { - var columns = response.Manifest?.Schema?.Columns?.Select(c => c.Name).ToArray() ?? []; - var totalRows = response.Manifest?.TotalRowCount ?? 0; - - // EXTERNAL_LINKS leaves DataArray null and puts the rows behind presigned URLs. Reporting - // that as a Success with an empty row list would be indistinguishable from a query that - // genuinely matched nothing. - var links = response.Result?.ExternalLinks?.ToArray() ?? []; - if (links.Length > 0) - { - return new StatementOutcome.LargeResult( - columns, - [.. links.Select(l => new Uri(l.ExternalLink))], - totalRows, - response.StatementId); - } - - var rows = response.Result?.DataArray? - .Select(IReadOnlyList (r) => r.ToArray()) - .ToArray() ?? []; - - return new StatementOutcome.Success(columns, rows, totalRows, response.StatementId); + return _session.CancelAsync(statementId, cancellationToken); } - - private static bool IsTransient(StatementExecutionErrorCode? code) => code switch - { - StatementExecutionErrorCode.TEMPORARILY_UNAVAILABLE => true, - StatementExecutionErrorCode.WORKSPACE_TEMPORARILY_UNAVAILABLE => true, - StatementExecutionErrorCode.SERVICE_UNDER_MAINTENANCE => true, - StatementExecutionErrorCode.RESOURCE_EXHAUSTED => true, - StatementExecutionErrorCode.INTERNAL_ERROR => true, - StatementExecutionErrorCode.IO_ERROR => true, - StatementExecutionErrorCode.ABORTED => true, - _ => false - }; } diff --git a/src/LakeWright.Databricks/DatabricksStatementSession.cs b/src/LakeWright.Databricks/DatabricksStatementSession.cs new file mode 100644 index 0000000..58be7c9 --- /dev/null +++ b/src/LakeWright.Databricks/DatabricksStatementSession.cs @@ -0,0 +1,146 @@ +using LakeWright.Core.Tenancy; +using Microsoft.Azure.Databricks.Client; +using Microsoft.Azure.Databricks.Client.Models; +using Microsoft.Extensions.Logging; + +namespace LakeWright.Databricks; + +internal interface IDatabricksStatementSession +{ + Task ExecuteAsync( + SqlStatement request, + TenantId tenantId, + CancellationToken cancellationToken); + + Task GetAsync( + TenantId tenantId, + string statementId, + CancellationToken cancellationToken); + + Task CancelAsync(string statementId, CancellationToken cancellationToken); +} + +/// +/// The shared Statement Execution API lifecycle used by tenant queries and privileged system +/// billing reads. +/// +internal sealed partial class DatabricksStatementSession( + DatabricksClient client, + ILogger logger) : IDatabricksStatementSession +{ + [LoggerMessage(Level = LogLevel.Warning, Message = "Databricks rejected a statement request for tenant {TenantId} (HTTP {StatusCode})")] + private static partial void LogRequestRejected(ILogger logger, TenantId tenantId, int statusCode); + + [LoggerMessage(Level = LogLevel.Warning, Message = "Statement {StatementId} for tenant {TenantId} ended {State}: {ErrorCode}")] + private static partial void LogStatementFailed(ILogger logger, string? statementId, TenantId tenantId, StatementExecutionState? state, string errorCode); + + [LoggerMessage(Level = LogLevel.Warning, Message = "Statement {StatementId} returned unrecognised state {State}; treating as pending")] + private static partial void LogUnrecognisedState(ILogger logger, string? statementId, StatementExecutionState? state); + + public async Task ExecuteAsync( + SqlStatement request, + TenantId tenantId, + CancellationToken cancellationToken) + { + try + { + var response = await client.SQL.StatementExecution.Execute(request, cancellationToken); + return Translate(response, tenantId); + } + catch (ClientApiException ex) + { + LogRequestRejected(logger, tenantId, (int)ex.StatusCode); + return new StatementOutcome.Failure("REQUEST_REJECTED", ex.Message, null, IsTransient: false) + { + StatusCode = ex.StatusCode + }; + } + } + + public async Task GetAsync( + TenantId tenantId, + string statementId, + CancellationToken cancellationToken) + { + try + { + var response = await client.SQL.StatementExecution.Get(statementId, cancellationToken); + return Translate(response, tenantId); + } + catch (ClientApiException ex) + { + return new StatementOutcome.Failure( + "REQUEST_REJECTED", ex.Message, statementId, IsTransient: false) + { + StatusCode = ex.StatusCode + }; + } + } + + public Task CancelAsync(string statementId, CancellationToken cancellationToken) => + client.SQL.StatementExecution.Cancel(statementId, cancellationToken); + + private StatementOutcome Translate(StatementExecution response, TenantId tenantId) + { + var state = response.Status?.State; + + switch (state) + { + case StatementExecutionState.SUCCEEDED: + return Succeeded(response); + + case StatementExecutionState.PENDING: + case StatementExecutionState.RUNNING: + return new StatementOutcome.Pending(response.StatementId); + + case StatementExecutionState.FAILED: + case StatementExecutionState.CLOSED: + case StatementExecutionState.CANCELED: + var error = response.Status?.Error; + var code = error?.ErrorCode.ToString() ?? state.ToString() ?? "UNKNOWN"; + LogStatementFailed(logger, response.StatementId, tenantId, state, code); + return new StatementOutcome.Failure( + code, + error?.Message ?? $"Statement ended in state {state}.", + response.StatementId, + IsTransient(error?.ErrorCode)); + + default: + LogUnrecognisedState(logger, response.StatementId, state); + return new StatementOutcome.Pending(response.StatementId); + } + } + + private static StatementOutcome Succeeded(StatementExecution response) + { + var columns = response.Manifest?.Schema?.Columns?.Select(c => c.Name).ToArray() ?? []; + var totalRows = response.Manifest?.TotalRowCount ?? 0; + var links = response.Result?.ExternalLinks?.ToArray() ?? []; + if (links.Length > 0) + { + return new StatementOutcome.LargeResult( + columns, + [.. links.Select(l => new Uri(l.ExternalLink))], + totalRows, + response.StatementId); + } + + var rows = response.Result?.DataArray? + .Select(IReadOnlyList (r) => r.ToArray()) + .ToArray() ?? []; + + return new StatementOutcome.Success(columns, rows, totalRows, response.StatementId); + } + + private static bool IsTransient(StatementExecutionErrorCode? code) => code switch + { + StatementExecutionErrorCode.TEMPORARILY_UNAVAILABLE => true, + StatementExecutionErrorCode.WORKSPACE_TEMPORARILY_UNAVAILABLE => true, + StatementExecutionErrorCode.SERVICE_UNDER_MAINTENANCE => true, + StatementExecutionErrorCode.RESOURCE_EXHAUSTED => true, + StatementExecutionErrorCode.INTERNAL_ERROR => true, + StatementExecutionErrorCode.IO_ERROR => true, + StatementExecutionErrorCode.ABORTED => true, + _ => false + }; +} diff --git a/src/LakeWright.Databricks/DatabricksTenantScopedExport.cs b/src/LakeWright.Databricks/DatabricksTenantScopedExport.cs index d044412..aecf362 100644 --- a/src/LakeWright.Databricks/DatabricksTenantScopedExport.cs +++ b/src/LakeWright.Databricks/DatabricksTenantScopedExport.cs @@ -32,13 +32,10 @@ namespace LakeWright.Databricks; /// public sealed partial class DatabricksTenantScopedExport : ITenantScopedExport { - [LoggerMessage(Level = LogLevel.Warning, Message = "Tenant-scoped export rejected: tenant {TenantId} (HTTP {StatusCode})")] - private partial void LogRequestRejected(TenantId? tenantId, int statusCode); - [LoggerMessage(Level = LogLevel.Warning, Message = "Tenant-scoped export failed: tenant {TenantId}, statement {StatementId}, code {ErrorCode}")] private partial void LogStatementFailed(TenantId? tenantId, string? statementId, string errorCode); - private readonly DatabricksClient _client; + private readonly IDatabricksStatementSession _session; private readonly DatabricksOptions _options; private readonly HttpClient _http; private readonly ILogger _logger; @@ -49,12 +46,24 @@ public DatabricksTenantScopedExport( HttpClient http, ILogger logger) { - _client = client; + _session = new DatabricksStatementSession(client, logger); _options = options.Value; _http = http; _logger = logger; } + internal DatabricksTenantScopedExport( + IDatabricksStatementSession session, + DatabricksOptions options, + HttpClient http, + ILogger logger) + { + _session = session; + _options = options; + _http = http; + _logger = logger; + } + public async IAsyncEnumerable StreamAsync( TenantScopedStatement statement, [EnumeratorCancellation] CancellationToken cancellationToken) @@ -85,58 +94,37 @@ public async IAsyncEnumerable StreamAsync( OnWaitTimeout = SqlStatementOnWaitTimeout.CONTINUE }; - StatementExecution response; - try - { - response = await _client.SQL.StatementExecution.Execute(request, cancellationToken).ConfigureAwait(false); - } - catch (ClientApiException ex) - { - LogRequestRejected(statement.Tenant.TenantId, (int)ex.StatusCode); - throw new HttpRequestException( - string.Create( - CultureInfo.InvariantCulture, - $"Databricks rejected the export for tenant {statement.Tenant.TenantId} (HTTP {(int)ex.StatusCode})."), - inner: ex, - statusCode: ex.StatusCode); - } + var outcome = await _session.ExecuteAsync( + request, + statement.Tenant.TenantId, + cancellationToken).ConfigureAwait(false); - if (response.Status is null - || response.Status.State is StatementExecutionState.FAILED - || response.Status.State is StatementExecutionState.CANCELED - || response.Status.State is StatementExecutionState.CLOSED) + if (outcome is StatementOutcome.Failure failure) { - var errorCode = response.Status?.Error?.ErrorCode.ToString() ?? "UNKNOWN"; - LogStatementFailed(statement.Tenant.TenantId, response.StatementId, errorCode); + LogStatementFailed(statement.Tenant.TenantId, failure.StatementId, failure.ErrorCode); throw new HttpRequestException( string.Create( CultureInfo.InvariantCulture, - $"Databricks ended the export in state {response.Status?.State} (code {errorCode})."), + $"Databricks rejected or failed the export (code {failure.ErrorCode})."), inner: null, - statusCode: null); + statusCode: failure.StatusCode); } - if (response.Status.State is StatementExecutionState.PENDING or StatementExecutionState.RUNNING) + if (outcome is StatementOutcome.Pending) { - // The executor already covered the polling path; the export is a single-call - // surface, so a not-yet-finished result is a programming error. Calling Get - // here would force the export to be aware of polling, which leaks the - // executor's state machine into a surface that is supposed to be a stream. throw new InvalidOperationException( "Databricks returned a still-running statement; the export is not a polling surface. " + "Use IStatementExecutor.ExecuteAsync and poll the returned statement id, then call " + "ITenantScopedExport.StreamAsync with a shorter statement or longer WaitTimeout."); } - if (response.Manifest is null || response.Result is null) + if (outcome is not StatementOutcome.LargeResult result) { throw new InvalidOperationException( - "Databricks returned a successful statement with no manifest or result."); + "Databricks export did not return an external-links result."); } - var columnNames = response.Manifest.Schema.Columns - .Select(c => c.Name) - .ToArray(); + var columnNames = result.ColumnNames.ToArray(); if (columnNames.Length == 0) { @@ -147,13 +135,10 @@ public async IAsyncEnumerable StreamAsync( // the header to write its column-name row, then write the values. yield return new ExportRow(new ExportColumn(columnNames), Array.Empty()); - var chunkLinks = response.Result.ExternalLinks - ?? throw new InvalidOperationException("EXTERNAL_LINKS disposition returned no chunk links."); - - foreach (var link in chunkLinks) + foreach (var link in result.Links) { cancellationToken.ThrowIfCancellationRequested(); - await foreach (var row in FetchChunkAsync(new Uri(link.ExternalLink), columnNames, cancellationToken).ConfigureAwait(false)) + await foreach (var row in FetchChunkAsync(link, columnNames, cancellationToken).ConfigureAwait(false)) { yield return row; } diff --git a/src/LakeWright.Databricks/LakeWright.Databricks.csproj b/src/LakeWright.Databricks/LakeWright.Databricks.csproj index 93519b0..fe16061 100644 --- a/src/LakeWright.Databricks/LakeWright.Databricks.csproj +++ b/src/LakeWright.Databricks/LakeWright.Databricks.csproj @@ -11,4 +11,7 @@ + + + diff --git a/src/LakeWright.Databricks/StatementOutcome.cs b/src/LakeWright.Databricks/StatementOutcome.cs index a932817..306a541 100644 --- a/src/LakeWright.Databricks/StatementOutcome.cs +++ b/src/LakeWright.Databricks/StatementOutcome.cs @@ -1,3 +1,5 @@ +using System.Net; + namespace LakeWright.Databricks; /// @@ -60,7 +62,11 @@ public sealed record Failure( string ErrorCode, string Message, string? StatementId, - bool IsTransient) : StatementOutcome; + bool IsTransient) : StatementOutcome + { + /// HTTP status for request-level failures, when the provider supplied one. + public HttpStatusCode? StatusCode { get; init; } + } /// /// The statement did not finish inside the wait timeout and is still running. The Statement diff --git a/src/LakeWright.Multitenancy/Cost/BillingCostAttribution.cs b/src/LakeWright.Multitenancy/Cost/BillingCostAttribution.cs new file mode 100644 index 0000000..0002ea3 --- /dev/null +++ b/src/LakeWright.Multitenancy/Cost/BillingCostAttribution.cs @@ -0,0 +1,144 @@ +using System.Globalization; +using LakeWright.Core.Cost; +using LakeWright.Core.Tenancy; +using Microsoft.EntityFrameworkCore; + +namespace LakeWright.Multitenancy.Cost; + +/// +/// Correlates tenant-owned operation records with priced Databricks job-run billing usage. +/// +public sealed class BillingCostAttribution( + LakeWrightDbContext db, + IBillingUsageReader billing) : ICostAttribution +{ + public async Task ResolveAsync( + TenantContext tenant, + DateTimeOffset from, + DateTimeOffset until, + CancellationToken cancellationToken) + { + ArgumentNullException.ThrowIfNull(tenant); + if (from >= until) + { + throw new ArgumentException("from must be earlier than until.", nameof(from)); + } + + var candidates = await db.Operations + .AsNoTracking() + .Where(operation => operation.OrganizationId == tenant.TenantId + && operation.ExternalId != null + && operation.ClaimedAt != null + && operation.CompletedAt != null + && operation.CompletedAt > from + && operation.ClaimedAt < until) + .GroupBy(operation => operation.ExternalId) + .Select(group => new + { + ExternalId = group.Key, + FirstKind = group.Min(operation => operation.Kind), + LastKind = group.Max(operation => operation.Kind) + }) + .Take(BillingUsageLimits.MaxJobRunsPerReport + 1) + .ToListAsync(cancellationToken); + + if (candidates.Count > BillingUsageLimits.MaxJobRunsPerReport) + { + throw new BillingUsageException("REPORT_TOO_LARGE", isTransient: false); + } + + var ownedRuns = new Dictionary(); + foreach (var candidate in candidates) + { + var kind = candidate.FirstKind; + if (kind is null + || candidate.LastKind is null + || !string.Equals(kind, candidate.LastKind, StringComparison.Ordinal)) + { + throw new BillingUsageException("AMBIGUOUS_RUN", isTransient: false); + } + + if (!long.TryParse( + candidate.ExternalId, + NumberStyles.None, + CultureInfo.InvariantCulture, + out var runId) + || runId <= 0) + { + throw new BillingUsageException("INVALID_OPERATION_RUN_ID", isTransient: false); + } + + if (ownedRuns.TryGetValue(runId, out var existingKind) + && !string.Equals(existingKind, kind, StringComparison.Ordinal)) + { + throw new BillingUsageException("AMBIGUOUS_RUN", isTransient: false); + } + + ownedRuns[runId] = kind; + } + + if (ownedRuns.Count == 0) + { + return Empty(tenant, from, until); + } + + var usage = await billing.ReadAsync( + tenant, + from, + until, + ownedRuns.Keys, + cancellationToken); + + if (usage.FirstOrDefault(row => !ownedRuns.ContainsKey(row.JobRunId)) is not null) + { + throw new BillingUsageException("UNEXPECTED_RUN", isTransient: false); + } + + var byKind = usage + .GroupBy(row => ownedRuns[row.JobRunId], StringComparer.Ordinal) + .Select(group => new CostByKind( + group.Key, + group.Select(row => row.JobRunId).Distinct().Count(), + ElapsedSeconds: 0, + DbusConsumed: Math.Round(group.Sum(row => row.DbusConsumed), 4)) + { + EstimatedListCost = CurrencyTotals(group) + }) + .OrderByDescending(row => row.DbusConsumed) + .ThenBy(row => row.Kind, StringComparer.Ordinal) + .ToArray(); + + return new TenantCostSummary( + tenant.TenantId, + from, + until, + CostSource.Billing, + WarehouseSku: null, + DbusConsumed: Math.Round(usage.Sum(row => row.DbusConsumed), 4), + byKind) + { + EstimatedListCost = CurrencyTotals(usage) + }; + } + + private static TenantCostSummary Empty( + TenantContext tenant, + DateTimeOffset from, + DateTimeOffset until) => new( + tenant.TenantId, + from, + until, + CostSource.Billing, + WarehouseSku: null, + DbusConsumed: 0, + ByKind: []); + + private static CurrencyAmount[] CurrencyTotals(IEnumerable usage) => + usage + .GroupBy(row => row.EstimatedListCost.CurrencyCode, StringComparer.Ordinal) + .Select(group => new CurrencyAmount( + group.Key, + Math.Round(group.Sum(row => row.EstimatedListCost.Amount), 4))) + .OrderBy(amount => amount.CurrencyCode, StringComparer.Ordinal) + .ToArray(); +} diff --git a/src/LakeWright.Multitenancy/LakeWrightDbContext.cs b/src/LakeWright.Multitenancy/LakeWrightDbContext.cs index b4cad44..4741bef 100644 --- a/src/LakeWright.Multitenancy/LakeWrightDbContext.cs +++ b/src/LakeWright.Multitenancy/LakeWrightDbContext.cs @@ -92,6 +92,11 @@ protected override void OnModelCreating(ModelBuilder modelBuilder) // which is the trailing column here, so it reads straight from the index. e.HasIndex(x => new { x.State, x.ClaimedAt }); + e.HasIndex(x => new { x.OrganizationId, x.CompletedAt }) + .IncludeProperties(x => new { x.ClaimedAt, x.Kind, x.ExternalId }) + .HasFilter("\"ExternalId\" IS NOT NULL AND \"ClaimedAt\" IS NOT NULL AND \"CompletedAt\" IS NOT NULL") + .HasDatabaseName("IX_operations_billing_window"); + // The claim query orders by CreatedAt, which the index above does not cover, so // Postgres sorted every pending row on every claim. Measured against Postgres 17: // 25.6ms at a 50,000-row backlog, and 141.7ms at 300,000 with the sort spilling 13.5MB diff --git a/tests/LakeWright.TenantIsolation.Tests/BillingCostAttributionTests.cs b/tests/LakeWright.TenantIsolation.Tests/BillingCostAttributionTests.cs new file mode 100644 index 0000000..7ef329e --- /dev/null +++ b/tests/LakeWright.TenantIsolation.Tests/BillingCostAttributionTests.cs @@ -0,0 +1,590 @@ +using LakeWright.Core.Cost; +using LakeWright.Core.Tenancy; +using LakeWright.Databricks; +using LakeWright.Multitenancy; +using LakeWright.Multitenancy.Cost; +using LakeWright.Multitenancy.Model; +using Microsoft.Azure.Databricks.Client.Models; +using Microsoft.Extensions.Time.Testing; +using NSubstitute; +using static LakeWright.TenantIsolation.Tests.TestApi; + +namespace LakeWright.TenantIsolation.Tests; + +[Trait("Category", "TenantIsolation")] +public class DatabricksBillingUsageReaderTests +{ + private static readonly DateTimeOffset From = + DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + private static readonly DateTimeOffset Until = + DateTimeOffset.Parse("2026-09-01T00:00:00Z", null); + + [Fact] + public async Task ReadAsync_binds_workspace_run_and_window_values() + { + var session = new StubStatementSession(Success([])); + var reader = Reader(session); + + var rows = await reader.ReadAsync( + Acme(), + From, + Until, + [22, 11], + TestContext.Current.CancellationToken); + + rows.ShouldBeEmpty(); + var request = session.Requests.Single(); + request.Catalog.ShouldBe("system"); + request.Schema.ShouldBe("billing"); + request.Disposition.ShouldBe(SqlStatementDisposition.INLINE); + request.Statement.ShouldContain("u.workspace_id = :workspace_id"); + request.Statement.ShouldContain("u.usage_metadata.job_run_id"); + request.Statement.ShouldContain("p.pricing.effective_list.default"); + request.Statement.ShouldContain("u.usage_date >= :from_date"); + request.Statement.ShouldContain("u.usage_date <= :until_date"); + request.Statement.ShouldNotContain("workspace-123"); + request.Statement.ShouldNotContain("11,22"); + + Parameters(request)["workspace_id"].ShouldBe("workspace-123"); + Parameters(request)["job_run_ids"].ShouldBe("11,22"); + Parameters(request)["from_date"].ShouldBe("2026-08-01"); + Parameters(request)["until_date"].ShouldBe("2026-09-01"); + } + + [Fact] + public async Task ReadAsync_rejects_a_report_that_would_require_repeated_system_table_scans() + { + var session = new StubStatementSession(Success([])); + var reader = Reader(session); + + var exception = await Should.ThrowAsync(async () => + await reader.ReadAsync( + Acme(), + From, + Until, + Enumerable.Range(1, BillingUsageLimits.MaxJobRunsPerReport + 1) + .Select(value => (long)value) + .ToArray(), + TestContext.Current.CancellationToken)); + + exception.Code.ShouldBe("REPORT_TOO_LARGE"); + session.Requests.ShouldBeEmpty(); + } + + [Fact] + public async Task ReadAsync_prorates_the_same_quantity_at_report_and_price_boundaries() + { + var session = new StubStatementSession(Success([])); + + await Reader(session).ReadAsync( + Acme(), From, Until, [11], TestContext.Current.CancellationToken); + + var sql = session.Requests.Single().Statement; + sql.ShouldContain("greatest(u.usage_start_time, :from, p.price_start_time)"); + sql.ShouldContain("least(u.usage_end_time, :until, coalesce(p.price_end_time, u.usage_end_time))"); + sql.ShouldContain("u.usage_end_time > p.price_start_time"); + sql.ShouldContain("u.usage_start_time < p.price_end_time"); + sql.ShouldContain(":until > p.price_start_time"); + sql.ShouldContain(":from < p.price_end_time"); + sql.ShouldContain("THEN WindowQuantity ELSE 0 END"); + sql.ShouldContain("WindowQuantity * EffectiveListPrice"); + sql.Split("WindowQuantity * EffectiveListPrice").Length.ShouldBe(2); + + // A quantity of 8 over 00:00-04:00 contributes 2 to each of the two price intervals + // intersecting a 01:00-03:00 report. The CTE emits one row per interval and both the DBU + // and price sums consume that same apportioned quantity: 2 + 2, never 8 or 16. + var quantity = 8m; + var usageDuration = TimeSpan.FromHours(4); + var firstPriceOverlap = TimeSpan.FromHours(1); + var secondPriceOverlap = TimeSpan.FromHours(1); + var expectedWindowQuantity = quantity * firstPriceOverlap.Ticks / usageDuration.Ticks + + quantity * secondPriceOverlap.Ticks / usageDuration.Ticks; + expectedWindowQuantity.ShouldBe(4m); + } + + [Fact] + public async Task ReadAsync_polls_a_pending_statement_to_completion() + { + var session = new StubStatementSession( + new StatementOutcome.Pending("statement-1"), + Success([["11", "2.5", "USD", "0.75"]])); + var reader = Reader(session); + + var rows = await reader.ReadAsync( + Acme(), + From, + Until, + [11], + TestContext.Current.CancellationToken); + + session.PolledStatementIds.ShouldBe(["statement-1"]); + rows.ShouldBe([ + new BillingRunUsage(11, 2.5m, new CurrencyAmount("USD", 0.75m)) + ]); + } + + [Fact] + public async Task ReadAsync_preserves_net_correction_values_and_normalizes_currency() + { + var session = new StubStatementSession(Success([ + ["11", "-1.2500", "usd", "-0.3125"] + ])); + + var rows = await Reader(session).ReadAsync( + Acme(), + From, + Until, + [11], + TestContext.Current.CancellationToken); + + rows.Single().ShouldBe( + new BillingRunUsage(11, -1.25m, new CurrencyAmount("USD", -0.3125m))); + session.Requests.Single().Statement.ShouldContain("u.usage_quantity"); + } + + [Theory] + [InlineData("not-a-run", "2.5", "USD", "0.75")] + [InlineData("11", "2,5", "USD", "0.75")] + [InlineData("11", "2.5", "", "0.75")] + [InlineData("11", "2.5", "USD", "not-money")] + public async Task ReadAsync_rejects_malformed_billing_rows( + string runId, + string dbus, + string currency, + string cost) + { + var session = new StubStatementSession(Success([[runId, dbus, currency, cost]])); + + var exception = await Should.ThrowAsync(async () => + await Reader(session).ReadAsync( + Acme(), + From, + Until, + [11], + TestContext.Current.CancellationToken)); + + exception.Code.ShouldBe("INVALID_ROW"); + } + + [Fact] + public async Task ReadAsync_surfaces_statement_failure_without_the_provider_message() + { + var session = new StubStatementSession(new StatementOutcome.Failure( + "PERMISSION_DENIED", + "message may contain tenant data", + "statement-1", + IsTransient: false)); + + var exception = await Should.ThrowAsync(async () => + await Reader(session).ReadAsync( + Acme(), + From, + Until, + [11], + TestContext.Current.CancellationToken)); + + exception.Code.ShouldBe("PERMISSION_DENIED"); + exception.Message.ShouldNotContain("tenant data"); + } + + [Fact] + public async Task ReadAsync_cancels_a_pending_statement_when_the_caller_cancels() + { + var session = new StubStatementSession(new StatementOutcome.Pending("statement-1")); + using var cancellation = new CancellationTokenSource(); + + var read = Reader(session).ReadAsync(Acme(), From, Until, [11], cancellation.Token); + await Task.Delay(10, TestContext.Current.CancellationToken); + cancellation.Cancel(); + + await Should.ThrowAsync(async () => await read); + session.CancelledStatementIds.ShouldBe(["statement-1"]); + } + + [Fact] + public async Task ReadAsync_keeps_caller_cancellation_when_best_effort_cancel_fails() + { + var session = new StubStatementSession(new StatementOutcome.Pending("statement-1")) + { + CancelException = new HttpRequestException("cancel transport failed") + }; + using var cancellation = new CancellationTokenSource(); + + var read = Reader(session).ReadAsync(Acme(), From, Until, [11], cancellation.Token); + await Task.Delay(10, TestContext.Current.CancellationToken); + cancellation.Cancel(); + + await Should.ThrowAsync(async () => await read); + session.CancelledStatementIds.ShouldBe(["statement-1"]); + } + + [Fact] + public async Task ReadAsync_cancels_a_pending_statement_at_the_overall_deadline() + { + var time = new FakeTimeProvider(From); + var session = new StubStatementSession( + new StatementOutcome.Pending("statement-1"), + new StatementOutcome.Pending("statement-1")); + var read = Reader(session, time, pollingTimeoutSeconds: 1) + .ReadAsync(Acme(), From, Until, [11], TestContext.Current.CancellationToken); + await Task.Yield(); + + time.Advance(TimeSpan.FromSeconds(1)); + + var exception = await Should.ThrowAsync(async () => await read); + exception.Code.ShouldBe("POLL_TIMEOUT"); + exception.IsTransient.ShouldBeTrue(); + session.CancelledStatementIds.ShouldBe(["statement-1"]); + } + + [Fact] + public async Task ReadAsync_cancels_a_pending_statement_when_poll_transport_fails() + { + var session = new StubStatementSession(new StatementOutcome.Pending("statement-1")) + { + GetException = new HttpRequestException("poll transport failed") + }; + + var exception = await Should.ThrowAsync(async () => + await Reader(session).ReadAsync( + Acme(), From, Until, [11], TestContext.Current.CancellationToken)); + + exception.Message.ShouldBe("poll transport failed"); + session.CancelledStatementIds.ShouldBe(["statement-1"]); + } + + private static DatabricksBillingUsageReader Reader( + IDatabricksStatementSession session, + TimeProvider? timeProvider = null, + int pollingTimeoutSeconds = 120) => new( + session, + new DatabricksOptions { WarehouseId = "warehouse-1", WorkspaceUrl = "https://example" }, + new BillingUsageOptions + { + WorkspaceId = "workspace-123", + PollIntervalMilliseconds = 50, + PollingTimeoutSeconds = pollingTimeoutSeconds + }, + timeProvider ?? TimeProvider.System); + + private static TenantContext Acme() => TenantContextFactory.ForTenant(AcmeId, "analytics"); + + private static Dictionary Parameters(SqlStatement request) => + request.Parameters.ToDictionary(parameter => parameter.Name, parameter => parameter.Value); + + private static StatementOutcome.Success Success(IReadOnlyList> rows) => + new( + ["JobRunId", "DbusConsumed", "CurrencyCode", "EstimatedListCost"], + rows, + rows.Count, + "statement-1"); + + private sealed class StubStatementSession(params StatementOutcome[] outcomes) + : IDatabricksStatementSession + { + private readonly Queue _outcomes = new(outcomes); + + public List Requests { get; } = []; + public List PolledStatementIds { get; } = []; + public List CancelledStatementIds { get; } = []; + public Exception? GetException { get; init; } + public Exception? CancelException { get; init; } + + public Task ExecuteAsync( + SqlStatement request, + TenantId tenantId, + CancellationToken cancellationToken) + { + Requests.Add(request); + return Task.FromResult(_outcomes.Dequeue()); + } + + public Task GetAsync( + TenantId tenantId, + string statementId, + CancellationToken cancellationToken) + { + PolledStatementIds.Add(statementId); + if (GetException is not null) + { + throw GetException; + } + return Task.FromResult(_outcomes.Dequeue()); + } + + public Task CancelAsync(string statementId, CancellationToken cancellationToken) + { + CancelledStatementIds.Add(statementId); + if (CancelException is not null) + { + throw CancelException; + } + return Task.CompletedTask; + } + } +} + +[Trait("Category", "TenantIsolation")] +[Collection(nameof(PostgresTests))] +public class BillingCostAttributionTests(PostgresFixture postgres) +{ + [Fact] + public async Task ResolveAsync_correlates_in_application_and_counts_distinct_runs() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + var until = from.AddDays(1); + await SeedAsync(db, from, cancellationToken); + + var billing = Substitute.For(); + billing.ReadAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any()) + .Returns([ + new BillingRunUsage(101, 2m, new CurrencyAmount("USD", 0.50m)), + new BillingRunUsage(101, -0.5m, new CurrencyAmount("USD", -0.125m)), + new BillingRunUsage(102, 1m, new CurrencyAmount("USD", 0.25m)) + ]); + + var summary = await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), + from, + until, + cancellationToken); + + summary.Source.ShouldBe(CostSource.Billing); + summary.DbusConsumed.ShouldBe(2.5m); + summary.EstimatedListCost.ShouldBe([new CurrencyAmount("USD", 0.625m)]); + summary.ByKind.Single().Operations.ShouldBe(2); + summary.ByKind.Single().EstimatedListCost.ShouldBe([ + new CurrencyAmount("USD", 0.625m) + ]); + + await billing.Received(1).ReadAsync( + Arg.Is(context => context.TenantId == AcmeId), + from, + until, + Arg.Is>(ids => + ids.Order().SequenceEqual(new long[] { 101, 102 })), + cancellationToken); + } + + [Fact] + public async Task ResolveAsync_refuses_a_run_not_owned_by_the_tenant() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + var until = from.AddDays(1); + await SeedAsync(db, from, cancellationToken); + + var billing = Substitute.For(); + billing.ReadAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any()) + .Returns([ + new BillingRunUsage(999, 10m, new CurrencyAmount("USD", 2.5m)) + ]); + + var exception = await Should.ThrowAsync(async () => + await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), + from, + until, + cancellationToken)); + + exception.Code.ShouldBe("UNEXPECTED_RUN"); + } + + [Fact] + public async Task ResolveAsync_refuses_a_malformed_stored_run_id() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + var until = from.AddDays(1); + db.Organizations.Add(Organization(AcmeId, "Acme", "acme", from)); + db.Operations.Add(Operation(AcmeId, "analysis", "not-a-run", from)); + await db.SaveChangesAsync(cancellationToken); + + var billing = Substitute.For(); + var exception = await Should.ThrowAsync(async () => + await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), + from, + until, + cancellationToken)); + + exception.Code.ShouldBe("INVALID_OPERATION_RUN_ID"); + await billing.DidNotReceiveWithAnyArgs().ReadAsync( + default!, default, default, default!, cancellationToken); + } + + [Fact] + public async Task ResolveAsync_rejects_more_runs_than_one_billing_query_can_bound() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + db.Organizations.Add(Organization(AcmeId, "Acme", "acme", from)); + db.Operations.AddRange(Enumerable + .Range(1, BillingUsageLimits.MaxJobRunsPerReport + 1) + .Select(runId => Operation(AcmeId, "analysis", runId.ToString(), from))); + await db.SaveChangesAsync(cancellationToken); + var billing = Substitute.For(); + + var exception = await Should.ThrowAsync(async () => + await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), from, from.AddDays(1), cancellationToken)); + + exception.Code.ShouldBe("REPORT_TOO_LARGE"); + await billing.DidNotReceiveWithAnyArgs().ReadAsync( + default!, default, default, default!, cancellationToken); + } + + [Fact] + public async Task ResolveAsync_bounds_distinct_run_ids_not_duplicate_operation_rows() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + db.Organizations.Add(Organization(AcmeId, "Acme", "acme", from)); + db.Operations.AddRange(Enumerable + .Range(1, BillingUsageLimits.MaxJobRunsPerReport) + .Select(runId => Operation(AcmeId, "analysis", runId.ToString(), from))); + db.Operations.AddRange(Enumerable + .Range(0, 10) + .Select(_ => Operation(AcmeId, "analysis", "1", from))); + await db.SaveChangesAsync(cancellationToken); + var billing = Substitute.For(); + billing.ReadAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any()) + .Returns([]); + + await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), from, from.AddDays(1), cancellationToken); + + await billing.Received(1).ReadAsync( + Arg.Any(), + from, + from.AddDays(1), + Arg.Is>(ids => + ids.Count == BillingUsageLimits.MaxJobRunsPerReport), + cancellationToken); + } + + [Fact] + public async Task ResolveAsync_rejects_one_run_recorded_for_conflicting_kinds() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + db.Organizations.Add(Organization(AcmeId, "Acme", "acme", from)); + db.Operations.AddRange( + Operation(AcmeId, "analysis", "101", from), + Operation(AcmeId, "export", "101", from)); + await db.SaveChangesAsync(cancellationToken); + var billing = Substitute.For(); + + var exception = await Should.ThrowAsync(async () => + await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), from, from.AddDays(1), cancellationToken)); + + exception.Code.ShouldBe("AMBIGUOUS_RUN"); + await billing.DidNotReceiveWithAnyArgs().ReadAsync( + default!, default, default, default!, cancellationToken); + } + + [Fact] + public async Task ResolveAsync_orders_kinds_by_dbus_without_adding_unlike_currencies() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + var from = DateTimeOffset.Parse("2026-08-01T00:00:00Z", null); + db.Organizations.Add(Organization(AcmeId, "Acme", "acme", from)); + db.Operations.AddRange( + Operation(AcmeId, "more-dbus", "101", from), + Operation(AcmeId, "more-money", "102", from)); + await db.SaveChangesAsync(cancellationToken); + var billing = Substitute.For(); + billing.ReadAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any()) + .Returns([ + new BillingRunUsage(101, 10m, new CurrencyAmount("EUR", 1m)), + new BillingRunUsage(102, 2m, new CurrencyAmount("USD", 999m)) + ]); + + var summary = await new BillingCostAttribution(db, billing).ResolveAsync( + Acme(), from, from.AddDays(1), cancellationToken); + + summary.ByKind.Select(row => row.Kind).ShouldBe(["more-dbus", "more-money"]); + summary.EstimatedListCost.ShouldBe([ + new CurrencyAmount("EUR", 1m), + new CurrencyAmount("USD", 999m) + ]); + } + + private static TenantContext Acme() => TenantContextFactory.ForTenant(AcmeId, "analytics"); + + private static async Task SeedAsync( + LakeWrightDbContext db, + DateTimeOffset now, + CancellationToken cancellationToken) + { + db.Organizations.AddRange( + Organization(AcmeId, "Acme", "acme", now), + Organization(GlobexId, "Globex", "globex", now)); + db.Operations.AddRange( + Operation(AcmeId, "analysis", "101", now), + Operation(AcmeId, "analysis", "102", now), + Operation(GlobexId, "other-tenant", "999", now)); + await db.SaveChangesAsync(cancellationToken); + } + + private static Organization Organization( + TenantId id, + string name, + string slug, + DateTimeOffset now) => new() + { + Id = id, + Name = name, + Slug = slug, + CreatedAt = now, + Schema = UnityCatalogIdentifier.SchemaForTenant(id), + State = OrganizationState.Active + }; + + private static Operation Operation( + TenantId tenantId, + string kind, + string externalId, + DateTimeOffset now) => new() + { + Id = Guid.CreateVersion7(), + OrganizationId = tenantId, + PrincipalId = Alice, + Kind = kind, + State = OperationState.Succeeded, + ExternalId = externalId, + IdempotencyKey = Guid.CreateVersion7().ToString("N"), + CreatedAt = now, + ClaimedAt = now, + CompletedAt = now.AddHours(1) + }; +} diff --git a/tests/LakeWright.TenantIsolation.Tests/CostEndpointTests.cs b/tests/LakeWright.TenantIsolation.Tests/CostEndpointTests.cs index 35d50bd..f009f53 100644 --- a/tests/LakeWright.TenantIsolation.Tests/CostEndpointTests.cs +++ b/tests/LakeWright.TenantIsolation.Tests/CostEndpointTests.cs @@ -3,7 +3,9 @@ using LakeWright.Core.Cost; using LakeWright.Multitenancy; using LakeWright.Multitenancy.Model; +using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Options; +using NSubstitute; using static LakeWright.TenantIsolation.Tests.TestApi; namespace LakeWright.TenantIsolation.Tests; @@ -135,4 +137,60 @@ public async Task A_window_ending_in_the_distant_future_is_rejected_with_400() // Assert response.StatusCode.ShouldBe(HttpStatusCode.BadRequest); } + + [Fact] + public async Task A_billing_provider_failure_answers_502_without_the_provider_message() + { + var cancellationToken = TestContext.Current.CancellationToken; + var cost = Substitute.For(); + cost.ResolveAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns>(_ => throw new BillingUsageException( + "PERMISSION_DENIED", + isTransient: false)); + var (host, client) = await StartAsync( + postgres, + services => services.AddScoped(_ => cost)); + using var _h = host; + + var response = await client.SendAsync( + As(Vera, HttpMethod.Get, $"/organizations/{AcmeId.Value}/cost"), + cancellationToken); + var body = await response.Content.ReadAsStringAsync(cancellationToken); + + response.StatusCode.ShouldBe(HttpStatusCode.BadGateway); + body.ShouldContain("PERMISSION_DENIED"); + body.ShouldNotContain("system.billing.usage"); + } + + [Fact] + public async Task An_oversized_billing_report_answers_422_with_the_enforced_run_limit() + { + var cancellationToken = TestContext.Current.CancellationToken; + var cost = Substitute.For(); + cost.ResolveAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns>(_ => throw new BillingUsageException( + "REPORT_TOO_LARGE", + isTransient: false)); + var (host, client) = await StartAsync( + postgres, + services => services.AddScoped(_ => cost)); + using var _h = host; + + var response = await client.SendAsync( + As(Vera, HttpMethod.Get, $"/organizations/{AcmeId.Value}/cost"), + cancellationToken); + var body = await response.Content.ReadAsStringAsync(cancellationToken); + + response.StatusCode.ShouldBe(HttpStatusCode.UnprocessableEntity); + body.ShouldContain("REPORT_TOO_LARGE"); + body.ShouldContain(BillingUsageLimits.MaxJobRunsPerReport.ToString()); + } } diff --git a/tests/LakeWright.TenantIsolation.Tests/OperationIndexTests.cs b/tests/LakeWright.TenantIsolation.Tests/OperationIndexTests.cs index 2616332..88afab5 100644 --- a/tests/LakeWright.TenantIsolation.Tests/OperationIndexTests.cs +++ b/tests/LakeWright.TenantIsolation.Tests/OperationIndexTests.cs @@ -38,4 +38,23 @@ public async Task The_claim_query_ordering_is_backed_by_a_partial_index() && i.Contains("WHERE", StringComparison.Ordinal), "the claim query orders by CreatedAt and sorts every pending row without this"); } + + [Fact] + public async Task The_billing_window_query_is_backed_by_a_tenant_time_index() + { + var cancellationToken = TestContext.Current.CancellationToken; + await using var db = await postgres.NewDatabaseAsync(); + + var indexes = await db.Database + .SqlQuery($"SELECT indexdef AS \"Value\" FROM pg_indexes WHERE tablename = 'operations'") + .ToListAsync(cancellationToken); + + indexes.ShouldContain( + index => index.Contains("IX_operations_billing_window", StringComparison.Ordinal) + && index.Contains("OrganizationId", StringComparison.Ordinal) + && index.Contains("CompletedAt", StringComparison.Ordinal) + && index.Contains("ExternalId", StringComparison.Ordinal) + && index.Contains("WHERE", StringComparison.Ordinal), + "billing attribution must not scan an unbounded tenant operation history"); + } } diff --git a/tests/LakeWright.TenantIsolation.Tests/OptionalDatabricksTests.cs b/tests/LakeWright.TenantIsolation.Tests/OptionalDatabricksTests.cs index 7a5e162..68920af 100644 --- a/tests/LakeWright.TenantIsolation.Tests/OptionalDatabricksTests.cs +++ b/tests/LakeWright.TenantIsolation.Tests/OptionalDatabricksTests.cs @@ -1,5 +1,7 @@ using LakeWright.AspNetCore; +using LakeWright.Core.Cost; using LakeWright.Databricks; +using LakeWright.Multitenancy.Cost; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Options; @@ -70,6 +72,36 @@ public void Databricks_accepts_a_complete_section() validate.ShouldNotThrow(); } + [Fact] + public void Billing_cost_rejects_a_missing_workspace_id() + { + var services = new ServiceCollection(); + services.AddLakeWrightBillingCostAttribution(Configuration([])); + + var validate = Validator(services); + + validate.ShouldThrow() + .Message.ShouldContain(nameof(BillingUsageOptions.WorkspaceId)); + } + + [Fact] + public void Billing_cost_registration_composes_the_reader_and_replaces_the_proxy() + { + var services = new ServiceCollection(); + services.AddLakeWrightCostAttribution(); + services.AddLakeWrightBillingCostAttribution(Configuration(new() + { + ["DatabricksBilling:WorkspaceId"] = "workspace-123" + })); + + var cost = services.Last(descriptor => descriptor.ServiceType == typeof(ICostAttribution)); + var reader = services.Last(descriptor => descriptor.ServiceType == typeof(IBillingUsageReader)); + + cost.ImplementationType.ShouldBe(typeof(BillingCostAttribution)); + reader.ImplementationType.ShouldBe(typeof(DatabricksBillingUsageReader)); + Validator(services).ShouldNotThrow(); + } + private static Action Validator(IServiceCollection services) { var provider = services.BuildServiceProvider(); diff --git a/tests/LakeWright.TenantIsolation.Tests/TestApi.cs b/tests/LakeWright.TenantIsolation.Tests/TestApi.cs index 125b897..1a5b7a3 100644 --- a/tests/LakeWright.TenantIsolation.Tests/TestApi.cs +++ b/tests/LakeWright.TenantIsolation.Tests/TestApi.cs @@ -68,7 +68,9 @@ protected override Task HandleAuthenticateAsync() } } - public static async Task<(IHost Host, HttpClient Client)> StartAsync(PostgresFixture postgres) + public static async Task<(IHost Host, HttpClient Client)> StartAsync( + PostgresFixture postgres, + Action? configureServices = null) { ArgumentNullException.ThrowIfNull(postgres); @@ -102,6 +104,7 @@ protected override Task HandleAuthenticateAsync() services.AddScoped(); services.AddScoped(); services.AddLakeWrightCostAttribution(); + configureServices?.Invoke(services); services.Configure(o => o.Catalog = "analytics"); services.AddAuthentication(StubAuth.SchemeName)