@comprehend/telemetry-node 0.1.4 → 0.2.1

package/src/ComprehendDevSpanProcessor.ts

@@ -1,19 +1,25 @@
 import {sha256} from "@noble/hashes/sha2";
 import {utf8ToBytes} from "@noble/hashes/utils";
-import {Context} from "@opentelemetry/api";
+import {Context, SpanKind} from "@opentelemetry/api";
 import {ReadableSpan, Span, SpanProcessor} from "@opentelemetry/sdk-trace-base";
 import {
-    DatabaseQueryObservation,
+    AttributeType,
+    CustomObservation,
+    CustomSpanObservationSpecification,
+    DatabaseQueryMessage,
     HttpClientObservation,
-    HttpServerObservation, NewObservedDatabaseConnectionMessage,
-    NewObservedDatabaseMessage, NewObservedDatabaseQueryMessage,
+    HttpServerObservation,
+    NewObservedDatabaseConnectionMessage,
+    NewObservedDatabaseMessage,
     NewObservedHttpRequestMessage,
     NewObservedHttpRouteMessage,
     NewObservedHttpServiceMessage,
     NewObservedServiceMessage,
-    ObservationInputMessage
+    ObservationInputMessage,
+    SpanMatcherRule,
+    TraceSpansMessage,
+    CustomMetricSpecification,
 } from "./wire-protocol";
