@oneuptime/common 10.5.9 → 10.5.18

package/Server/Infrastructure/PostgresDatabase.ts

@@ -4,6 +4,7 @@ import Sleep from "../../Types/Sleep";
 import { DataSource, DataSourceOptions } from "typeorm";
 import { createDatabase, dropDatabase } from "typeorm-extension";
 import CaptureSpan from "../Utils/Telemetry/CaptureSpan";
+import GracefulShutdown, { ShutdownPriority } from "../Utils/GracefulShutdown";
 
 export type DatabaseSourceOptions = DataSourceOptions;
 export type DatabaseSource = DataSource;
@@ -30,6 +31,15 @@ export default class Database {
 
   @CaptureSpan()
   public static async connect(): Promise<DataSource> {
+    /*
+     * Idempotent: a second connect() must not overwrite (and thereby orphan)
+     * the existing pool. Return the live DataSource instead of building a new
+     * one.
+     */
+    if (this.dataSource) {
+      return this.dataSource;
+    }
+
     let retry: number = 0;
 
     const dataSourceOptions: DataSourceOptions = this.getDatasourceOptions();
@@ -64,7 +74,23 @@ export default class Database {
           }
         };
 
-      return await connectToDatabase();
+      const dataSource: DataSource = await connectToDatabase();
+
+      /*
+       * Drain the pool on shutdown. Registered here (after a successful
+       * connect) so we never register cleanup for a pool that was never
+       * created, and — thanks to GracefulShutdown deduping by name — exactly
+       * once even if connect() is somehow reached twice.
+       */
+      GracefulShutdown.registerHandler(
+        "PostgresDatabase",
+        ShutdownPriority.DataStores,
+        () => {
+          return this.disconnect();
+        },
+      );
+
+      return dataSource;
     } catch (err) {
       logger.error("Postgres Database Connection Failed");
       logger.error(err);

package/Server/Infrastructure/QueueWorker.ts

@@ -8,7 +8,14 @@ import {
 import { Worker } from "bullmq";
 import CaptureSpan from "../Utils/Telemetry/CaptureSpan";
 import AppMetrics from "../Utils/Telemetry/AppMetrics";
+import TelemetryContext from "../Utils/Telemetry/TelemetryContext";
+import Telemetry, {
+  Span,
+  SpanException,
+  SpanStatusCode,
+} from "../Utils/Telemetry";
 import Redis from "./Redis";
+import GracefulShutdown, { ShutdownPriority } from "../Utils/GracefulShutdown";
 
 export default class QueueWorker {
   @CaptureSpan()
@@ -45,7 +52,40 @@ export default class QueueWorker {
       let outcome: "success" | "failure" | "timeout" = "success";
 
       try {
-        await onJobInQueue(job);
+        /*
+         * Seed a telemetry-context scope for this job so every span and log it
+         * produces inherits the queue/job name plus any tenant identifiers
+         * carried in the job payload (projectId, monitorId, incidentId, ...).
+         */
+        await TelemetryContext.runWithContext(
+          {
+            queueName: queueName,
+            jobName: job.name || "unknown",
+            ...TelemetryContext.pickKnownAttributes(job.data),
+          },
+          () => {
+            /*
+             * Wrap the job in an explicit root span so every background job has
+             * a consistent, named trace root that carries the seeded context —
+             * the @CaptureSpan service calls it makes become children of this.
+             */
+            return Telemetry.startActiveSpan<Promise<void>>({
+              name: `worker.job ${queueName}/${job.name || "unknown"}`,
+              fn: async (span: Span): Promise<void> => {
+                try {
+                  await onJobInQueue(job);
+                  span.setStatus({ code: SpanStatusCode.OK });
+                } catch (err) {
+                  span.recordException(err as SpanException);
+                  span.setStatus({ code: SpanStatusCode.ERROR });
+                  throw err;
+                } finally {
+                  span.end();
+                }
+              },
+            });
+          },
+        );
       } catch (err) {
         outcome =
           err instanceof TimeoutException ||
@@ -77,9 +117,19 @@ export default class QueueWorker {
         : {}),
     });
 
-    process.on("SIGINT", async () => {
-      await worker.close();
-    });
+    /*
+     * Stop pulling new jobs and let in-flight ones finish on shutdown. Runs in
+     * the Workers tier — before datastores are drained — so jobs mid-flight can
+     * still reach Postgres / Redis. Replaces a SIGINT-only handler that never
+     * fired in containers (Kubernetes / docker stop send SIGTERM).
+     */
+    GracefulShutdown.registerHandler(
+      `QueueWorker:${queueName}`,
+      ShutdownPriority.Workers,
+      () => {
+        return worker.close();
+      },
+    );
 
     return worker;
   }

package/Server/Infrastructure/Redis.ts

@@ -15,6 +15,7 @@ import logger from "../Utils/Logger";
 import Sleep from "../../Types/Sleep";
 import { Redis as RedisClient, RedisOptions } from "ioredis";
 import CaptureSpan from "../Utils/Telemetry/CaptureSpan";
+import GracefulShutdown, { ShutdownPriority } from "../Utils/GracefulShutdown";
 
 export type ClientType = RedisClient;
 export type RedisOptionsType = RedisOptions;
@@ -122,6 +123,16 @@ export default abstract class Redis {
       logger.debug(
         `Redis connected on ${RedisHostname}:${RedisPort.toNumber()}`,
       );
+
+      // Close the Redis connection on shutdown.
+      GracefulShutdown.registerHandler(
+        "Redis",
+        ShutdownPriority.DataStores,
+        () => {
+          return this.disconnect();
+        },
+      );
+
       return this.client;
     } catch (err) {
       logger.error("Redis Connection Failed");

package/Server/Services/AnalyticsDatabaseService.ts

@@ -366,6 +366,93 @@ export default class AnalyticsDatabaseService<
     return await this._findBy(findBy);
   }
 
+  /**
+   * Group telemetry rows by (serviceId, serviceType) for a project over a
+   * time window, returning the row count and an estimate of the ingested
+   * byte size (ClickHouse `byteSize(*)`, the uncompressed in-memory size of
+   * each row's columns). This is the enumeration source for usage billing:
+   * a single aggregation scan surfaces EVERY resource that emitted
+   * telemetry — real Services, Hosts, Docker hosts, Kubernetes clusters,
+   * Monitors and unattributed (serviceId = projectId) — without needing a
+   * Postgres row per resource. The caller decides which serviceTypes to
+   * bill and how to attribute retention.
+   */
+  @CaptureSpan()
+  public async groupTelemetryUsageByService(data: {
+    projectId: ObjectID;
+    timestampColumnName: keyof TBaseModel | string;
+    startDate: Date;
+    endDate: Date;
+  }): Promise<
+    Array<{
+      serviceId: string;
+      serviceType: string | null;
+      rowCount: number;
+      estimatedBytes: number;
+    }>
+  > {
+    const timestampColumnName: string = data.timestampColumnName.toString();
+
+    if (!this.model.getTableColumn(timestampColumnName)) {
+      throw new BadDataException(
+        `Invalid timestampColumnName: ${timestampColumnName}`,
+      );
+    }
+
+    if (!this.database) {
+      this.useDefaultDatabase();
+    }
+    const databaseName: string =
+      this.database!.getDatasourceOptions().database!;
+
+    const statement: Statement = SQL`SELECT serviceId AS serviceId, serviceType AS serviceType, count() AS rowCount, sum(byteSize(*)) AS estimatedBytes FROM ${databaseName}.${this.model.tableName} WHERE projectId = ${{
+      type: TableColumnType.ObjectID,
+      value: data.projectId,
+    }} AND ${timestampColumnName} >= ${{
+      type: TableColumnType.DateTime64,
+      value: data.startDate,
+    }} AND ${timestampColumnName} <= ${{
+      type: TableColumnType.DateTime64,
+      value: data.endDate,
+    }} GROUP BY serviceId, serviceType`;
+
+    statement.append(
+      " SETTINGS max_execution_time = 60, timeout_overflow_mode = 'break'",
+    );
+
+    const dbResult: ResultSet<"JSON"> = await this.executeQuery(statement);
+    const responseJSON: ResponseJSON<JSONObject> =
+      await dbResult.json<JSONObject>();
+    const items: Array<JSONObject> = responseJSON.data ? responseJSON.data : [];
+
+    const results: Array<{
+      serviceId: string;
+      serviceType: string | null;
+      rowCount: number;
+      estimatedBytes: number;
+    }> = [];
+
+    for (const item of items) {
+      const serviceId: string = (item["serviceId"] as string) || "";
+      if (!serviceId) {
+        continue;
+      }
+      const serviceTypeRaw: unknown = item["serviceType"];
+      const serviceType: string | null =
+        typeof serviceTypeRaw === "string" && serviceTypeRaw.trim()
+          ? serviceTypeRaw
+          : null;
+      results.push({
+        serviceId,
+        serviceType,
+        rowCount: Number(item["rowCount"]) || 0,
+        estimatedBytes: Number(item["estimatedBytes"]) || 0,
+      });
+    }
+
+    return results;
+  }
+
   @CaptureSpan()
   public async aggregateBy(
     aggregateBy: AggregateBy<TBaseModel>,

package/Server/Services/DatabaseService.ts

@@ -54,6 +54,7 @@ import HashedString from "../../Types/HashedString";
 import { JSONObject, JSONValue } from "../../Types/JSON";
 import JSONFunctions from "../../Types/JSONFunctions";
 import ObjectID from "../../Types/ObjectID";
+import TelemetryContext from "../Utils/Telemetry/TelemetryContext";
 import PositiveNumber from "../../Types/PositiveNumber";
 import Text from "../../Types/Text";
 import Typeof from "../../Types/Typeof";
@@ -681,6 +682,69 @@ class DatabaseService<TBaseModel extends BaseModel> extends BaseService {
     }
   }
 
+  /**
+   * Derive the telemetry attribute key for this model's primary id, e.g.
+   * `Incident` -> `incidentId`, `Monitor` -> `monitorId`. Matches the keys in
+   * TelemetryContextAttributes so dashboards/queries stay consistent.
+   */
+  private getTelemetryEntityIdKey(): string {
+    const name: string = this.modelName || "entity";
+    return name.charAt(0).toLowerCase() + name.slice(1) + "Id";
+  }
+
+  /**
+   * Seed the ambient telemetry context with the project (tenant) of an
+   * operation so worker/cron/service spans and logs — which run outside the
+   * HTTP request scope — still carry projectId. Best-effort and safe to call
+   * anywhere.
+   */
+  protected setTelemetryContextFromProps(
+    props: DatabaseCommonInteractionProps | undefined,
+  ): void {
+    try {
+      if (props?.tenantId) {
+        TelemetryContext.setAttributes({
+          projectId: props.tenantId.toString(),
+        });
+      }
+    } catch {
+      // Telemetry must never break a database operation.
+    }
+  }
+
+  /**
+   * Seed the ambient telemetry context from a concrete model instance: the
+   * project (tenant) and the entity's own id (e.g. incidentId, monitorId).
+   * Used on create, where a new entity id is minted. Best-effort and safe.
+   */
+  protected setTelemetryContextFromItem(item: TBaseModel | undefined): void {
+    try {
+      if (!item) {
+        return;
+      }
+
+      const attributes: { [key: string]: string } = {};
+
+      const tenantColumn: string | null = this.model.getTenantColumn();
+      if (tenantColumn) {
+        const tenantValue: unknown = item.getColumnValue(tenantColumn);
+        if (tenantValue) {
+          attributes["projectId"] = String(tenantValue);
+        }
+      }
+
+      if (item.id) {
+        attributes[this.getTelemetryEntityIdKey()] = item.id.toString();
+      }
+
+      if (Object.keys(attributes).length > 0) {
+        TelemetryContext.setAttributes(attributes);
+      }
+    } catch {
+      // Telemetry must never break a database operation.
+    }
+  }
+
   @CaptureSpan()
   public async create(createBy: CreateBy<TBaseModel>): Promise<TBaseModel> {
     const onCreate: OnCreate<TBaseModel> = createBy.props.ignoreHooks
@@ -742,6 +806,9 @@ class DatabaseService<TBaseModel extends BaseModel> extends BaseService {
     try {
       createBy.data = await this.getRepository().save(createBy.data);
 
+      // Seed telemetry context with projectId + <model>Id for this create.
+      this.setTelemetryContextFromItem(createBy.data);
+
       if (!createBy.props.ignoreHooks) {
         createBy.data = await this.onCreateSuccess(
           {
@@ -1216,6 +1283,8 @@ class DatabaseService<TBaseModel extends BaseModel> extends BaseService {
 
   private async _deleteBy(deleteBy: DeleteBy<TBaseModel>): Promise<number> {
     try {
+      this.setTelemetryContextFromProps(deleteBy.props);
+
       if (this.doNotAllowDelete && !deleteBy.props.isRoot) {
         throw new BadDataException("Delete not allowed");
       }
@@ -1444,6 +1513,8 @@ class DatabaseService<TBaseModel extends BaseModel> extends BaseService {
     withDeleted?: boolean | undefined,
   ): Promise<Array<TBaseModel>> {
     try {
+      this.setTelemetryContextFromProps(findBy.props);
+
       let automaticallyAddedCreatedAtInSelect: boolean = false;
 
       if (!findBy.sort || Object.keys(findBy.sort).length === 0) {
@@ -1638,6 +1709,8 @@ class DatabaseService<TBaseModel extends BaseModel> extends BaseService {
 
   private async _updateBy(updateBy: UpdateBy<TBaseModel>): Promise<number> {
     try {
+      this.setTelemetryContextFromProps(updateBy.props);
+
       const onUpdate: OnUpdate<TBaseModel> = updateBy.props.ignoreHooks
         ? { updateBy, carryForward: [] }
         : await this.onBeforeUpdate(updateBy);

package/Server/Services/TelemetryAttributeService.ts

@@ -364,6 +364,7 @@ export class TelemetryAttributeService {
     telemetryType: TelemetryType;
     metricName?: string | undefined;
     attributeKey: string;
+    searchText?: string | undefined;
   }): Promise<string[]> {
     const source: TelemetrySource | null = this.getTelemetrySource(
       data.telemetryType,
@@ -378,15 +379,17 @@ export class TelemetryAttributeService {
       source,
       metricName: data.metricName,
       attributeKey: data.attributeKey,
+      searchText: data.searchText,
     });
   }
 
-  private static async fetchAttributeValuesFromDatabase(data: {
+  private static buildAttributeValuesStatement(data: {
     projectId: ObjectID;
     source: TelemetrySource;
     metricName?: string | undefined;
     attributeKey: string;
-  }): Promise<Array<string>> {
+    searchText?: string | undefined;
+  }): Statement {
     const lookbackStartDate: Date =
       TelemetryAttributeService.getLookbackStartDate();
 
@@ -419,6 +422,26 @@ export class TelemetryAttributeService {
       );
     }
 
+    /*
+     * Case-insensitive substring filter so the value autocomplete keeps
+     * narrowing server-side as the user types. Without it only the first
+     * ATTRIBUTE_VALUES_LIMIT values (alphabetically) are ever reachable,
+     * which hides matches on high-cardinality keys (host.name, url, ...).
+     * Mirrors the ILIKE idiom used for bodySearchText / nameSearchText.
+     */
+    if (data.searchText && data.searchText.trim().length > 0) {
+      statement.append(
+        SQL`
+        AND ${data.source.attributesColumn}[${{
+          type: TableColumnType.Text,
+          value: data.attributeKey,
+        }}] ILIKE ${{
+          type: TableColumnType.Text,
+          value: `%${data.searchText.trim()}%`,
+        }}`,
+      );
+    }
+
     statement.append(
       SQL`
       ORDER BY attributeValue ASC
@@ -428,6 +451,19 @@ export class TelemetryAttributeService {
       }}`,
     );
 
+    return statement;
+  }
+
+  private static async fetchAttributeValuesFromDatabase(data: {
+    projectId: ObjectID;
+    source: TelemetrySource;
+    metricName?: string | undefined;
+    attributeKey: string;
+    searchText?: string | undefined;
+  }): Promise<Array<string>> {
+    const statement: Statement =
+      TelemetryAttributeService.buildAttributeValuesStatement(data);
+
     const dbResult: Results = await data.source.service.executeQuery(statement);
     const response: DbJSONResponse = await dbResult.json<{
       data?: Array<JSONObject>;

package/Server/Services/TelemetryExceptionService.ts

@@ -1,10 +1,9 @@
 import DatabaseService from "./DatabaseService";
 import Model from "../../Models/DatabaseModels/TelemetryException";
-import TelemetryServiceModel from "../../Models/DatabaseModels/Service";
+import ServiceType from "../../Types/Telemetry/ServiceType";
 import AIAgentTask from "../../Models/DatabaseModels/AIAgentTask";
 import AIAgentTaskTelemetryException from "../../Models/DatabaseModels/AIAgentTaskTelemetryException";
 import ObjectID from "../../Types/ObjectID";
-import PositiveNumber from "../../Types/PositiveNumber";
 import SortOrder from "../../Types/BaseDatabase/SortOrder";
 import BadDataException from "../../Types/Exception/BadDataException";
 import AIAgentTaskType from "../../Types/AI/AIAgentTaskType";
@@ -13,7 +12,6 @@ import { FixExceptionTaskMetadata } from "../../Types/AI/AIAgentTaskMetadata";
 import DatabaseCommonInteractionProps from "../../Types/BaseDatabase/DatabaseCommonInteractionProps";
 import AIAgentTaskService from "./AIAgentTaskService";
 import AIAgentTaskTelemetryExceptionService from "./AIAgentTaskTelemetryExceptionService";
-import ServiceService from "./ServiceService";
 import QueryHelper from "../Types/Database/QueryHelper";
 import ModelPermission from "../Types/Database/Permissions/Index";
 import CaptureSpan from "../Utils/Telemetry/CaptureSpan";
@@ -24,7 +22,13 @@ export interface CreateAIAgentTaskForExceptionParams {
 }
 
 export interface DashboardServiceSummary {
-  service: TelemetryServiceModel;
+  /*
+   * Polymorphic: a real Service, a Host/DockerHost/KubernetesCluster id, or
+   * the projectId for unattributed telemetry. The client resolves the
+   * display name per serviceType.
+   */
+  serviceId: string;
+  serviceType: ServiceType | null;
   unresolvedCount: number;
   totalOccurrences: number;
 }
@@ -182,11 +186,9 @@ export class Service extends DatabaseService<Model> {
       lastSeenAt: true,
       firstSeenAt: true,
       environment: true,
-      service: {
-        _id: true,
-        name: true,
-        serviceColor: true,
-      },
+      // serviceId is polymorphic; the client resolves it per serviceType.
+      serviceId: true,
+      serviceType: true,
     };
 
     const [
@@ -284,6 +286,7 @@ export class Service extends DatabaseService<Model> {
 
     interface AggregateRow {
       serviceId: string | null;
+      serviceType: string | null;
       unresolvedCount: string;
       totalOccurrences: string | null;
     }
@@ -292,6 +295,7 @@ export class Service extends DatabaseService<Model> {
       "TelemetryException",
     )
       .select(`"TelemetryException"."serviceId"`, "serviceId")
+      .addSelect(`"TelemetryException"."serviceType"`, "serviceType")
       .addSelect(`COUNT(*)`, "unresolvedCount")
       .addSelect(
         `COALESCE(SUM("TelemetryException"."occuranceCount"), 0)`,
@@ -305,6 +309,7 @@ export class Service extends DatabaseService<Model> {
       .andWhere(`"TelemetryException"."deletedAt" IS NULL`)
       .andWhere(`"TelemetryException"."serviceId" IS NOT NULL`)
       .groupBy(`"TelemetryException"."serviceId"`)
+      .addGroupBy(`"TelemetryException"."serviceType"`)
       .orderBy(`"unresolvedCount"`, "DESC")
       .getRawMany()) as Array<AggregateRow>;
 
@@ -312,52 +317,22 @@ export class Service extends DatabaseService<Model> {
       return [];
     }
 
-    const serviceIds: Array<string> = [];
-    for (const row of rows) {
-      if (row.serviceId) {
-        serviceIds.push(row.serviceId);
-      }
-    }
-
-    if (serviceIds.length === 0) {
-      return [];
-    }
-
-    const services: Array<TelemetryServiceModel> = await ServiceService.findBy({
-      query: {
-        projectId,
-        _id: QueryHelper.any(serviceIds),
-      },
-      select: {
-        _id: true,
-        name: true,
-        serviceColor: true,
-      },
-      limit: new PositiveNumber(serviceIds.length),
-      skip: new PositiveNumber(0),
-      props,
-    });
-
-    const serviceById: Map<string, TelemetryServiceModel> = new Map();
-    for (const service of services) {
-      if (service._id) {
-        serviceById.set(service._id, service);
-      }
-    }
-
+    /*
+     * serviceId is polymorphic — do NOT resolve it to a Service here. The
+     * old code looked each serviceId up in the Service table and dropped
+     * any that didn't match, which silently excluded Host / DockerHost /
+     * KubernetesCluster and unattributed (Unknown) telemetry. Return the
+     * raw (serviceId, serviceType) + counts; the client resolves the
+     * display name per serviceType.
+     */
     const summaries: Array<DashboardServiceSummary> = [];
     for (const row of rows) {
       if (!row.serviceId) {
         continue;
       }
-      const service: TelemetryServiceModel | undefined = serviceById.get(
-        row.serviceId,
-      );
-      if (!service) {
-        continue;
-      }
       summaries.push({
-        service,
+        serviceId: row.serviceId,
+        serviceType: (row.serviceType as ServiceType | null) ?? null,
         unresolvedCount: parseInt(row.unresolvedCount, 10) || 0,
         totalOccurrences: parseInt(row.totalOccurrences || "0", 10) || 0,
       });