DEV Community: Dmtro Harazdovskiy

Mastering Cross-Account DynamoDB Access with Terraform and Node.js: A Step-by-step Guide!

Dmtro Harazdovskiy — Mon, 08 May 2023 16:11:44 +0000

Introduction

Establishing cross-account access for DynamoDB can be a complex task, particularly when navigating the intricacies of configuring multiple AWS accounts. However, by employing the power of Terraform and Node.js, you can overcome this challenge and create a robust and efficient multi-account setup. In this article, we will guide you through the process of setting up cross-account access for DynamoDB using Terraform and Node.js, enabling you to manage and fetch data across two separate AWS accounts with ease.

The benefits of this setup are numerous. It not only streamlines your project’s resource management but also significantly improves security by segregating resources between accounts. Additionally, it allows for better scalability, enabling your infrastructure to grow and evolve alongside your project’s needs. So, let’s embark on this exciting journey and explore the world of cross-account DynamoDB access!

Prerequisites

Before we begin, it’s crucial to ensure you have the necessary tools and services at your disposal. For this tutorial, you’ll need the following:

Two AWS accounts: Main with DDB Table and Consumer — can be blank.
Terraform: We will use Terraform to automate the creation and management of AWS resources. To install Terraform, follow the instructions in the official documentation.
Node.js: Our example application for fetching data from DynamoDB will be built using Node. Also serverless will be used to deploy the app.

While this guide assumes you have a basic understanding of AWS Components, Terraform, and Javascript syntax.

Now that you have all the necessary tools and services, let’s move forward and set up our AWS environment for cross-account DynamoDB access.

AWS Setup Overview

Let’s review the AWS setup on a higher level before moving out to the details:

Main account — account holds DDB Table we need to access to

Consumer account — an account that should grant access to DDB Table from the Main account.

Inside a Main AWS Account create a Role (CONSUMER_ACCOUNT_ROLE) that would allow the Consumer Account to assume credentials for the DDB of the Main account.
Inside a Main account create a policy (MAIN_ACCOUNT_DDB_POLICIES_FULL) that would allow access to DDB Table.
Inside a Main account attach MAIN_ACCOUNT_DDB_POLICIES_FULL to CONSUMER_ACCOUNT_ROLE.
Using 1 step, we created a Trust policy rule that allows Consumer accounts to use CONSUMER_ACCOUNT_ROLE.
Inside the Consumer account, we created a policy that would allow Consumer Lambda to assume the secure credentials using CONSUMER_ACCOUNT_ROLE.

Application Layer overview

Since Example Consumer Lambda has access to CONSUMER_ROLE via the attached CONSUMER_ACCOUNT_ACCESS_POLICY we need to issue temporary credentials using Security Token Service (STS) to access DDB Table.
An issued credential that was returned from the STS call must be used to instantiate a DDB client and query data.

Now that we are familiar with all the moving pieces, the funniest part is left — code all of this!

Terraform Config

There are 3 things we need to create for Main Account using Terraform:

aws_iam_role — that Consumer account would use
aws_iam_policy — it would allow the role below to access the DDB table to do specific actions.
aws_iam_policy_attachment — to link roles and policies together.

This module could have 3 params to input:

Consumer account id.
DDB Table name.
Environment name (just in case).

It looks pretty fine for a one-time setup. 1 policy -> 1 attachment -> 1 role. But what if we need to Read from one table and Full access to the other?

Should we duplicate this setup one more time but for the other table?

Let’s create a module that could handle these cases:



# main.tf

resource "aws_iam_role" "ddb_access_role" {
  name = "${var.env}-ddb_cross_acc_access_role-${var.consumer_account_id}"
  description = "Role that provides access to DynamoDB tables for account ${var.consumer_account_id} on ${var.env} env"

  assume_role_policy = jsonencode({
    Version   = "2012-10-17"
    Statement = [
      {
        Action    = "sts:AssumeRole"
        Effect    = "Allow"
        Principal = {
          AWS = "arn:aws:iam::${var.consumer_account_id}:root"
        }
      }
    ]
  })
}

resource "aws_iam_policy" "ddb_table_policy" {
  for_each = { for table_access in var.table_access_list : table_access.table_name => table_access }

  name        = "${var.env}-DynamoDBTableAccess-${var.consumer_account_id}-${each.key}"
  description = "Access policy to DynamoDB ${each.key} table for account ${var.consumer_account_id} on ${var.env} env"

  policy = jsonencode({
    Version   = "2012-10-17"
    Statement = [
      {
        Action    = lookup({
          "Read"            = [
            "dynamodb:BatchGetItem",
            "dynamodb:GetItem",
            "dynamodb:Query",
            "dynamodb:Scan"
          ],
          "ReadWriteUpdate" = [
            "dynamodb:BatchWriteItem",
            "dynamodb:DeleteItem",
            "dynamodb:GetItem",
            "dynamodb:PutItem",
            "dynamodb:UpdateItem"
          ],
          "Full" = ["dynamodb:*"]
        }, each.value.access_type, [])

        Effect    = "Allow"
        Resource  = [
          "arn:aws:dynamodb:*:${data.aws_caller_identity.current.account_id}:table/${each.key}",
          "arn:aws:dynamodb:*:${data.aws_caller_identity.current.account_id}:table/${each.key}/index/*"
        ]
      }
    ]
  })
}

resource "aws_iam_policy_attachment" "ddb_access_policy_attachment" {
  for_each = aws_iam_policy.ddb_table_policy

  policy_arn = each.value.arn
  roles      = [aws_iam_role.ddb_access_role.name]
  name       = "cross_acc_access_policy_attachment-${var.consumer_account_id}-${each.key}"
}

data "aws_caller_identity" "current" {}

The IAM Role

aws_iam_role.ddb_access_role that is created has a name with the following format: [var.env]-ddb_cross_acc_access_role-[var.consumer_account_id].

