[Issue #2094] Update the way we setup the connection to OpenSearch

api/src/adapters/search/opensearch_client.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Any, Generator, Iterable
 
+import boto3
 import opensearchpy
 
 from src.adapters.search.opensearch_config import OpensearchConfig, get_opensearch_config
@@ -252,14 +253,24 @@ def scroll(
 
 
 def _get_connection_parameters(opensearch_config: OpensearchConfig) -> dict[str, Any]:
-    # TODO - we'll want to add the AWS connection params here when we set that up
     # See: https://opensearch.org/docs/latest/clients/python-low-level/#connecting-to-amazon-opensearch-serverless
+    # for further details on configuring the connection to OpenSearch
 
-    return dict(
+    params = dict(
         hosts=[{"host": opensearch_config.host, "port": opensearch_config.port}],
         http_compress=True,
         use_ssl=opensearch_config.use_ssl,
         verify_certs=opensearch_config.verify_certs,
-        ssl_assert_hostname=False,
-        ssl_show_warn=False,
+        connection_class=opensearchpy.RequestsHttpConnection,
+        pool_maxsize=opensearch_config.connection_pool_size,
     )
+
+    # If an AWS region is set, we assume we're running non-locally
+    # and will attempt to authenticate with AOSS
+    if opensearch_config.aws_region is not None:
+        # Get credentials and authorize with AWS Opensearch Serverless (aoss)
+        credentials = boto3.Session().get_credentials()
+        auth = opensearchpy.AWSV4SignerAuth(credentials, opensearch_config.aws_region, "aoss")
+        params["http_auth"] = auth
+
+    return params

api/src/adapters/search/opensearch_config.py

@@ -15,6 +15,10 @@ class OpensearchConfig(PydanticBaseEnvConfig):
     port: int  # OPENSEARCH_PORT
     use_ssl: bool = Field(default=True)  # OPENSEARCH_USE_SSL
     verify_certs: bool = Field(default=True)  # OPENSEARCH_VERIFY_CERTS
+    connection_pool_size: int = Field(default=10)  # OPENSEARCH_CONNECTION_POOL_SIZE
+
+    # AWS configuration
+    aws_region: str | None = Field(default=None)  # OPENSEARCH_AWS_REGION
 
 
 def get_opensearch_config() -> OpensearchConfig:
@@ -27,6 +31,7 @@ def get_opensearch_config() -> OpensearchConfig:
             "port": opensearch_config.port,
             "use_ssl": opensearch_config.use_ssl,
             "verify_certs": opensearch_config.verify_certs,
+            "connection_pool_size": opensearch_config.connection_pool_size,
         },
     )
 

api/tests/src/adapters/search/test_opensearch_client.py

@@ -1,7 +1,11 @@
 import uuid
 
+import opensearchpy
 import pytest
 
+from src.adapters.search import get_opensearch_config
+from src.adapters.search.opensearch_client import _get_connection_parameters
+
 ########################################################################
 # These tests are primarily looking to validate
 # that our wrappers around the OpenSearch client
@@ -193,3 +197,19 @@ def test_scroll(search_client, generic_index):
     assert len(results[0].records) == 3
     assert len(results[1].records) == 3
     assert len(results[2].records) == 2
+
+
+def test_get_connection_parameters():
+    # Just validating this builds as expected for local mode
+    config = get_opensearch_config()
+    params = _get_connection_parameters(config)
+
+    # Mostly validating defaults get used
+    assert params == {
+        "hosts": [{"host": config.host, "port": 9200}],
+        "http_compress": True,
+        "use_ssl": False,
+        "verify_certs": False,
+        "connection_class": opensearchpy.RequestsHttpConnection,
+        "pool_maxsize": 10,
+    }