Normalize Lambda and OpenSearch client config

.env.test

@@ -6,4 +6,5 @@ AWS_ENDPOINT_URL_LAMBDA=http://localhost:9200
 AWS_ACCESS_KEY_ID=test-key
 AWS_SECRET_ACCESS_KEY=test-secret
 AWS_REGION=us-east-1
+AWS_SESSION_TOKEN=test-session-token
 TIMDEX_SEMANTIC_BUILDER_FUNCTION_NAME=timdex-semantic-builder-prod:live

README.md

@@ -20,6 +20,7 @@ This application interfaces with an OpenSearch backend and exposes a GraphQL end
 - [AWS Configuration](#aws-configuration)
   - [OpenSearch Configuration](#opensearch-configuration)
   - [AWS Credentials (Used for AWS-based OpenSearch and timdex-semantic-builder)](#aws-credentials-used-for-aws-based-opensearch-and-timdex-semantic-builder)
+  - [TIMDEX Semantic Builder Lambda Authentication](#timdex-semantic-builder-lambda-authentication)
   - [AWS OpenSearch Service (Legacy)](#aws-opensearch-service-legacy)
   - [AWS OpenSearch Serverless (AOSS)](#aws-opensearch-serverless-aoss)
   - [TIMDEX Semantic Builder Lambda](#timdex-semantic-builder-lambda)
@@ -212,9 +213,22 @@ locally.
 - `AWS_ACCESS_KEY_ID`: AWS access key for OpenSearch and Lambda
 - `AWS_SECRET_ACCESS_KEY`: AWS secret key for OpenSearch and Lambda
 - `AWS_REGION`: AWS region for OpenSearch and Lambda services
+- `AWS_ROLE_ARN`: IAM role ARN to assume when using role-based AWS authentication.
+  Used by OpenSearch (AOSS) and TIMDEX Semantic Builder Lambda when `AWS_SESSION_TOKEN` is not set.
 - `AWS_SESSION_TOKEN`: (Optional) AWS session token for temporary credentials when using expiring AWS credentials.
   Use this with temporary AWS credentials for AWS-based OpenSearch access and Lambda.
-  For AOSS, when this is set, temporary credentials are used directly and `AWS_AOSS_ROLE_ARN` is not needed.
+  When this is set, temporary credentials are used directly and `AWS_ROLE_ARN` is not needed.
+
+### TIMDEX Semantic Builder Lambda Authentication
+
+Credential behavior for the Lambda client is:
+
+1. If `AWS_SESSION_TOKEN` is set, use static credentials with the session token.
+2. Otherwise, if `AWS_ROLE_ARN` is set, assume that role.
+3. Otherwise, if `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` are set, use long-lived static credentials.
+4. If required credential env vars are missing, initialization fails with an explicit configuration error.
+
+When both `AWS_ROLE_ARN` and `AWS_SESSION_TOKEN` are set, Lambda uses `AWS_SESSION_TOKEN` directly.
 
 ### AWS OpenSearch Service (Legacy)
 
@@ -227,8 +241,8 @@ This is our legacy AWS OpenSearch Service Cluster. All production instances shou
 This is our upcoming configuration once migration is complete. This uses a different [authentication mechanism](https://github.com/awsdocs/amazon-opensearch-service-developer-guide/blob/master/doc_source/serverless-clients.md#ruby) than our legacy AWS OpenSearch Service.
 
 - `AWS_AOSS`: boolean. Set to `true` to enable AWS OpenSearch Serverless (AOSS).
-- `AWS_AOSS_ROLE_ARN`: AWS IAM role ARN to assume for AOSS authentication. **Required when** `AWS_AOSS=true` **and** `AWS_SESSION_TOKEN` is not set. This enables automatic credential refresh via role assumption.
-  When `AWS_SESSION_TOKEN` is present, temporary credentials are used directly and `AWS_AOSS_ROLE_ARN` is not needed. This is only used in local development. `AWS_AOSS_ROLE_ARN` is used in production.
+- `AWS_ROLE_ARN`: AWS IAM role ARN to assume for AOSS authentication. **Required when** `AWS_AOSS=true` **and** `AWS_SESSION_TOKEN` is not set. This enables automatic credential refresh via role assumption.
+  When `AWS_SESSION_TOKEN` is present, temporary credentials are used directly and `AWS_ROLE_ARN` is not needed.
 
 ### TIMDEX Semantic Builder Lambda
 

config/initializers/lambda.rb

@@ -1,16 +1,35 @@
 require 'aws-sdk-lambda'
+require 'aws_auth'
+require 'aws_config_validator'
+
+def validate_lambda_config!
+  AwsConfigValidator.validate_lambda_config
+end
+
+def lambda_credentials
+  if ENV['AWS_SESSION_TOKEN'].present?
+    Rails.logger.debug 'Configuring Lambda client with temporary static credentials (session token)'
+    return AwsAuth.static_credentials
+  end
+
+  Rails.logger.debug 'Configuring Lambda client with assumed role credentials'
+  AwsAuth.assume_role_credentials(role_session_name: 'timdex-lambda')
+end
 
 def configure_lambda_client
-  options = {
-    region: ENV.fetch('AWS_REGION', 'us-east-1'),
-    access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID'),
-    secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY')
-  }
-  options[:session_token] = ENV['AWS_SESSION_TOKEN'] if ENV['AWS_SESSION_TOKEN'].present?
+  validate_lambda_config!
+
+  Rails.logger.debug 'Configuring AWS Lambda client'
+
+  options = { region: ENV.fetch('AWS_REGION', 'us-east-1') }
+  options[:credentials] = lambda_credentials
 
   # AWS SDK sets this env in prod. However, we need to conditionally set it for tests so VCR can
   # intercept the requests with a fake URL.
-  options[:endpoint] = ENV['AWS_ENDPOINT_URL_LAMBDA'] if ENV['AWS_ENDPOINT_URL_LAMBDA'].present?
+  if ENV['AWS_ENDPOINT_URL_LAMBDA'].present?
+    Rails.logger.debug 'Using AWS_ENDPOINT_URL_LAMBDA override for Lambda client endpoint'
+    options[:endpoint] = ENV['AWS_ENDPOINT_URL_LAMBDA']
+  end
   Aws::Lambda::Client.new(options)
 end
 

config/initializers/opensearch.rb

@@ -1,7 +1,8 @@
 require 'faraday_middleware/aws_sigv4' if ENV['AWS_OPENSEARCH'] == 'true' && ENV.fetch('AWS_AOSS', 'false') == 'false'
 require 'opensearch-aws-sigv4'
 require 'aws-sigv4'
-require 'opensearch_config_validator'
+require 'aws_auth'
+require 'aws_config_validator'
 
 # Helper method to parse OPENSEARCH_LOG as a boolean
 # Environment variables are always strings, so 'false' is truthy
@@ -13,10 +14,10 @@ def opensearch_logging_enabled?
 # Priority is given to AWS AOSS, then AWS OpenSearch, and finally vanilla OpenSearch
 def configure_opensearch
   if ENV['AWS_AOSS'] == 'true'
-    OpensearchConfigValidator.validate_aws_aoss_config
+    AwsConfigValidator.validate_aws_aoss_config
     aws_aoss_client
   elsif ENV['AWS_OPENSEARCH'] == 'true'
-    OpensearchConfigValidator.validate_aws_os_config
+    AwsConfigValidator.validate_aws_os_config
     aws_os_client
   else
     os_client
@@ -44,10 +45,10 @@ def os_client
 # obtained by assuming a role.
 def aws_os_client
   OpenSearch::Client.new log: opensearch_logging_enabled?, url: ENV.fetch('OPENSEARCH_URL', nil) do |config|
-    Rails.logger.debug "Configuring Legacy AWS OpenSearch Service client"
+    Rails.logger.debug 'Configuring Legacy AWS OpenSearch Service client'
     # personal keys use expiring credentials with tokens
     if ENV['AWS_SESSION_TOKEN'].present?
-      Rails.logger.debug 'Using temporary credentials with session token'
+      Rails.logger.debug 'Using temporary credentials with session token for OpenSearch Service client'
       config.request :aws_sigv4,
                      service: 'es',
                      region: ENV.fetch('AWS_REGION', nil),
@@ -56,7 +57,7 @@ def aws_os_client
                      session_token: ENV['AWS_SESSION_TOKEN']
     # application keys don't use tokens
     else
-      Rails.logger.debug 'Using long-lived credentials without session token'
+      Rails.logger.debug 'Using long-lived credentials without session token for OpenSearch Service client'
       config.request :aws_sigv4,
                      service: 'es',
                      region: ENV.fetch('AWS_REGION', nil),
@@ -74,12 +75,22 @@ def aws_os_client
 # @note this configuration uses temporary credentials obtained by assuming a role or via the AWS console, unlike
 # AWS OpenSearch Service which can use long-lived access keys directly.
 def aws_aoss_client
-  Rails.logger.debug "Configuring AWS AOSS client"
+  Rails.logger.debug 'Configuring AWS OpenSearch Serverless (AOSS) client'
+
+  credentials_provider = if ENV.fetch('AWS_SESSION_TOKEN', false).present?
+                           Rails.logger.debug 'Using temporary credentials with session token for OpenSearch AOSS ' \
+                                              'client'
+                           AwsAuth.static_credentials
+                         else
+                           Rails.logger.debug 'Using long-lived credentials and assuming role for OpenSearch AOSS ' \
+                                              'client'
+                           AwsAuth.assume_role_credentials(role_session_name: 'timdex-opensearch')
+                         end
 
   signer = Aws::Sigv4::Signer.new(
     service: 'aoss',
-    region: ENV.fetch('AWS_REGION', nil),
-    credentials_provider: credentials
+    region: ENV.fetch('AWS_REGION', 'us-east-1'),
+    credentials_provider: credentials_provider
   )
 
   OpenSearch::Aws::Sigv4Client.new(
@@ -91,42 +102,4 @@ def aws_aoss_client
   )
 end
 
-def credentials
-  if ENV.fetch('AWS_SESSION_TOKEN', false).present?
-    Rails.logger.debug 'Using temporary credentials with session token'
-    temporary_credentials
-  else
-    Rails.logger.debug 'Using long-lived credentials and assuming role'
-    assume_role_credentials
-  end
-end
-
-# personal keys use expiring credentials with tokens, so we use them directly without assuming a role
-# application keys use long-lived credentials and assume a role to get temporary credentials for AOSS
-def temporary_credentials
-  Aws::Credentials.new(
-    ENV.fetch('AWS_ACCESS_KEY_ID', nil),
-    ENV.fetch('AWS_SECRET_ACCESS_KEY', nil),
-    ENV.fetch('AWS_SESSION_TOKEN', nil)
-  )
-end
-
-# AWS AOSS uses temporary credentials that are obtained by assuming a role. The
-# Aws::AssumeRoleCredentials class is used to get these temporary credentials. It requires the ARN of
-# the role to assume, a session name, and a client for the AWS Security Token Service (STS) which is
-# used to perform the AssumeRole operation. It uses the AWS region and access keys from the
-# environment variables to create the STS client. When the session token expires, the
-# Aws::AssumeRoleCredentials will automatically refresh the credentials by calling AssumeRole again.
-def assume_role_credentials
-  Aws::AssumeRoleCredentials.new(
-    role_arn: ENV.fetch('AWS_AOSS_ROLE_ARN', nil),
-    role_session_name: 'timdex-opensearch',
-    client: Aws::STS::Client.new(
-      region: ENV.fetch('AWS_REGION', nil),
-      access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID', nil),
-      secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
-    )
-  )
-end
-
 Timdex::OSClient = configure_opensearch

lib/aws_auth.rb

@@ -0,0 +1,51 @@
+# AwsAuth centralizes AWS credential and role-assumption helpers shared by
+# OpenSearch and Lambda initialization.
+#
+# It supports two auth patterns:
+# 1) direct static credentials from environment variables
+# 2) temporary assumed-role credentials via AWS STS
+module AwsAuth
+  module_function
+
+  # Builds static credentials directly from environment variables.
+  #
+  # @return [Aws::Credentials] Credentials built from AWS_ACCESS_KEY_ID,
+  #   AWS_SECRET_ACCESS_KEY, and optional AWS_SESSION_TOKEN.
+  def static_credentials
+    Aws::Credentials.new(
+      ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+      ENV.fetch('AWS_SECRET_ACCESS_KEY', nil),
+      ENV.fetch('AWS_SESSION_TOKEN', nil)
+    )
+  end
+
+  # Builds an AWS STS client used for role-assumption flows.
+  #
+  # Region defaults to us-east-1 when AWS_REGION is not set.
+  # Optional AWS_SESSION_TOKEN is forwarded when present.
+  #
+  # @return [Aws::STS::Client] Configured STS client instance.
+  def sts_client
+    options = {
+      region: ENV.fetch('AWS_REGION', 'us-east-1'),
+      access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+      secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+    }
+    options[:session_token] = ENV['AWS_SESSION_TOKEN'] if ENV['AWS_SESSION_TOKEN'].present?
+
+    Aws::STS::Client.new(options)
+  end
+
+  # Builds auto-refreshing assumed-role credentials backed by STS.
+  #
+  # @param role_session_name [String] Session identifier used in STS and
+  #   CloudTrail for traceability.
+  # @return [Aws::AssumeRoleCredentials] Refreshing credentials provider.
+  def assume_role_credentials(role_session_name:)
+    Aws::AssumeRoleCredentials.new(
+      role_arn: ENV.fetch('AWS_ROLE_ARN', nil),
+      role_session_name: role_session_name,
+      client: sts_client
+    )
+  end
+end

lib/aws_config_validator.rb

@@ -0,0 +1,77 @@
+# AwsConfigValidator validates AWS-related environment configuration for
+# OpenSearch and Lambda initialization.
+#
+# Validation is split by integration pathway:
+# 1) Lambda runtime credentials
+# 2) AWS AOSS (OpenSearch Serverless)
+# 3) AWS-managed OpenSearch
+class AwsConfigValidator
+  class << self
+    # Validates required Lambda credential environment variables.
+    #
+    # AWS_ROLE_ARN becomes required only when AWS_SESSION_TOKEN is not present.
+    # This supports both temporary session-token auth and role-assumption auth.
+    #
+    # @return [nil]
+    # @raise [RuntimeError] when required env vars are missing.
+    def validate_lambda_config
+      required_vars = {
+        'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+        'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+      }
+
+      required_vars['AWS_ROLE_ARN'] = ENV.fetch('AWS_ROLE_ARN', nil) if ENV['AWS_SESSION_TOKEN'].blank?
+
+      validate_required_vars!(required_vars, error_prefix: 'AWS Lambda Config Error')
+    end
+
+    # Validates required configuration for AWS AOSS connections.
+    #
+    # AWS_ROLE_ARN becomes required only when AWS_SESSION_TOKEN is not present.
+    # This supports both temporary session-token auth and role-assumption auth.
+    #
+    # @return [nil]
+    # @raise [RuntimeError] when required env vars are missing.
+    def validate_aws_aoss_config
+      required_vars = {
+        'OPENSEARCH_URL' => ENV.fetch('OPENSEARCH_URL', nil),
+        'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+        'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+      }
+
+      # Required only when AWS_SESSION_TOKEN is not present (using role assumption)
+      required_vars['AWS_ROLE_ARN'] = ENV.fetch('AWS_ROLE_ARN', nil) if ENV['AWS_SESSION_TOKEN'].blank?
+
+      validate_required_vars!(required_vars, error_prefix: 'AWS AOSS Config Error')
+    end
+
+    # Validates required configuration for AWS-managed OpenSearch.
+    #
+    # @return [nil]
+    # @raise [RuntimeError] when required env vars are missing.
+    def validate_aws_os_config
+      validate_required_vars!({
+                                'OPENSEARCH_URL' => ENV.fetch('OPENSEARCH_URL', nil),
+                                'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+                                'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+                              }, error_prefix: 'AWS OpenSearch Config Error')
+    end
+
+    private
+
+    # Raises a standardized configuration error for missing env vars.
+    #
+    # @param required_vars [Hash{String => Object}] Mapping of env var names
+    #   to their current values.
+    # @param error_prefix [String] Prefix used to identify the validator path.
+    # @return [nil]
+    # @raise [RuntimeError] when one or more values are blank.
+    def validate_required_vars!(required_vars, error_prefix:)
+      missing_vars = required_vars.select { |_key, value| value.blank? }.keys
+
+      return unless missing_vars.any?
+
+      raise "#{error_prefix}: These required environment variables are not set: #{missing_vars.join(', ')}"
+    end
+  end
+end