The assume_role_policy the attribute of the IAM role is set to a JSON-encoded string that defines the permissions for assuming the role. In this case, the policy grants permission to the specified AWS account (arn:aws:iam::${var.consumer_account_id}:root) to assume the role. After applying Terraform you can check out this policy in the Trusted relationships tab on the IAM policy page.

The IAM Policy

aws_iam_policy.ddb_table_policy provides access to a list of DynamoDB tables specified by the var.table_access_list input variable. Each table in the list has an associated access_type attribute that specifies what kind of access (Read, ReadWriteUpdate, Full)should be granted.

The policy grants access to tables based on the access_type attribute. Specifically, the policy allows certain DynamoDB actions (e.g. BatchGetItem, PutItem, etc.) on the tables, as defined by the lookup() function, which returns the appropriate value for the corresponding access_type .

The for_each attribute is used to create a separate IAM policy for each table in the list. Each policy has a name with the following format: [var.env]-DynamoDBTableAccess-[var.consumer_account_id]-${each.key}.

It was required to make a policy to have a unique name for each new table in the list for each account so that we could easily reuse modules for different accounts and environments without worrying about name overlap.

The IAM Policy Attachment

aws_iam_policy_attachment.ddb_access_policy_attachment attaches the IAM policy to the IAM role created earlier. The attachment is created for each IAM policy created as part of the aws_iam_policy.ddb_table_policy resource block.

Finally, the data.aws_caller_identity.current the data source is used to retrieve the AWS account ID of the Terraform user. The account ID is used in the IAM policy resources to grant access to DynamoDB tables for that specific account.

Don’t forget about variables.tf



variable "consumer_account_id" {
  type        = string
  description = "The AWS account ID of the account which gains access."
}

variable "env" {
  type        = string
  description = "Environment name"
}

variable "table_access_list" {
  description = "List of DynamoDB table names and access types"
  type = list(object({
    table_name = string,
    access_type = string  #allowed values are Read, ReadWriteUpdate, and Full
  }))
}

Usage

Go to your Main account terraform setup and use the module we just created:



#.....

module "ddb_cross_acc_access" {
  source = "path/to/module"

  table_access_list = [
    {
      table_name  = "my-table-1"
      access_type = "ReadWriteUpdate"
    },
    {
      table_name  = "my-table-2"
      access_type = "Read"
    }
  ]

  consumer_account_id = "123456789012"

  env = "dev"
}

#.....

Now as you can see we can specify which access level the Main account would Give the DDB Tables for the Consumer account!

Node.js Application Setup

You can use an awesome serverless-iam-roles-per-function serverless plugin to create policy per function to give access to STS: ```yaml

ConsumerAccountExampleFunctions:
handler: src/handler.handler
name: ${self:provider.stage}-${self:service}-consumerAccountExampleFunctions
iamRoleStatements:
- Effect: 'Allow'
Action:
- 'sts:AssumeRole'
Resource:
- 'arn:aws:iam::${env:MAIN_ACCOUNT_ID}:role/${env:ENVIRONMENT}-ddb_cross_acc_access_role-${env:AWS_ACCOUNT_ID}'


It’s a good security practice to specify policy per function since we have strict control of what external sources Lambda has access to.

1. Grant credentials using STS service:
```ts


// get-sts-creds.ts
import {STS} from 'aws-sdk';

type CredentialsParams = {
  RoleArn: string; // Optional - only required if assuming a role
  RoleSessionName: string; // Optional - only required if assuming a role
};

export const getCredentials = async (params: CredentialsParams) => {
  const stsClient = new STS();

  // Assume the IAM role of the account that owns the target resource, if specified
  const assumedRoleObject = await stsClient
    .assumeRole({
      RoleArn: params.RoleArn,
      RoleSessionName: params.RoleSessionName,
    })
    .promise();

  // Get temporary credentials
  const credentials = {
    accessKeyId: assumedRoleObject?.Credentials?.AccessKeyId || '',
    secretAccessKey: assumedRoleObject?.Credentials?.SecretAccessKey || '',
    sessionToken: assumedRoleObject?.Credentials?.SessionToken || '',
  };

  if (!credentials.accessKeyId && !credentials.secretAccessKey && !credentials.sessionToken) {
    throw new Error('Credentials are not defined!');
  }

  return credentials;
};

Use credentials to perform cross-account DDB calls!



import { DocumentClient } from 'aws-sdk/clients/dynamodb';

import { getCredentials } from './get-sts-creds.ts';

const tableName = 'ExampleDDBTableName';

// AWS_ACCOUNT_ID = Consumer account id

const roleArn = arn:aws:iam::</span><span class="p">${</span><span class="nx">process</span><span class="p">.</span><span class="nx">env</span><span class="p">.</span><span class="nx">MAIN_ACCOUNT_ID</span><span class="p">}</span><span class="s2">:role/</span><span class="p">${</span><span class="nx">process</span><span class="p">.</span><span class="nx">env</span><span class="p">.</span><span class="nx">ENVIRONMENT</span><span class="p">}</span><span class="s2">-ddb_cross_acc_access_role-</span><span class="p">${</span><span class="nx">process</span><span class="p">.</span><span class="nx">env</span><span class="p">.</span><span class="nx">AWS_ACCOUNT_ID</span><span class="p">}</span><span class="s2">;

const roleSessionName = 'my-session-name';

export const handler = async () => {

  const credentials = await getCredentials({ RoleArn: roleArn, RoleSessionName: roleSessionName });

const ddbClient = new DocumentClient({

    credentials

  });

const params = { TableName: tableName };

  const data = await ddbClient.scan(params).promise();

return data

};

Summary

Now you know how to create terraform module for granular cross-account DDB permission:

On the Main account create a Role with Consumer Account trusted relationship
On the Main account create a policy with DDB accesses attached to the Role
On Consumer Account create a policy to allow using the Mains’ Account Role
Use STS client in your application to grant temporary credentials from the Main account role
Perform DDB Queries with granted credentials!

Next-level performance with Distributed Map!

Dmtro Harazdovskiy — Thu, 02 Feb 2023 09:04:20 +0000

Problem