-import {analyzeSQL} from "./sql-analyzer";
 import {WebSocketConnection} from "./WebSocketConnection";
 
 interface ObservedService {
@@ -32,7 +38,7 @@ interface ObservedHTTPRoute {
 
 interface ObservedDatabase {
     system: string;
-    connection?: string; // Connection string, hopefully scrubbed of user info
+    connection?: string;
     host?: string;
     port?: number;
     name?: string;
@@ -56,11 +62,6 @@ interface ObservedDatabaseConnection {
     user?: string;
 }
 
-interface ObservedDatabaseQuery {
-    hash: string;
-    query: string;
-}
-
 const sqlDbSystems = new Set([
     'mysql', 'postgresql', 'mssql', 'oracle', 'db2', 'sqlite', 'hsqldb', 'h2',
     'informix', 'cockroachdb', 'redshift', 'tidb', 'trino', 'greenplum'
@@ -73,33 +74,45 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
     private observedHttpServices: ObservedHttpService[] = [];
     private observedInteractions: Map<string, Map<string, {
         dbConnections: ObservedDatabaseConnection[],
-        dbQueries: ObservedDatabaseQuery[],
         httpRequest?: ObservedHttpRequest
     }>> = new Map();
-    private observationsSeq = 1;
+    private processContextSet = false;
+    private spanObservationSpecs: CustomSpanObservationSpecification[] = [];
+
+    constructor(connection: WebSocketConnection) {
+        this.connection = connection;
+    }
 
-    constructor(options: { organization: string, token: string, debug?: boolean | ((message: string) => void) }) {
-        this.connection = new WebSocketConnection(options.organization, options.token,
-            options.debug === true ? console.log : options.debug === false ? undefined : options.debug);
+    updateCustomMetrics(specs: CustomMetricSpecification[]): void {
+        this.spanObservationSpecs = specs.filter(
+            (s): s is CustomSpanObservationSpecification => s.type === 'span'
+        );
     }
 
     onStart(span: Span, parentContext: Context): void {
     }
 
     onEnd(span: ReadableSpan): void {
+        const process = () => this.processSpan(span);
+        if (span.resource.asyncAttributesPending) {
+            span.resource.waitForAsyncAttributes?.().then(process);
+        } else {
+            process();
+        }
+    }
+
+    private processSpan(span: ReadableSpan): void {
         const currentService = this.discoverService(span);
         if (!currentService)
             return;
 
         const attrs = span.attributes;
-        if (span.kind === 1) {
-            // Server span, see if it's something to ingest.
+        if (span.kind === 1) { // Server span
             if (attrs['http.route'] && attrs['http.method']) {
                 this.processHTTPRoute(currentService, attrs['http.route'] as string, attrs['http.method'] as string, span);
             }
         }
-        else if (span.kind === 2) {
-            // Client span, see if it's something to ingest.
+        else if (span.kind === 2) { // Client span
             if (attrs['http.url']) {
                 this.processHttpRequest(currentService, attrs['http.url'] as string, span);
             }
@@ -107,10 +120,15 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         if (attrs['db.system']) {
             this.processDatabaseOperation(currentService, span);
         }
+
+        // Report trace span for every span
+        this.reportTraceSpan(span);
+
+        // Check custom span observation matching
+        this.checkCustomSpanMatching(span);
     }
 
     private discoverService(span: ReadableSpan): ObservedService | undefined {
-        // Look for an existing matching entry.
         const resAttrs = span.resource.attributes;
         const name = resAttrs['service.name'] as string | undefined;
         if (!name)
@@ -125,7 +143,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         if (existing)
             return existing;
 
-        // New; hash it and add it to the observed services.
         const idString = `service:${name}:${namespace ?? ''}:${environment ?? ''}`;
         const hash = hashIdString(idString);
         const newService: ObservedService = {
@@ -134,7 +151,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         };
         this.observedServices.push(newService);
 
-        // Ingest its existence.
         const message: NewObservedServiceMessage = {
             event: "new-entity",
             type: "service",
@@ -144,22 +160,27 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             ...(environment ? { environment } : {})
         };
         this.ingestMessage(message);
+
+        // Set process context on first service discovery
+        if (!this.processContextSet) {
+            this.processContextSet = true;
+            const resources = flattenResourceAttributes(resAttrs);
+            this.connection.setProcessContext(hash, resources);
+        }
+
         return newService;
     }
 
     private processHTTPRoute(service: ObservedService, route: string, method: string, span: ReadableSpan): void {
-        // Check if this route+method already exists under the service
         let observedRoute = service.httpRoutes.find(r =>
             r.route === route && r.method === method
         );
         if (!observedRoute) {
-            // It's new; hash it and add it to our collection.
             const idString = `http-route:${service.hash}:${method}:${route}`;
             const hash = hashIdString(idString);
             observedRoute = { method, route, hash };
             service.httpRoutes.push(observedRoute);
 
-            // Emit observation message
             const message: NewObservedHttpRouteMessage = {
                 event: "new-entity",
                 type: "http-route",
@@ -170,14 +191,13 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             this.ingestMessage(message);
         }
 
-        // Extract the request path, making sure we only get that, without any query string.
+        // Extract the request path, making sure we strip any query string.
         const attrs = span.attributes;
         let path: string;
         if (attrs['http.target']) {
             try {
-                // This might be just a path like "/search?q=foo"
                 const rawTarget = attrs['http.target'] as string;
-                const fakeUrl = new URL(rawTarget, 'http://placeholder'); // placeholder base
+                const fakeUrl = new URL(rawTarget, 'http://placeholder');
                 path = fakeUrl.pathname;
             } catch {
                 path = '/';
@@ -196,7 +216,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             path = '/';
         }
 
-        // Build and ingest observation.
         const status = attrs['http.status_code'] as number | undefined ?? 0;
         const duration = span.duration;
         const httpVersion = attrs['http.flavor'] as string | undefined;
@@ -204,9 +223,12 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         const requestBytes = attrs['http.request_content_length'] as number | undefined;
         const responseBytes = attrs['http.response_content_length'] as number | undefined;
         const { message: errorMessage, type: errorType, stack: stack } = extractErrorInfo(span);
+        const { traceId, spanId } = getSpanContext(span);
         const observation: HttpServerObservation = {
             type: 'http-server',
             subject: observedRoute.hash,
+            spanId,
+            traceId,
             timestamp: span.startTime,
             path,
             status,
@@ -221,13 +243,12 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         };
         this.ingestMessage({
             event: "observations",
-            seq: this.observationsSeq++,
+            seq: this.connection.nextSeq(),
             observations: [observation]
         });
     }
 
     private processDatabaseOperation(currentService: ObservedService, span: ReadableSpan): void {
-        // Parse the connection string.
         const attrs = span.attributes;
         const system = attrs['db.system'] as string;
         const rawConnection = attrs['db.connection_string'] as string | undefined;
@@ -241,11 +262,9 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
                 name: attrs['db.name'] as string | undefined,
             };
 
-        // See if we already have an entry for this database.
         const hash = hashIdString(`database:${system}:${parsed.host ?? ''}:${parsed.port ?? ''}:${parsed.name ?? ''}`);
         let observedDatabase = this.observedDatabases.find(db => db.hash === hash);
 
-        // If we see it for the first time, add and ingest it.
         if (!observedDatabase) {
             observedDatabase = {
                 system,
@@ -256,7 +275,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             };
             this.observedDatabases.push(observedDatabase);
 
-            // The existence of the database.
             const message: NewObservedDatabaseMessage = {
                 event: "new-entity",
                 type: "database",
@@ -271,7 +289,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             this.ingestMessage(message);
         }
 
-        // The connection to the database should have an interaction.
         const interactions = this.getInteractions(currentService.hash, observedDatabase.hash);
         let connectionInteraction = interactions.dbConnections.find(c => c.connection === parsed.scrubbed && c.user === parsed.user);
         if (!connectionInteraction) {
@@ -293,69 +310,40 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             this.ingestMessage(message);
         }
 
-        // The query of the database from the service (only for SQL for now).
+        // For SQL databases, send raw query to server for analysis
         if (sqlDbSystems.has(system) && attrs['db.statement']) {
-            // The interaction, based upon the normalized query.
-            let queryInfo = analyzeSQL(attrs['db.statement'] as string);
-            let queryInteraction = interactions.dbQueries.find(q => q.query === queryInfo.normalizedQuery);
-            if (!queryInteraction) {
-                queryInteraction = {
-                    hash: hashIdString(`db-query:${currentService.hash}:${observedDatabase.hash}:${queryInfo.normalizedQuery}`),
-                    query: queryInfo.normalizedQuery
-                };
-                interactions.dbQueries.push(queryInteraction);
-                const selects = getTablesWithOperation(queryInfo.tableOperations, 'SELECT');
-                const inserts = getTablesWithOperation(queryInfo.tableOperations, 'INSERT');
-                const updates = getTablesWithOperation(queryInfo.tableOperations, 'UPDATE');
-                const deletes = getTablesWithOperation(queryInfo.tableOperations, 'DELETE');
-                const message: NewObservedDatabaseQueryMessage = {
-                    event: "new-interaction",
-                    type: "db-query",
-                    hash: queryInteraction.hash,
-                    from: currentService.hash,
-                    to: observedDatabase.hash,
-                    query: queryInfo.presentableQuery,
-                    ...(selects ? { selects } : {}),
-                    ...(inserts ? { inserts } : {}),
-                    ...(updates ? { updates } : {}),
-                    ...(deletes ? { deletes } : {}),
-                };
-                this.ingestMessage(message);
-            }
-
-            // Build and ingest observation.
             const duration = span.duration;
             const returnedRows = (attrs['db.response.returned_rows'] as number | undefined)
                 ?? (attrs['db.sql.rows'] as number | undefined);
             const { message: errorMessage, type: errorType, stack } = extractErrorInfo(span);
-            const observation: DatabaseQueryObservation = {
-                type: "db-query",
-                subject: queryInteraction.hash,
+            const { traceId, spanId } = span.spanContext();
+            const dbQueryMessage: DatabaseQueryMessage = {
+                event: "db-query",
+                seq: this.connection.nextSeq(),
+                query: attrs['db.statement'] as string,
+                from: currentService.hash,
+                to: observedDatabase.hash,
                 timestamp: span.startTime,
                 duration,
+                traceId,
+                spanId,
                 ...(errorMessage ? { errorMessage } : {}),
-                ...(errorType ? {  errorType } : {}),
+                ...(errorType ? { errorType } : {}),
                 ...(stack ? { stack } : {}),
                 ...(returnedRows !== undefined ? { returnedRows } : {})
             };
-            this.ingestMessage({
-                event: "observations",
-                seq: this.observationsSeq++,
-                observations: [observation]
-            });
+            this.ingestMessage(dbQueryMessage);
         }
     }
 
     private processHttpRequest(currentService: ObservedService, url: string, span: ReadableSpan): void {
-        // Build identity based upon protocol, host, and port.
         const parsed = new URL(url);
-        const protocol = parsed.protocol.replace(':', ''); // Remove trailing colon
+        const protocol = parsed.protocol.replace(':', '');
         const host = parsed.hostname;
         const port = parsed.port ? parseInt(parsed.port) : (protocol === 'https' ? 443 : 80);
         const idString = `http-service:${protocol}:${host}:${port}`;
         const hash = hashIdString(idString);
 
-        // Ingest it if it's not already observed.
         let observedHttpService = this.observedHttpServices.find(s =>
             s.protocol === protocol && s.host === host && s.port === port
         );
@@ -363,7 +351,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             observedHttpService = { protocol, host, port, hash };
             this.observedHttpServices.push(observedHttpService);
 
-            // The existence of the service.
             const message: NewObservedHttpServiceMessage = {
                 event: "new-entity",
                 type: "http-service",
@@ -375,7 +362,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             this.ingestMessage(message);
         }
 
-        // Ingest the interaction if first observed.
         const interactions = this.getInteractions(currentService.hash, observedHttpService.hash);
         if (!interactions.httpRequest) {
             const idString = `http-request:${currentService.hash}:${observedHttpService.hash}`;
@@ -391,11 +377,10 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             this.ingestMessage(message);
         }
 
-        // Build and ingest observation.
         const attrs = span.attributes;
         const path = parsed.pathname || '/';
         const method = span.attributes['http.method'] as string;
-        if (!method) // Really should always be there
+        if (!method)
             return;
         const status = attrs['http.status_code'] as number | undefined;
         const duration = span.duration;
@@ -403,9 +388,12 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         const requestBytes = attrs['http.request_content_length'] as number | undefined;
         const responseBytes = attrs['http.response_content_length'] as number | undefined;
         const { message: errorMessage, type: errorType, stack } = extractErrorInfo(span);
+        const { traceId, spanId } = getSpanContext(span);
         const observation: HttpClientObservation = {
             type: "http-client",
             subject: interactions.httpRequest.hash,
+            spanId,
+            traceId,
             timestamp: span.startTime,
             path,
             method,
@@ -415,16 +403,65 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             ...(requestBytes !== undefined ? { requestBytes } : {}),
             ...(responseBytes !== undefined ? { responseBytes } : {}),
             ...(errorMessage ? { errorMessage } : {}),
-            ...(errorType ? {  errorType } : {}),
+            ...(errorType ? { errorType } : {}),
             ...(stack ? { stack } : {})
         };
         this.ingestMessage({
             event: "observations",
-            seq: this.observationsSeq++,
+            seq: this.connection.nextSeq(),
             observations: [observation]
         });
     }
 
+    private reportTraceSpan(span: ReadableSpan): void {
+        const { traceId, spanId } = getSpanContext(span);
+        const parent = getParentSpanId(span);
+        const traceSpanMessage: TraceSpansMessage = {
+            event: "tracespans",
+            seq: this.connection.nextSeq(),
+            data: [{
+                trace: traceId,
+                span: spanId,
+                parent,
+                name: span.name,
+                timestamp: span.startTime,
+            }]
+        };
+        this.ingestMessage(traceSpanMessage);
+    }
+
+    private checkCustomSpanMatching(span: ReadableSpan): void {
+        for (const spec of this.spanObservationSpecs) {
+            if (matchSpanRule(span, spec.rule)) {
+                const { traceId, spanId } = getSpanContext(span);
+                const { message: errorMessage, type: errorType, stack } = extractErrorInfo(span);
+                const collectedAttrs: Record<string, AttributeType> = {};
+                for (const [key, value] of Object.entries(span.attributes)) {
+                    if (value !== undefined) {
+                        collectedAttrs[key] = value as AttributeType;
+                    }
+                }
+                const observation: CustomObservation = {
+                    type: "custom",
+                    subject: spec.subject,
+                    id: spec.subject,
+                    spanId,
+                    traceId,
+                    timestamp: span.startTime,
+                    attributes: collectedAttrs,
+                    ...(errorMessage ? { errorMessage } : {}),
+                    ...(errorType ? { errorType } : {}),
+                    ...(stack ? { stack } : {}),
+                };
+                this.ingestMessage({
+                    event: "observations",
+                    seq: this.connection.nextSeq(),
+                    observations: [observation]
+                });
+            }
+        }
+    }
+
     private getInteractions(from: string, to: string) {
         let fromMap = this.observedInteractions.get(from);
         if (!fromMap) {
@@ -433,7 +470,7 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
         }
         let interactions = fromMap.get(to);
         if (!interactions) {
-            interactions = { httpRequest: undefined, dbConnections: [], dbQueries: [] };
+            interactions = { httpRequest: undefined, dbConnections: [] };
             fromMap.set(to, interactions);
         }
         return interactions;
@@ -447,7 +484,6 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
     }
 
     async shutdown(): Promise<void> {
-        this.connection.close()
     }
 }
 
@@ -457,8 +493,36 @@ function hashIdString(idString: string) {
         .join('');
 }
 
-/** Try to extract data from the database connection string. Many are URL based, fall back on
- * some kind of comma-seperated style for the rest. */
+function getSpanContext(span: ReadableSpan): { traceId: string, spanId: string } {
+    const ctx = span.spanContext?.();
+    return {
+        traceId: ctx?.traceId ?? '',
+        spanId: ctx?.spanId ?? '',
+    };
+}
+
+function getParentSpanId(span: ReadableSpan): string {
+    return span.parentSpanContext?.spanId ?? '';
+}
+
+function flattenResourceAttributes(attrs: Record<string, any>): Record<string, AttributeType> {
+    const result: Record<string, AttributeType> = {};
+    function flatten(obj: any, prefix: string) {
+        for (const [key, value] of Object.entries(obj)) {
+            const fullKey = prefix ? `${prefix}.${key}` : key;
+            if (value !== null && typeof value === 'object' && !Array.isArray(value)) {
+                flatten(value, fullKey);
+            } else {
+                result[fullKey] = String(value);
+            }
+        }
+    }
+    flatten(attrs, '');
+    return result;
+}
+
+/** Try to extract data from the database connection string. Many are URL-based; falls back to
+ * semicolon-separated key=value parsing for the rest (e.g. MSSQL style). */
 function parseDatabaseConnectionStringRaw(conn: string): {
     scrubbed: string;
     user?: string;
@@ -467,7 +531,6 @@ function parseDatabaseConnectionStringRaw(conn: string): {
     name?: string;
 } {
     try {
-        // Try URL-style parsing first; scrub the user details if present.
         const url = new URL(conn);
         const user = url.username || undefined;
         const host = url.hostname || undefined;
@@ -483,7 +546,6 @@ function parseDatabaseConnectionStringRaw(conn: string): {
             name: dbName
         };
     } catch {
-        // Not a URL-style string; try semi-structured parsing
         const parts = conn.split(';');
         const kv: Record<string, string> = {};
 
@@ -497,7 +559,6 @@ function parseDatabaseConnectionStringRaw(conn: string): {
         const port = kv['port'];
         const name = kv['database'] || kv['initial catalog'];
 
-        // Reconstruct a scrubbed connection string without credentials
         const scrubbed = Object.entries(kv)
             .filter(([key]) => !['user id', 'uid', 'password', 'pwd'].includes(key))
             .map(([k, v]) => `${k}=${v}`)
@@ -513,8 +574,8 @@ function parseDatabaseConnectionStringRaw(conn: string): {
     }
 }
 
-/** Try pretty hard to get error info out of a span, if it has any. Handles it being there as an event,
- * directly on the span with error semantics, and some other more ad-hoc cases. */
+/** Try pretty hard to get error info out of a span, if it has any. Handles it being there as an
+ * event, directly on the span with error semantics, and some other more ad-hoc cases. */
 function extractErrorInfo(span: ReadableSpan): {
     message?: string;
     type?: string;
@@ -522,7 +583,6 @@ function extractErrorInfo(span: ReadableSpan): {
 } {
     const attrs = span.attributes;
 
-    // Try to extract from a structured 'exception' event, as it should have more detail
     const exceptionEvent = span.events.find(e => e.name === 'exception');
     if (exceptionEvent?.attributes) {
         const message = exceptionEvent.attributes['exception.message'] as string | undefined;
@@ -533,7 +593,6 @@ function extractErrorInfo(span: ReadableSpan): {
         }
     }
 
-    // Fallback to attributes directly on the span.
     const isError = span.status.code === 2;
     const message =
         (attrs['exception.message'] as string | undefined) ??
@@ -553,13 +612,27 @@ function extractErrorInfo(span: ReadableSpan): {
     };
 }
 
-function getTablesWithOperation(
-    tableOps: Record<string, string[]>,
-    operation: string
-): string[] | undefined {
-    const op = operation.toUpperCase();
-    const result = Object.entries(tableOps)
-        .filter(([_, ops]) => ops.includes(op))
-        .map(([table]) => table);
-    return result.length > 0 ? result : undefined;
+export function matchSpanRule(span: ReadableSpan, rule: SpanMatcherRule): boolean {
+    switch (rule.kind) {
+        case 'type': {
+            const kindMap: Record<string, SpanKind> = {
+                'client': SpanKind.CLIENT,
+                'server': SpanKind.SERVER,
+                'internal': SpanKind.INTERNAL,
+            };
+            return span.kind === kindMap[rule.value];
+        }
+        case 'attribute-present':
+            return span.attributes[rule.key] !== undefined;
+        case 'attribute-absent':
+            return span.attributes[rule.key] === undefined;
+        case 'attribute-equals':
+            return String(span.attributes[rule.key]) === rule.value;
+        case 'attribute-not-equals':
+            return String(span.attributes[rule.key]) !== rule.value;
+        case 'all':
+            return rule.rules.every(r => matchSpanRule(span, r));
+        case 'any':
+            return rule.rules.some(r => matchSpanRule(span, r));
+    }
 }