[SRE-1128] Create generic wrapper for allowing spans to be wrapped

JupiterOne · Jul 27, 2023 · 7fd93d7 · 7fd93d7
1 parent ccf0f4f
commit 7fd93d7
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 3 deletions.
diff --git a/packages/integration-sdk-core/src/types/config.ts b/packages/integration-sdk-core/src/types/config.ts
@@ -1,5 +1,5 @@
 import { IntegrationInstanceConfig } from './instance';
-import { GetStepStartStatesFunction, Step } from './step';
+import { GetStepStartStatesFunction, Step, StepWrapperFunction } from './step';
 import { InvocationValidationFunction } from './validation';
 import {
   ExecutionContext,
@@ -86,6 +86,7 @@ export interface InvocationConfig<
    *
    */
   ingestionConfig?: IntegrationIngestionConfigFieldMap;
+  stepWrapper?: StepWrapperFunction<TStepExecutionContext>;
 }
 
 export interface IntegrationInvocationConfig<

diff --git a/packages/integration-sdk-core/src/types/step.ts b/packages/integration-sdk-core/src/types/step.ts
@@ -240,3 +240,10 @@ export type StepMetadata = StepGraphObjectMetadataProperties & {
    */
   ingestionSourceId?: string;
 };
+
+export type StepWrapperFunction<
+  TStepExecutionContext extends StepExecutionContext,
+> = (
+  step: Step<TStepExecutionContext>,
+  stepFunction: () => Promise<void>,
+) => Promise<void>;
diff --git a/packages/integration-sdk-runtime/src/execution/dependencyGraph.ts b/packages/integration-sdk-runtime/src/execution/dependencyGraph.ts
@@ -15,6 +15,7 @@ import {
   StepExecutionContext,
   StepResultStatus,
   StepStartStates,
+  StepWrapperFunction,
 } from '@jupiterone/integration-sdk-core';
 
 import { timeOperation } from '../metrics';
@@ -60,7 +61,6 @@ export function buildStepDependencyGraph<
 
   return dependencyGraph;
 }
-
 /**
  * This function takes a step dependency graph and executes
  * the steps in order based on the values of their `dependsOn`.
@@ -75,6 +75,7 @@ export function buildStepDependencyGraph<
  * created more leaf nodes and executes them. This continues
  * until there are no more nodes to execute.
  */
+
 export function executeStepDependencyGraph<
   TExecutionContext extends ExecutionContext,
   TStepExecutionContext extends StepExecutionContext,
@@ -90,6 +91,7 @@ export function executeStepDependencyGraph<
   beforeAddRelationship,
   afterAddEntity,
   afterAddRelationship,
+  stepWrapper,
 }: {
   executionContext: TExecutionContext;
   inputGraph: DepGraph<Step<TStepExecutionContext>>;
@@ -102,6 +104,7 @@ export function executeStepDependencyGraph<
   beforeAddRelationship?: BeforeAddRelationshipHookFunction<TExecutionContext>;
   afterAddEntity?: AfterAddEntityHookFunction<TExecutionContext>;
   afterAddRelationship?: AfterAddRelationshipHookFunction<TExecutionContext>;
+  stepWrapper?: StepWrapperFunction<TStepExecutionContext>;
 }): Promise<IntegrationStepResult[]> {
   // create a clone of the dependencyGraph because mutating
   // the input graph is icky
@@ -267,7 +270,10 @@ export function executeStepDependencyGraph<
                   stepId,
                   cached: hasCachePath(stepId).toString(),
                 },
-                operation: () => executeStep(step),
+                operation: () =>
+                  stepWrapper
+                    ? stepWrapper(step, async () => executeStep(step))
+                    : executeStep(step),
               }).catch(handleUnexpectedError),
             );
           } else {

diff --git a/packages/integration-sdk-runtime/src/execution/executeIntegration.ts b/packages/integration-sdk-runtime/src/execution/executeIntegration.ts
@@ -266,6 +266,7 @@ export async function executeWithContext<
         afterAddEntity: createAfterAddEntityInternalHook(logger),
         afterAddRelationship: createAfterAddRelationshipInternalHook(logger),
         dependencyGraphOrder: config.dependencyGraphOrder,
+        stepWrapper: config.stepWrapper,
       });
 
       const partialDatasets = determinePartialDatasetsFromStepExecutionResults(

diff --git a/packages/integration-sdk-runtime/src/execution/step.ts b/packages/integration-sdk-runtime/src/execution/step.ts
@@ -17,6 +17,7 @@ import {
   StepExecutionContext,
   StepResultStatus,
   StepStartStates,
+  StepWrapperFunction,
 } from '@jupiterone/integration-sdk-core';
 
 import { GraphObjectStore } from '../storage';
@@ -48,6 +49,7 @@ export async function executeSteps<
   afterAddEntity,
   afterAddRelationship,
   dependencyGraphOrder,
+  stepWrapper,
 }: {
   executionContext: TExecutionContext;
   integrationSteps: Step<TStepExecutionContext>[];
@@ -61,6 +63,7 @@ export async function executeSteps<
   afterAddEntity?: AfterAddEntityHookFunction<TExecutionContext>;
   afterAddRelationship?: AfterAddRelationshipHookFunction<TExecutionContext>;
   dependencyGraphOrder?: string[];
+  stepWrapper?: StepWrapperFunction<TStepExecutionContext>;
 }): Promise<IntegrationStepResult[]> {
   const stepsByGraphId = seperateStepsByDependencyGraph(integrationSteps);
   let allStepResults: IntegrationStepResult[] = [];
@@ -93,6 +96,7 @@ export async function executeSteps<
         beforeAddRelationship,
         afterAddEntity,
         afterAddRelationship,
+        stepWrapper,
       }),
     );
   }