Let's consider such a case — you are developing functionality that processes thousands of records in any data storage and your weapon of choice is the AWS Step Function service.

One of the most powerful tools you can use there is Map task type since there you can maximize your function load and process chunks of data concurrently with maxConcurrency param.

Let’s assume the logic of the steps:

CaptureInput — will accept some params from the user/service fetch data from storage and pass it to be processed
ProcessData — consume data process it and update/delete whatever logic requires

We also should keep in mind the main limitation of Inline Map for the step function

Max payload size — 256 KB (whole payload, not just one array item)
Max execution history — 25,000 events
Accepts only a JSON array as input.

Code

An obvious part of the work that the CaptureInput Lamda should do is to fetch and pass some data into the Map state in the form from-to JSON Array:

// capture-input/handler.js
exports.handler = async ({number = 10, chunk = 1}) => {

  return paginate(number, chunk)
}

function paginate(number, chunk){
  const result = [];
  const numPages = Math.ceil(number / chunk);
  for (let i = 0; i < numPages; i++) {
    result.push({
      from: i * chunk,
      to: Math.min((i + 1) * chunk, number),
    });
  }
  //  [{from: 1, to: 2}, {from: 2, to: 3}.....]
  return result;
};

For simplicity, imagine we are not passing raw ids but only a form — to cursor for a chunk of data. Lambda receives the following params:

number — the number of processed data items
chunk — count of the size that each map would process

The ProcessData Lambda would log it and somehow process the data:

// process-output/handler.js
exports.handler = (event) => {
  console.log(JSON.stringify(event))
  // ...do some processing
}

For deploying the step function serverless-step-function the plugin is used, the YAML file would look like this:

#serverless.yml
service: limit-payload-step-function

provider:
  name: aws
  runtime: nodejs18.x
  stage: dev
  region: us-east-1
  accountId: 670313496061
  architecture: arm64

package:
  individually: true

plugins:
  - serverless-step-functions

functions:
  captureInput:
    handler: capture-input/handler.handler
    name: ${self:provider.stage}-${self:service}-captureInput
    timeout: 900
    memorySize: 1024
    tags:
      name: captureInput

  processOutput:
    handler: process-output/handler.handler
    name: ${self:provider.stage}-${self:service}-processOutput
    timeout: 900
    memorySize: 128
    tags:
      name: processOutput

stepFunctions:
  stateMachines:
    ProcessDataStateMachine:
      name: process-data-state-machine
      definition:
        StartAt: CaptureInput
        States:
          CaptureInput:
            Type: Task
            Resource: arn:aws:lambda:${self:provider.region}:${self:provider.accountId}:function:${self:provider.stage}-${self:service}-captureInput
            Next: MapData
          MapData:
            MaxConcurrency: 5
            Type: Map
            InputPath: $
            Iterator:
              StartAt: ProcessData
              States:
                ProcessData:
                  Type: Task
                  Resource: arn:aws:lambda:${self:provider.region}:${self:provider.accountId}:function:${self:provider.stage}-${self:service}-processOutput
                  End: true
            End: true

Then run to deploy:

serverless deploy --aws-profile=personal

Great! We have our function up and running and everything seems fine:

Until we try some bigger input to be passed to Map: Let’s try running it with:

{
  "number": 10000,
  "chunk": 5
}

And now we are in trouble:

The state/task returned a result with a size* exceeding the maximum number of bytes service limit*.

The output of CaptureInput is bigger than the step function can process (256 KB). Even if we try somehow to optimize params passing we still can hit 25k execution events when we would need to add other steps into the MapData task. And that is the time for Distributed Map to rescue us!

Inline VS Distributed

There is a neat illustration from AWS that outlines the work of the **Inline Map **we used before:

Runs a set of steps for each element in a JSON array passed from the previous state.
Step Functions adds the execution history of each Map iteration to the state machine’s execution history.
Map iterations run in parallel, with limited concurrency after 40 iterations.

As we can see we are pretty much limited using the default Map, and that's why the AWS team created a New Distributed Mode that would overcome the limits of the previous one.

Runs a set of steps for each item in a JSON array, S3 object list, or CSV file.
Map iterations are fanned out as separate child executions with their own execution history, enabling very high concurrency.
Iterates over and processes millions of items by running thousands of parallel child executions.
Can cascade multiple Distributed Map states and batch items to create map-reduce type jobs.
A Distributed Map state launches a Map Run resource, which you can access through the console, APIs, and SDK.

As you can see basically every one of these items is going to make work with larger datasets easier and faster. By checking out Step Function Studio, you can see that there are 3 main types of Distributed Map:

In this case, I guess the most suitable is the second option: Let Map iterate over JSON that is stored in S3. This is an example of how the step function can be realized using this approach:

CaptureInput Lambda receives params and generates from — to pairs
Saves JSON with the pairs into Bucket
Pass the Key (path) of JSON and Bucket name to Distributed Map
Map will automatically prefetch JSON and loop over each item in it, the item in the array is going to be passed directly into the Map body (CaptureInput)
CaptureInput receive the item of JSON array and process it like input for further data manipulations.

Coding the solution

Creating bucket

As was shown in the diagram above we require AWS S3 Bucket to store items for Map. You can do it manually via the console or by adding Resource section in serverless.yml file:

resources:
  Resources:
    PayloadBucket:
      Type: AWS::S3::Bucket
      DeletionPolicy: Retain
      Properties:
        BucketName: step-function-payloads

Adding S3 permission

Then since CaptureInputfunction required S3 permission, I'm using serverless-iam-role-per-function plugin to do that, but you can also do it manually.

yarn add serverless-iam-roles-per-function

Add it to the plugins section:

provider:
  # .....
  environment:
    BUCKET_NAME: step-function-payloads
  iam:
    role:
      statements:
        - Effect: 'Allow'
          Action:
            - 's3:*'
          Resource:
            - 'arn:aws:s3:::step-function-payloads/*'
            - 'arn:aws:s3:::step-function-payloads'
  # .....
  plugins:
    - serverless-step-functions
    - serverless-iam-roles-per-function
  # .....
  custom:
    serverless-iam-roles-per-function:
      defaultInherit: true

