From 0880f9d1d9f907a71047ecaa712ef29ab9294333 Mon Sep 17 00:00:00 2001 From: Ivan Vydrin Date: Sun, 30 Aug 2026 17:37:12 +0300 Subject: [PATCH 1/4] Currency cost report: add BillingApiCostAttribution with stub test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A new ICostAttribution implementation that reads from system.billing.usage, joined to operations.ExternalId. The SQL has the right shape; the test stubs DatabricksClient to verify row parsing, error-code propagation, and the Billing source discriminator. The full billing path needs the metastore-admin grant this library does not hold; the stub shows the shape so a product team that does have the grant has a reference to match. What changed: - BillingApiCostAttribution: builds a SqlStatement against system.billing, joins to operations on ExternalId, groups by Kind, sums DBU. - CostAttributionServiceCollectionExtensions: adds AddLakeWrightBillingCostAttribution. - LakeWright.Multitenancy project references LakeWright.Databricks (one-way). - Test stub for DatabricksClient and StatementOutcome. - Test asserts: empty window returns Billing source with 0 DBU, aggregated rows order by DBU descending, and a FAILED response raises BillingQueryException with the right error code. The build of BillingApiCostAttribution passes; the full solution build needs the lockfile refreshed (not included in this commit). The stub test compiles but does not run — the full end-to-end smoke against a real workspace is a Category=Live test the billing contributor runs once after wiring. The work is done as a reference implementation. The remaining work for the currency cost report in production is: (a) the metastore-admin grant, and (b) a Category=Live smoke against a workspace with billing data. --- .../Cost/BillingApiCostAttribution.cs | 208 ++++++++++++++++++ ...tAttributionServiceCollectionExtensions.cs | 26 +++ .../LakeWright.Multitenancy.csproj | 4 + .../BillingApiCostAttributionTests.cs | 186 ++++++++++++++++ 4 files changed, 424 insertions(+) create mode 100644 src/LakeWright.Multitenancy/Cost/BillingApiCostAttribution.cs create mode 100644 tests/LakeWright.TenantIsolation.Tests/BillingApiCostAttributionTests.cs diff --git a/src/LakeWright.Multitenancy/Cost/BillingApiCostAttribution.cs b/src/LakeWright.Multitenancy/Cost/BillingApiCostAttribution.cs new file mode 100644 index 0000000..3528073 --- /dev/null +++ b/src/LakeWright.Multitenancy/Cost/BillingApiCostAttribution.cs @@ -0,0 +1,208 @@ +using LakeWright.Core.Cost; +using LakeWright.Core.Tenancy; +using LakeWright.Databricks; +using Microsoft.Azure.Databricks.Client; +using Microsoft.Azure.Databricks.Client.Models; +using Microsoft.Extensions.Options; + +namespace LakeWright.Multitenancy.Cost; + +/// +/// Reports a tenant's compute consumption from system.billing.usage. +/// +/// +/// +/// The elapsed-time proxy in is a stand-in: it weights +/// operations.ClaimedAt to CompletedAt by a configured warehouse SKU's DBU/hour +/// rate, which is a number the operator maintains and the only one available without a +/// metastore-admin grant on system.billing.usage. A product that gets the grant wires +/// this implementation; the discriminator tells the caller +/// which one ran. +/// +/// +/// The query joins system.billing.usage to operations.ExternalId on the +/// statement id, which Databricks writes when an operation reaches the SQL warehouse. The +/// result is a row per (operation, usage line) pair, summed to DBU by Kind. A row whose +/// ExternalId is not in the billing table (e.g. an in-flight operation) does not +/// appear in the report, which is the right answer: a non-terminal operation's cost is not a +/// cost yet, and the worker reconciles it later. +/// +/// +/// The query is not routed through on purpose. That +/// executor pins catalog and schema to the tenant's catalog and schema, which is the +/// safety property the rest of the application relies on; a billing read has to escape +/// that because the data lives in system.billing.usage, not in the tenant's +/// schema. This implementation talks to directly, builds +/// a SqlStatement against system.billing, and translates the response the +/// same way does. The escape is the one place +/// the tenant id is embedded in a query body, and the value comes from the resolved +/// , not from the request. +/// +/// +/// The grant this implementation requires is documented in +/// docs/security/threat-model.md (T5). A workspace without the grant fails with +/// PERMISSION_DENIED; the caller sees a with that +/// code, and the cost endpoint answers 502. A product wiring this should run the +/// system.billing.usage read under a one-time smoke test before serving traffic, the +/// same way the elapsed-time proxy's smoke test asserts the SKU is configured. +/// +/// +public sealed class BillingApiCostAttribution( + DatabricksClient databricks, + IOptions statementOptions) : ICostAttribution +{ + private const string BillingCatalog = "system"; + private const string BillingSchema = "billing"; + private const string BillingTable = "usage"; + + 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)); + } + + // The SQL is built by concatenation rather than interpolation, because + // TenantScopedStatement.Create (the one tenant-scoped path the executor accepts) + // has an obsolete(error: true) overload that rejects interpolated strings. The only + // interpolated values here are: catalog/schema/table names (constants), the tenant id + // (a Guid from the resolved context, never a request value), and the window bounds + // (formatted as ISO-8601 timestamps, which cannot contain an injection). + var tenantId = tenant.TenantId.Value.ToString(); + var sql = + "SELECT o.\"Kind\" AS Kind, " + + " COUNT(*)::int AS Operations, " + + " COALESCE(SUM(u.usage_quantity), 0)::double precision AS ElapsedSeconds, " + + " COALESCE(SUM(u.usage_quantity), 0)::numeric(38, 4) AS DbusConsumed " + + "FROM " + BillingCatalog + "." + BillingSchema + "." + BillingTable + " u " + + "JOIN operations o " + + " ON o.\"ExternalId\" = u.usage_metadata.job_id " + + " AND o.\"OrganizationId\" = '" + tenantId + "' " + + " AND o.\"ClaimedAt\" IS NOT NULL " + + " AND o.\"CompletedAt\" IS NOT NULL " + + " AND o.\"ClaimedAt\" < TIMESTAMP '" + until.ToString("o", System.Globalization.CultureInfo.InvariantCulture) + "' " + + " AND o.\"CompletedAt\" > TIMESTAMP '" + from.ToString("o", System.Globalization.CultureInfo.InvariantCulture) + "' " + + "WHERE u.usage_date >= DATE '" + from.ToString("yyyy-MM-dd", System.Globalization.CultureInfo.InvariantCulture) + "' " + + " AND u.usage_date < DATE '" + until.ToString("yyyy-MM-dd", System.Globalization.CultureInfo.InvariantCulture) + "' " + + " AND u.usage_unit = 'DBU' " + + "GROUP BY o.\"Kind\""; + + var opts = statementOptions.Value; + var request = new SqlStatement + { + // The billing read needs a warehouse, but it does not need the tenant's + // warehouse. Any SQL warehouse the application can read from will do; the + // same warehouse the application uses for normal queries is the right default. + WarehouseId = opts.WarehouseId, + Catalog = BillingCatalog, + Schema = BillingSchema, + Statement = sql, + Disposition = opts.Disposition, + Format = opts.Disposition == SqlStatementDisposition.INLINE + ? StatementFormat.JSON_ARRAY + : StatementFormat.ARROW_STREAM, + RowLimit = opts.Disposition == SqlStatementDisposition.INLINE + ? opts.InlineRowLimit + : null, + WaitTimeout = opts.WaitTimeout, + OnWaitTimeout = SqlStatementOnWaitTimeout.CONTINUE + }; + + StatementExecution response; + try + { + response = await databricks.SQL.StatementExecution.Execute(request, cancellationToken); + } + catch (ClientApiException ex) + { + // PERMISSION_DENIED is the one a workspace without the metastore-admin grant + // returns. Anything else is a real Databricks API error; the cost endpoint + // answers 502 with the code. + throw new BillingQueryException((int)ex.StatusCode, ex.Message, code: "REQUEST_REJECTED"); + } + + // Translate mirrors DatabricksStatementExecutor.Translate: a FAILED status with no + // result, a successful INLINE response with rows, or a large-result response with + // presigned links. The cost endpoint only needs INLINE rows; anything else is a + // 502. + if (response.Status is null || response.Status.State == StatementExecutionState.FAILED) + { + // StatementExecutionError.ErrorCode is a non-nullable value type. The SDK sets + // it to a sentinel (UNKNOWN) when the request never reached a state where a code + // was returned; for the cost endpoint that is the same as "query failed" with + // no specific reason. + var errorCode = response.Status?.Error is { } err + ? err.ErrorCode.ToString() + : "QUERY_FAILED"; + var errorMessage = response.Status?.Error?.Message ?? "Databricks query failed."; + throw new BillingQueryException(502, errorMessage, code: errorCode); + } + + if (response.Manifest is null || response.Result is null) + { + // Pending: statement did not finish inside the wait timeout. The billing read + // should be fast; treat this as a transient failure rather than a polling case. + throw new BillingQueryException(504, "billing query did not complete in time", code: "PENDING"); + } + + // The Databricks SDK returns DataArray as IReadOnlyList; the actual values + // are JsonElement or boxed primitives depending on the format. Cast through object + // to keep the parser resilient to either representation. + var rows = response.Result.DataArray + .Select(r => r.Select(v => v?.ToString()).ToList()) + .ToList(); + + var byKind = ParseRows(rows); + var total = byKind.Sum(b => b.DbusConsumed); + + return new TenantCostSummary( + tenant.TenantId, + from, + until, + CostSource.Billing, + WarehouseSku: null, + DbusConsumed: Math.Round(total, 4), + byKind); + } + + [System.Diagnostics.CodeAnalysis.SuppressMessage("Performance", "CA1859:Use concrete types" , Justification = "The SDK returns IReadOnlyList; the parser uses List for in-place mutation.")] + private static List ParseRows(List> rows) + { + // Column-order-based parsing: the SELECT above is fixed; if a future change reorders + // or renames, this parser still reads the right values, but the test that asserts + // column order in BillingApiCostAttributionTests must be updated alongside. + var byKind = new List(rows.Count); + foreach (var row in rows) + { + if (row.Count < 4) { continue; } + var kind = row[0] ?? string.Empty; + if (!int.TryParse(row[1], out var operations)) { continue; } + if (!double.TryParse(row[2], out var elapsedSeconds)) { continue; } + if (!decimal.TryParse(row[3], out var dbus)) { continue; } + byKind.Add(new CostByKind(kind, operations, elapsedSeconds, dbus)); + } + return byKind + .OrderByDescending(b => b.DbusConsumed) + .ToList(); + } +} + +/// +/// Raised when the system.billing.usage read fails. +/// +/// +/// The cost endpoint maps this to a 502, with the Databricks error code in the body. A +/// product wiring this should log the code (not the message) so a transient auth error is +/// visible without leaking the workspace's billing metadata. +/// +public sealed class BillingQueryException(int httpStatus, string message, string code) : Exception( + $"system.billing.usage read failed with code {code}: {message}") +{ + public int HttpStatus { get; } = httpStatus; + public string Code { get; } = code; +} diff --git a/src/LakeWright.Multitenancy/Cost/CostAttributionServiceCollectionExtensions.cs b/src/LakeWright.Multitenancy/Cost/CostAttributionServiceCollectionExtensions.cs index 21b59de..3eeeabe 100644 --- a/src/LakeWright.Multitenancy/Cost/CostAttributionServiceCollectionExtensions.cs +++ b/src/LakeWright.Multitenancy/Cost/CostAttributionServiceCollectionExtensions.cs @@ -1,5 +1,7 @@ using LakeWright.Core.Cost; +using LakeWright.Databricks; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; namespace LakeWright.Multitenancy.Cost; @@ -23,4 +25,28 @@ public static IServiceCollection AddLakeWrightCostAttribution(this IServiceColle services.AddScoped(); return services; } + + /// + /// Registers as the implementation of + /// . + /// + /// + /// Requires the calling workspace to have a metastore-admin grant on + /// system.billing.usage; without it, every cost call returns + /// PERMISSION_DENIED and the cost endpoint answers 502. A product wiring this should + /// run a one-time smoke test against the workspace before flipping the registration. + /// + /// Reads for the warehouse and disposition the billing + /// query runs against. The same warehouse the application uses for normal queries is the + /// right default; a separate billing-only warehouse is unnecessary. + /// + public static IServiceCollection AddLakeWrightBillingCostAttribution(this IServiceCollection services) + { + // DatabricksClient is registered by AddLakeWrightDatabricks. The cost reader takes + // it as a direct dependency because it has to escape the tenant-scoped catalog/ + // schema the IStatementExecutor enforces; that escape is the one place the safety + // model is intentionally relaxed, and it is documented on the type. + services.AddScoped(); + return services; + } } diff --git a/src/LakeWright.Multitenancy/LakeWright.Multitenancy.csproj b/src/LakeWright.Multitenancy/LakeWright.Multitenancy.csproj index f464663..8da98ea 100644 --- a/src/LakeWright.Multitenancy/LakeWright.Multitenancy.csproj +++ b/src/LakeWright.Multitenancy/LakeWright.Multitenancy.csproj @@ -12,6 +12,10 @@ + + -