feat: enhance HTTP extractor's script options (#470)

- Allow configuration of the max_allocs and max_const_objects that limit memory allocation in the script env. - Pass the recipe's scope as a script global so that it can be used to construct the asset's URN if needed.
raystack · Feb 2, 2023 · 256db4e · 256db4e
1 parent 7ce6a55
commit 256db4e
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 16 deletions.
diff --git a/plugins/extractors/http/README.md b/plugins/extractors/http/README.md
@@ -57,13 +57,15 @@ source:
 
 ## Inputs
 
-| Key             | Value    | Example                                | Description                                                                                     | Required? |
-|:----------------|:---------|:---------------------------------------|:------------------------------------------------------------------------------------------------|:----------|
-| `request`       | `Object` | see [Request](#request)                | The configuration for constructing and sending HTTP request.                                    | ✅         |
-| `success_codes` | `[]int`  | `[200]`                                | The list of status codes that would be considered as a successful response. Default is `[200]`. | ✘         |
-| `concurrency`   | `int`    | `5`                                    | Number of concurrent child requests to execute. Default is `5`                                  | ✘         |
-| `script.engine` | `string` | `tengo`                                | Script engine. Only `"tengo"` is supported currently                                            | ✅         |
-| `script.source` | `string` | see [Worked Example](#worked-example). | [Tengo][tengo] script used to map the response into 0 or more assets.                           | ✅         |
+| Key                        | Value    | Example                                | Description                                                                                     | Required? |
+|:---------------------------|:---------|:---------------------------------------|:------------------------------------------------------------------------------------------------|:----------|
+| `request`                  | `Object` | see [Request](#request)                | The configuration for constructing and sending HTTP request.                                    | ✅         |
+| `success_codes`            | `[]int`  | `[200]`                                | The list of status codes that would be considered as a successful response. Default is `[200]`. | ✘         |
+| `concurrency`              | `int`    | `5`                                    | Number of concurrent child requests to execute. Default is `5`                                  | ✘         |
+| `script.engine`            | `string` | `tengo`                                | Script engine. Only `"tengo"` is supported currently                                            | ✅         |
+| `script.source`            | `string` | see [Worked Example](#worked-example). | [Tengo][tengo] script used to map the response into 0 or more assets.                           | ✅         |
+| `script.max_allocs`        | `int`    | 10000                                  | The max number of object allocations allowed during the script run time. Default is `5000`.     | ✘         |
+| `script.max_const_objects` | `int`    | 1000                                   | The maximum number of constant objects in the compiled script. Default is `500`.                | ✘         |
 
 ### Request
 
@@ -92,6 +94,29 @@ source:
 
 ### Script Globals
 
+- [`recipe_scope`](#recipe_scope)
+- [`response`](#response)
+- [`new_asset(string): Asset`](#new_assetstring-asset)
+- [`emit(Asset)`](#emitasset)
+- [`execute_request(...requests): []Response`](#executerequestrequests-response)
+- [`exit`](#exit)
+
+#### `recipe_scope`
+
+The value of the scope specified in the recipe (string).
+
+With the following example recipe:
+
+```yaml
+source:
+  scope: integration
+  type: http
+  config:
+  #...
+```
+
+The value of `recipe_scope` will be `integration`.
+
 #### `response`
 
 HTTP response received with the `status_code`, `header` and `body`. Ex:
@@ -150,7 +175,7 @@ asset.data.full_name = "Daiyamondo Jozu"
 Takes an asset and emits the asset that can then be consumed by the
 processor/sink.
 
-#### `execute_request(...requests)`
+#### `execute_request(...requests): []Response`
 
 Takes 1 or more requests and executes the requests with the concurrency defined
 in the recipe. The results are returned as an array. Each item in the array can
@@ -315,8 +340,8 @@ source:
         }
 ```
 
-This would emit a 'User' asset for each user object in `response.data`. Note 
-that the response headers can be accessed under `response.header` and can be 
+This would emit a 'User' asset for each user object in `response.data`. Note
+that the response headers can be accessed under `response.header` and can be
 used as needed.
 
 ## Caveats

diff --git a/plugins/extractors/http/execute_script.go b/plugins/extractors/http/execute_script.go
@@ -20,13 +20,17 @@ import (
 )
 
 func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plugins.Emit) error {
+	scriptCfg := e.config.Script
 	s, err := tengoutil.NewSecureScript(
-		([]byte)(e.config.Script.Source), e.scriptGlobals(ctx, res, emit),
+		([]byte)(scriptCfg.Source), e.scriptGlobals(ctx, res, emit),
 	)
 	if err != nil {
 		return err
 	}
 
+	s.SetMaxAllocs(scriptCfg.MaxAllocs)
+	s.SetMaxConstObjects(scriptCfg.MaxConstObjects)
+
 	c, err := s.Compile()
 	if err != nil {
 		return fmt.Errorf("compile: %w", err)
@@ -41,7 +45,8 @@ func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plu
 
 func (e *Extractor) scriptGlobals(ctx context.Context, res interface{}, emit plugins.Emit) map[string]interface{} {
 	return map[string]interface{}{
-		"response": res,
+		"recipe_scope": &tengo.String{Value: e.UrnScope},
+		"response":     res,
 		"new_asset": &tengo.UserFunction{
 			Name:  "new_asset",
 			Value: newAssetWrapper(),

diff --git a/plugins/extractors/http/http_extractor.go b/plugins/extractors/http/http_extractor.go
@@ -34,8 +34,10 @@ type Config struct {
 	SuccessCodes []int         `mapstructure:"success_codes" validate:"dive,gte=200,lt=300" default:"[200]"`
 	Concurrency  int           `mapstructure:"concurrency" validate:"gte=1,lte=100" default:"5"`
 	Script       struct {
-		Engine string `mapstructure:"engine" validate:"required,oneof=tengo"`
-		Source string `mapstructure:"source" validate:"required"`
+		Engine          string `mapstructure:"engine" validate:"required,oneof=tengo"`
+		Source          string `mapstructure:"source" validate:"required"`
+		MaxAllocs       int64  `mapstructure:"max_allocs" validate:"gt=100" default:"5000"`
+		MaxConstObjects int    `mapstructure:"max_const_objects" validate:"gt=10" default:"500"`
 	} `mapstructure:"script"`
 }
 

diff --git a/plugins/extractors/http/http_extractor_test.go b/plugins/extractors/http/http_extractor_test.go
@@ -293,7 +293,7 @@ func TestExtract(t *testing.T) {
 						body := response.body
 						asset := new_asset("user")
 						// URN format: "urn:{service}:{scope}:{type}:{id}"
-						asset.urn = format("urn:%s:staging:user:%s", "my_usr_svc", body.employee_id)
+						asset.urn = format("urn:%s:%s:user:%s", "my_usr_svc", recipe_scope, body.employee_id)
 						asset.name = body.fullname
 						asset.service = "my_usr_svc"
 						// asset.type = "user" // not required, new_asset("user") sets the field.
@@ -322,7 +322,7 @@ func TestExtract(t *testing.T) {
 				)
 			},
 			expected: []*v1beta2.Asset{{
-				Urn:     "urn:my_usr_svc:staging:user:395f8292-d48b-431b-9e2d-63b3dcd4b986",
+				Urn:     "urn:my_usr_svc:test-http:user:395f8292-d48b-431b-9e2d-63b3dcd4b986",
 				Name:    "Van Stump",
 				Service: "my_usr_svc",
 				Type:    "user",