This config will attach iam role with access to step-function-payload to each function as a default. Also, it will help create granular permissions in further development.

Saving Iterable JSON

Now we need to write our output JSON into S3 so that later Map state will iterate over it. It will be done in CaptureInput lambda:

// capture-input/handler.js
const AWS = require('aws-sdk')
const s3 = new AWS.S3();
const payloadFolder = 'temp-sf-payload'
const Bucket = process.env.BUCKET_NAME;

exports.handler = async ({number = 10, chunk = 1}) => {

  const data = paginate(number, chunk)

  const path = `${payloadFolder}/${Date.now()}.json`

  const res = await s3.putObject({
    Bucket,
    Key: path,
    Body: JSON.stringify(data)
  }).promise()

  console.log('Saved output: ', {Bucket, path})

  return {bucket: Bucket, path}
}

function paginate(number, chunk){
  const result = [];
  const numPages = Math.ceil(number / chunk);
  for (let i = 0; i < numPages; i++) {
    result.push({
      from: i * chunk,
      to: Math.min((i + 1) * chunk, number),
    });
  }

  return result;
}

We just save the from-to[] payload into an s3 file and pass bucket and path to the Map State. Now we need to consume it!

Distributed Map Definition

There are not many changes in a Step function definition:

    # ...
    MapData:
      Type: Map
      MaxConcurrency: 5
      ItemReader:
        Resource: arn:aws:states:::s3:getObject
        ReaderConfig:
          InputType: JSON
        Parameters:
          # Assign output of CaptureInput lambda for Map to iterate over
          Bucket.$: $.bucket
          Key.$: $.path
      ItemProcessor:
        ProcessorConfig:
          Mode: DISTRIBUTED
          ExecutionType: STANDARD
        StartAt: ProcessData
        States:
          ProcessData:
            Type: Task
            Resource: arn:aws:lambda:${self:provider.region}:${self:provider.accountId}:function:${self:provider.stage}-${self:service}-processOutput
            End: true
      End: true

There are two new key components here:

ItemReader used to specify the source of iteration and its type - either JSON CSV or S3 Folder
ItemProcessor — the same as Iterator in inline mode but with specifying nested workflow type

I’ve added all the references regarding the new Distributed mode including the blueprints repository in the last section of the article.

Step function Role

The last step is to add a custom role to the step function. New Map mode requires Map state to have getObject permissions to read from bucket we specified.

resources:
  Resources:
    ProcessDataStateMachineStepFunctionRole:
      Type: "AWS::IAM::Role"
      Properties:
        RoleName: ${self:provider.stage}-ProcessDataStateMachine-StepFunctionRole
        AssumeRolePolicyDocument:
          Statement:
            - Effect: Allow
              Principal:
                Service:
                  - states.amazonaws.com
              Action:
                - sts:AssumeRole
        Policies:
          # Allow invoking lambdas as a tasks in step function
          - PolicyName: ${self:provider.stage}-${self:service}-lambda-invoke
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
                - Effect: 'Allow'
                  Action:
                    - 'lambda:*'
                  Resource:
                    - 'arn:aws:lambda:*:*:function:${self:provider.stage}-*'
          # Allow starting execution
          - PolicyName: ${self:provider.stage}-${self:service}-start-execution
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
                - Effect: 'Allow'
                  Action:
                    - 'states:StartExecution'
                  Resource:
                    - arn:aws:states:${self:provider.region}:${aws:accountId}:stateMachine:${self:stepFunctions.stateMachines.ProcessDataStateMachine.name}
          # Allow logging
          - PolicyName: ${self:provider.stage}-${self:service}-default-logging
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
                - Effect: 'Allow'
                  Action:
                    - 'logs:CreateLogGroup'
                    - 'logs:CreateLogStream'
                    - 'logs:PutLogEvents'
                  Resource:
                    - 'arn:aws:logs:*:*:log-group:/aws/lambda/*'

                - Effect: 'Allow'
                  Action:
                    - 'xray:PutTraceSegments'
                    - 'xray:PutTelemetryRecords'
                  Resource:
                    - '*'
          # Full access to S3 bucket
          - PolicyName: ${self:provider.stage}-${self:service}-read-s3-bucket
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
                - Effect: 'Allow'
                  Action:
                    - 's3:*'
                  Resource:
                    - arn:aws:s3:::step-function-payloads/*

Demo

Started function with these params:

{
  "number": 100000,
  "chunk": 5
}

Map successfully started!

Now we can monitor the progress in a separate menu:

This Map is processing 20k steps in a snap of a finger, and it’s not the last feature of Distributed mode. We also can increase MaxConcurrency and use ItemBatcher parameters to make it even faster!

Summary

Now you know how to overcome the limitations of the Inline Map of your step function workflow with Distributed Map task state. Check out the blueprint in the section below!

Thanks for reading!

Resources

Repository: Step Function Distributed Map Example
DOC: Map state processing modes
DOC: Using Map state in Distributed mode
DOC: ItemBatcher

Test Postgres using only this jest plugin (no DB mocks/Docker)!

Dmtro Harazdovskiy — Thu, 12 Jan 2023 14:57:28 +0000

Originally published here

All of us love working with Postgres. It’s a powerful database with lots of functionality. You can write very complex queries that would join, and group many tables.

I guess with great power comes great responsibility and the bigger the project gets the easier it is to break something. That’s why you should test your Postgres business logic.

Jest is the weapon of choice for this article. We can mock everything with jest? But how to mock all the variety of your Postgres calls? Is it even necessary? Which plugin to choose and what is better for your needs, let’s find out!

Available approaches

To test Postgres with the Jest framework you can do all of the heavy liftings yourself.

Before running tests pre-install docker, spin up the Postgres container with the required params, create a schema, and after tests are finished, stop Postgres, delete the container — done. Then put all of this in a bash script and that's it.

As for me this approach sounds awful and would take at least half a minute to start and shut down properly on your local machine. You would not be able to iterate quickly and run it every time you made a change to query since it would take CPU resources -> your time.

Why should you consider it — you would have a fully capable database up and running for your tests. However with a reasonable effort.

Even easier option — look up some available npm plugins for Postgres there. But basically, all of them have limited capabilities. Sooner or later the queries, you test may become more complex. Therefore relying on lib that does not provide you with limited functionality is a bad idea.

Now you may ask, ok what to do? We have a simple answer…

@shelfio/jest-postgres

The current company I’m working for is using SQL a lot and we love Jest too! That’s why we spend some time creating open-source jest plugins.

Why is this one better?

This plugin is fully self-written, uses up-to-date dependencies, with pretty straightforward functionality.

Apart from the approaches below this plugin utilizes a full-featured, up-and-running Postgres instance on your machine. You can even choose the version and preconfigure scheme.

Since we are using this plugin a lot, we care about the work of this plugin in the first place and can recommend it with no doubts.

Downsides

Currently, only Linux(Ubuntu) and macOS are supported and tested out. Also, you should run your system as a root used in the case of macOS to prevent using sudo. For Linux, all the commands are executed under the postgres user that is installed after apt-get.

How to set up and use it?

Install a plugin:

npm i @shelfio/jest-postgres --save-dev

Add configuration to jest that would run the plugin only for specific files that have Postgres logic:

module.exports = {
  preset: '@shelf/jest-postgres',
  testRegex: `.*(?<!integration)\\.pg\\.test\\.js$`,
  maxWorkers: 1
};

This means that when you are testing create.js you can create create.pg.test.js the file that would test the logic and run Postgres for it.

Create jest-postgres-config.js a file that would contain configs.

const cwd = require('cwd');

module.exports = {
  seedPath: `${cwd()}/schema.sql`,
  version: 12,
  port: 5555
};

**Create schema.sql — it **is a required file for seeding Postgres with the schema of tables before running the tests.

create schema test;

create table test.model1
(
  id         varchar(36) not null,
  type       text,
  text       text,
  vector     double precision[],
  json       jsonb,
  updated_at timestamp,
  someBool   bool,
  constraint message
    primary key (id)
);

create index some_cool_index on test.model1 using btree(id);

Suggestion: just create a sctructure of the table in this file to keep it crear. Make insertion in code directly.

Profit, write some tests!

@shelfio/postgres-local under the hood

To run Postgres locally using node.js we followed the path of our first child elasticsearch-local (read more about it here).

The npm plugin has two functions — to start and stop Postgres.

When we run start() :

Raw bash commands executed via execSync download and install Postgres on your machine.
Creating /tmp/postgres-local directory
Initializing that temporary directory
Start Postgres

When we run stop() :

Stop Postgres
Delete /tmp/postgres-local

In the case of Linux users, we need to run all the commands as a postgres user to avoid permissions issues, also we need to create a user and database since ubuntu apt-get does not make it by default.

Conclusion

Now you have an alternative, buttle-proved approach for testing Postgres business logic. View examples of ready tests here.

Hope this plugin will help you test your queries better and iterate development faster!

Read me on Twitter!

How to mock Elasticsearch with Jest?

Dmtro Harazdovskiy — Fri, 02 Dec 2022 20:37:10 +0000

Originally published here

Intro

I guess you are probably using Elasticsearch or planning to since you've opened this article 🙂 Another assumption - you definitely found it a great idea to lock your functionality with some integration tests. And it really is!

The current company I’m working for has 90+% of code coverage with both unit and integration tests! I would recommend everyone cover their code base with tests since as one wise man said:

[Without unit tests] You’re not refactoring, you’re just changing shit.
— Hamlet D’Arcy

Setup

Imagine you have a simple RESTful server with some logic that uses Elasticsearch. In the current showcase — CRUD server.

const Hapi = require('@hapi/hapi');
const Qs = require('qs');
const {createHandler} = require("./create/index.js");
const {readAllHandler, readHandler} = require("./read/index.js");
const {updateHandler} = require("./update/index.js");
const {deleteAllHandler, deleteHandler} = require("./delete/index.js");

const init = async () => {

    const server = Hapi.server({
    port: 3000,
    host: 'localhost',
    query: {
        parser: (query) => Qs.parse(query)
    }
    });

    server.route({
    method: 'POST',
    path: '/',
    handler: createHandler
    });

    server.route({
    method: 'GET',
    path: '/{id}',
    handler: readHandler
    });

    server.route({
    method: 'GET',
    path: '/',
    handler: readAllHandler
    });

    server.route({
    method: 'PATCH',
    path: '/{id}',
    handler: updateHandler
    });

    server.route({
    method: 'DELETE',
    path: '/{id}',
    handler: deleteHandler
    });

    server.route({
    method: 'DELETE',
    path: '/',
    handler: deleteAllHandler
    });

    await server.start();

    server.events.on('log', (event, tags) => {
    console.log({event}, {tags})
    if (tags.error) {
        console.log(`Server error: ${event.error ? event.error.message : 'unknown'}`);
    }
    });
    console.log('Server running on %s', server.info.uri);
};

process.on('unhandledRejection', (err) => {

    console.log(err);
    process.exit(1);
});

init();

Now you need to cover the logic of each route with some tests to lock functionality and prevent business logic from being broken.

One obvious but not simple solution is to use Docker and spin up Elastic for tests every time.

However, does it worth it? I mean you really want to have a longer spin-up time for the pipeline environment? Maybe there are already-built solutions for that?

This plugin downloads and caches Elasticsearch binary when jest starts, then the plugin automatically starts Elastic on defined ports and tears it down when tests are done.

How to add tests?

I have an example of with/without jest tests setup in the pull request so you can compare all the changes. But now let’s go through it step by step.

1. Install additional modules

yarn add --dev jest @shelf/jest-elasticsearch @types/jest

2. Add jest-config.js

touch jest.config.js

module.exports = {
    preset: '@shelf/jest-elasticsearch',
    clearMocks: true,
    collectCoverage: true,
    coverageDirectory: "coverage",
    coverageProvider: "v8"
};

Alternatively, you can generate jest config on your own using the CLI tool.

3. Add jest-es-config.js for the plugin

touch jest-es-config.js

const {index} = require('./src/elastic.js');

module.exports = () => {
    return {
    esVersion: '8.4.0',
    clusterName: 'things-cluster',
    nodeName: 'things-node',
    port: 9200,
    indexes: [
        {
        name: index,
        body: {
            settings: {
            number_of_shards: '1',
            number_of_replicas: '1'
            },
            mappings: {
            dynamic: false,
            properties: {
                id: {
                type: 'keyword'
                },
                value: {
                type: 'integer'
                },
                type: {
                type: 'keyword'
                },
                name: {
                type: 'keyword'
                },
            }
            }
        }
        }
    ]
    };
};

4. Extend package.json script to run tests

{
    "scripts": {
    "test": "jest"
    "serve": "node src/index.js"
    }
}

5. Tune elastic client

const dotenv = require('dotenv')
dotenv.config()
const {Client} = require('@elastic/elasticsearch');

module.exports.client = new Client({
    node: process.env.NODE_ENV === 'test' ? 'http://localhost:9200' : process.env.ES_URL
})

module.exports.index = 'things'

Add a condition that will use NODE_ENV to connect to local spined-up elastic whenever we are running tests.

Profit!

Now all the things are ready to write and run tests. All routes are fully covered and stored here:
elastic-jest-example

As an example let’s cover create function business logic.

const {ulid} = require('ulid');
const {client, index} = require("../elastic.js");

module.exports.createHandler = async (request, h) => {
    if (Object.keys(request.payload))
    try {
        const res = await this.create(request.payload)
        return h.response(res).code(200);
    } catch (e) {
        console.log({e})
        return h.response({e}).code(400);
    }
}

// let's cover this function with some tests
module.exports.create = async (entity) => {
    const {
    type,
    value,
    name,
    } = entity;

    const document = {
    id: ulid(),
    type: type.trim().toLowerCase(),
    value: +value.toFixed(0),
    name: name.trim()
    }

    await client.index({
    index,
    document
    });
    return document.id
}

Create a test file and add a couple of statements.

touch  src/create/index.test.js

const {create} = require("./index.js");
const {client, index} = require("../elastic");

describe('#create', () => {

// clear elastic every time before running it the statement.
// It's really important since each test would be idempotent.
    beforeEach(async () => {
    await client.deleteByQuery({
        index,
        query: {
        match_all: {}
        }
    })
    await client.indices.refresh({index})
    })

    it('should insert data', async () => {
    expect.assertions(3);
    const res = await create({type: 'some', value: 100, name: 'jacket'})
    await client.indices.refresh();
    const data = await client.search({
        index,
        query: {
        match: {
            "id": res
        }
        }
    })

    expect(res).toEqual(expect.any(String))
    expect(res).toHaveLength(26);
    expect(data.hits.hits[0]._source).toEqual({
        "id": res,
        "name": "jacket",
        "type": "some",
        "value": 100
        }
    );
    })

    it('should insert and process the inserted fields', async () => {
    const res = await create({type: 'UPPERCASE', value: 25.99, name: ' spaces '})
    await client.indices.refresh();
    const data = await client.search({
        index,
        query: {
        match: {
            "id": res
        }
        }
    })
    expect(data.hits.hits[0]._source).toEqual({
        "id": res,
        "name": "spaces",
        "type": "uppercase",
        "value": 26
        }
    );
    })
});

A basic testing flow for each business logic function can be simply described like this:

insert data-> run tested function -> check outputs -> clear data -> repeat

Data insertion/deletion can be improved by unifying them into helpers and using additional mooching libs.

The elastic teardown is managed by @shelf/jest-elasticsearch lib itself.

One more good convention to follow is to cover each function you are testing with describeblock, so that later you can easily run a specific test with IDE helper without reruning whole suite:

Resources

Now you know how to test your Elasticsearch queries using jest.
jest-elasticsearch

Also, you have a small blueprint repo with ready to use setup!
elastic-jest-example

Hope this article will help you set up and test your elastic project!

Want to connect?

Follow me on Twitter!

Optimizing massive MongoDB inserts, load 50 million records faster by 33%!

Dmtro Harazdovskiy — Sun, 16 Oct 2022 17:55:11 +0000

Intro

Original article here

Working with 100k — 1m database records is almost not a problem with current Mongo Atlas pricing plans. You get the most out of it without any hustle, just by enough hardware, simply use indexes and pagination.

But imagine your project recently got a massive customer dataset for the new ETL pipeline. Now you need to make a one-time upload script and later deal with KPI calculation for all (or part) of the dataset.

What about having a dataset larger than 1 million records? 7 million? Or even more than 50🍋! Not all of the standard approaches would work efficiently for such volumes.

Therefore, let’s take a look at how we can handle this volume by gradually increasing dataset size and tackling performance issues!

Disclaimer — All the code is located here — Github

Setup

To make this experiment simple, I choose to use the M30 *Dedicated *cluster on Mongo Atlas. There is a serverless Mongo plan but let’s omit it for this experiment.

Also, there is an already built-in tool in Atlas for loading dataset samples, but I recommend following along with a *rough *approach to doing custom scripts. You might have a self-hosted MongoDB server without Atlas, or atlas CLI can not support your dataset extension type, moreover, you will control the flow of data on your own.

One more thing to mention is internet speed. Experiments were completed on my local Macbook Pro with a 25–35 Mbps wifi connection. The results of the current experiment may differ from running on a prod-ready EC2 instance since the network is much better there.

Where could you find some sample data?

To play around with data generation and make a custom dataset I can recommend using — https://generatedata.com/. I’ve used it to generate 1🍋 records of the data. At the moment of writing this article, the basic yearly plan costs 25$ and you would not regret it.

This tool works in the browser so you need to be ready for heavy CPU load during generation.

Don’t forget to *strip whitespace from generated content *tick to make your file ~20% smaller

It took me ~3h to generate a 1🍋 dataset of 254MB size. It was a JSON file that we are going to use later on to load it into Mongo.

You will find links to a ready to use datasets in the insertion paragraph.

I recommend using this tool for no more than 2m data generation since it took a reasonable time and your laptop CPU.

Want more data with less effort? Go to — Kaggle

A colleague of mine recommended me a great DS platform with already prepared datasets for testing data. After digging it, I found two datasets that suit the best.

Yelp Dataset

You can review all the contents of it, but generally, there is one big 5GB JSON file — yelp_academic_dataset_review.json it contains 6,9 🍋 JSON records. Use Dadroid to review massive JSON files.

2. Steam Store Apps Reviews

This is a big one — 45GB of steam reviews. I did not know how many of the records are inside while downloading it since all the files were in the paginated folders. Later on, we will find out how to write a script, to summarize, the count of all the records. The exact number of them is going to be revealed, but believe me, this can be called Big(little) Data.

How to insert the data?

I recommend putting your dataset outside of the working folder. It will prevent your code editor from indexing all of that tree of files and slowing down the text search.

Inserting 1 million records

Let’s start with inserting 1m JSON using just a simple updateMany operation. Template of generated data:

Data sample:

I decided to try inserting all the docs in one piece without any chunks:

In this script, we are simply calling an anonymous function that will be executed immediately. I’m going to use a native mongodb driver since all other libs are built on top of that. On line 3 I’ve imported promisified version fs to use await in JSON parse input.

It was interesting to try what the free plan is capable of, so this script was run on the free Atlas cluster.

> node insert/1m/index.js                                                                                                                               
Reading json: 3.279s
Connected successfully to server
Started insertion process successfully to server
Inserting records: 2:18.415 (m:ss.mmm)

Mongo did not even feel that! It took node ~3 seconds to read the file and about 2,5 minutes for Mongo to write 1m records. Also, there were no more than 3 connections created for that.

What about the M30 cluster? Let’s test how it would perform running the same script.

node insert/1m/index.js                                                                                                                               
Reading json: 2.874s
Connected successfully to server
Started insertion process successfully to server
Inserting records: 2:20.470 (m:ss.mmm)

No changes in performance here. The current approach is the *fine *approach, you can use it for one time insert. But the easy way is not our one. Moving on to the next section.

Inserting 7 million

Now let’s work with 7 million JSON. Records example:

One review here is less massive than one record in a previously generated dataset, nevertheless, this volume is enough to put some stress both on Node and Mongo.

Would the previous approach work out for such a data volume? Let’s test. I’ve changed a path to a new dataset and started my script.

node insert/1m/index.js

RangeError [ERR_FS_FILE_TOO_LARGE]: File size (5341868833) is greater than 2 GB
    at new NodeError (node:internal/errors:372:5)
    at readFileHandle (node:internal/fs/promises:377:11)
    at async /Users/mongo-performance/insert/1m/index.js:16:36 {
  code: 'ERR_FS_FILE_TOO_LARGE'

2GB is a limit size for files that can be stored in a Node.js buffer entity. That’s why we should use Streams here!

A quick reference from the doc:

A stream is an abstract interface for working with streaming data in Node.js. There are many stream objects provided by Node.js. For instance, a request to an HTTP server and process.stdout are both stream instances. Streams can be readable, writable, or both. All streams are instances of EventEmitter.

Now we need to make some rough changes to the script.

Line 23 — we have created a readable stream and used stream-json lib to pipe it. This was done since consuming a raw Buffer from Node.js would make the process of reading more complicated. We need somehow to stringify that Buffer and parse JSON out from string on our own — stream-json is going to do this for us.

Line 25 — stream is consumed via Event Emitter. The function is subscribed to a *data *event and each data variable that contains exactly one JSON record from our 7m dataset file.

Line 27 — all the records are pushed into one accumulator array — arrayToInsert

Line 29 has a remainder operator — it is used just for triggering if branch whenever a new 100k of records are pushed to the array.

Line 30 — when we entered the if statement branch, pipeline should be paused to prevent new array pushes. Nothing new further — insertMany and wait 100ms after each call. Resume pipeline and do it until end the event fires to exit the script.

In this example, we are inserting 100k sequentially with a slight delay. Testing time!

node ./insert/insert-7mil-m30.js

Connected successfully to server
arrayToInsert size - 354.85045185293717 mb
Inserting time - 100000: 45.277s
--------------


.................


arrayToInsert size - *142*.*23186546517442* mb
Inserting time - *6900000*: *25*.911s
--------------

Operation took -  *1749*.*997*  seconds = ~29 minutes


Process finished with exit code 0

It took about 29 minutes to insert. And now it seems to be reasonable too much time to run.

As you can see here we were using about 30% of the system CPU and 66 IOPs were the highest value during the insertion.

According to 29 minutes per 7m docs, when you need to put 50m records it would take almost an infinity (~3,5h to run)😊.

We definitely can improve it by adding some parallel insert operations, and probably increasing the 100k chunk size to let more parallel operations be handled:

Lines 36 - 42 — radash.parallel is used to make 5 parallel insertions with a chunk of 20k records. Chunks were created by lodash.chunk method.

Later on, this approach will be reused for larger insertions so take your time to figure out how it is working.

Results:

node insert/7m/7m-m30-per200k-5X20k.js                                                                                                                
Connected successfully to server
arrayToInsert size - 269.14077478740194 mb

id: 1664486093579 - stats: size 26.909 mb, records: 20000 - took: : 7.258s
id: 1664486093497 - stats: size 26.581 mb, records: 20000 - took: : 11.450s

................

id: 1664486950224 - stats: size 27.520 mb, records: 20000 - took: : 6.973s
id: 1664486951694 - stats: size 27.230 mb, records: 20000 - took: : 5.562s
--------------

Operation took -  867.229  seconds // ~14 minutes

As we can see here performance improved to 50%! I guess It’s much better now. What about Mongo?

System CPU has risen to 40%, it’s 10% more than the sequenced 100k insertion.

IOPS also almost doubled from 50 to 100, and connections also increased from 40 to ~60.

I’m sure there is a big way to improve here but this result is fine for most cases. Now using this parallelization we can get on our Big-little data.

Insertion of more than 50 🍋

To make this test happen we need lots of data, right? So that I’ve chosen the steam dataset — user reviews of about 50000+ games on Steam. Also, I’ve mentioned that before messing around with the insertions script we need just to be aware of how exactly records are there.

Quick overview of the dataset folders

This dataset consists of folders, each of them containing 10 to 2000 JSON files. Each JSON may contain any number of reviews.

In the current example, the script should primarily iterate over each folder, read the file and save a counter of reviews array length there. Take a look at one review inside of JSON file:

As you may see this record is the biggest document among compared datasets.

One thing I want to make in this script is to add some input parameters for the limitation of our folder's pagination. It will let us test the script on a subset of data and parallelize insertion by selecting a range of folders for each process. Now let’s check the script itself!

Line 2 — minimist package was used to make it easy to read CLI parameters inside of the node process. It can be now used like this* node inde.js --from=1 --to=55*

Line 11 — lists directories using fs.promises.readdir

Line 13 — select paths that are inside the datasets’ folder and cut them to the exact range we will be working with.

Line 18 — every path that is inside of a loop contains a folder without JSON files and the same operation should be done here — list all JSON files with fs.promises.readdir.

innerPaths contains all paths to JSON files we need to iterate over. I’ve used radash.parallel to read files in dir quicker. Also, no streams here since almost every file is not bigger than 100kb.

Line 27 — parsing JSON out of Buffer string and reading length of reviews the key. Summing doneReadingReview var.

There is an important console log on line 33. Since data insertion can be parallelized later, we need to know how many folders were read to accumulate the next million records. Time to run this weirdo script!

node insert/50m+/read-size-of-dataset/index.js --from=1 --to=60000

 All paths 51936
done paths - 1953 done reviews 1000066
done paths - 3339 done reviews 4000078

........

done paths - 44528 done reviews 59000025
done paths - 51410 done reviews 63000010

Reviews 63199505
Reviews 63199505
Operation took -  429.955  seconds // ~ 7 minutes

My laptop fans reached max RPS while reading this dataset. When do we have here:

51936 folders
incredible 63 199 505 records!
read time of 7 minutes (and that’s just counting the length)

Let’s count how long it would take to insert 63🍋 using our best results on previous steps (7🍋/14m) — 2h (about 126 minutes). I don’t think we can give up on that.

Insertion of 63 million records

As we saw before — power is in parallelization. Using the previous section's log output we should divide our logs into 6 even parts. Then reuse done paths parts as --from and — toparams inputs for our next script. It will result in these pairs:

1 - 14454

14454 - 25104

25104 - 28687

28687 - 34044

34044 - 42245

42245 - 52000

Since we already prepared our counting script to use ranges for the folder, the next step is just to add the insertion. Let’s review the ready script and its main parts!

The are 3 main parts that were changed:

Lines 46–47 — we are simply pushing new reviews into an array

Lines 55–65 — when there are more reviews than the threshold they are inserted into Mogno, and when the loop is on the last iteration we forcibly insert data to make sure all data is inserted

Lines 87–100 — insertion function that is chunking all reviews per 20k and parallelly inserts it

Since we already defined pairs for chunks of the same dataset size earlier I guess we can try running the insertion script in 6 parallel processes. To make things simple let’s just open 6 terminal windows and run it! You can always use built-in worker threads modules in Node.js and the pm2 tool, but our experiment can satisfy the most primitive solution. Testing time!

Results of each parallel execution:

--from=*1* --to=*14454*
Done reading reviews:  *10 963 467
*Script: took -  *4963*.*341*  seconds ~ 1h 22m

--------

--from=*14454* --to=*25104
*Done reading reviews:  *10 498 700
*Script: took -  *5016*.*944*  seconds ~ 1h 23m

--------

--from=*25104* --to=*28687*
Done reading reviews:  *10 874 942
*Script: took -  *5050*.*838*  seconds ~ 1h 24m

---------

--from=*28687* --to=*34044
*Done reading reviews:  *11 659 485
*Script: took -  *5088*.*016 * seconds ~ 1h 24m

---------

--from=*34044* --to=*42245
*Done reading reviews:  *10 044 895
*Script: took -  *4796*.*953*  seconds ~ 1h 19m

---------

--from=*42245* --to=*52000*

Done reading reviews:  *9 158 016
*Script: took -  *4197*  seconds ~ 1h 9m

It took about 1 hour and 24 minutes to insert 63🍋 of records! It’s 33% faster than our previous prediction for sequential insert! We loaded almost 45GB of data in under one and a half an hour. What about Mongo?

These inserts loaded the cluster pretty heavily. Max System CPU ~65%, process CPU was about 40% during the whole period. Connections stayed at 100 almost the entire time and IOPS averaged 125. I guess these results are good enough now!

One more thing

All the approaches described are general performance improvement. But there is one more driver-specific hack you can use depending on your project.

MongoDB insertMany method has options param you can change to make inserts even faster:

writeConcern
ordered

Using them may also speed up insertion.

But you need to keep in mind that using: ordered: false will not allow you to rely on the document's insertion order.

writeConcern: {w: 0} requests no acknowledgment of the write operation (source).

If none of these imply the app requirements you are welcome to use it!

Wrap up!

you can easily insert 1m of records using insertMany . It should not make any dramatic impact on your app performance.
loading from 1m to 10m would probably require you to use Stream API and some parallelization of inserts. It may impact the performance of the App because of insertion time and increased write load.
to load more than 10 million you need parallelization of both processes and insertions to make it fast. App performance can degrade dramatically since cluster CPU and IOPS are used a lot. Schedule script run or tune how many operations/processes can execute in parallel

Immense appreciation for you reading up to this moment, hope It will help you become a better dev.

Subscribe to my Medium, follow me on Twitter, get connected on LinkedIn, and see code examples on GitHub!