TIMX 611 - AOSS support

README.md

@@ -135,13 +135,15 @@ WORKSPACE=### Set to `dev` for local development, this will be set to `stage` an
 
 ```shell
 AWS_REGION=### Only needed if AWS region changes from the default of us-east-1.
-OPENSEARCH_BULK_MAX_CHUNK_BYTES=### Chunk size limit for sending requests to the bulk indexing endpoint, in bytes. Defaults to 104857600 (100 * 1024 * 1024) if not set.
-OPENSEARCH_BULK_MAX_RETRIES=### Maximum number of retries when sending requests to the bulk indexing endpoint. Defaults to 50 if not set.
+OPENSEARCH_BULK_CHUNK_SIZE=### Chunk size limit for sending requests to the bulk indexing endpoint, in documents. Defaults to 100.
+OPENSEARCH_BULK_MAX_CHUNK_BYTES=### Chunk size limit for sending requests to the bulk indexing endpoint, in bytes. Defaults to 104857600 (100 * 1024 * 1024).
+OPENSEARCH_BULK_MAX_RETRIES=### Maximum number of retries when sending requests to the bulk indexing endpoint. Defaults to 5.
 OPENSEARCH_INITIAL_ADMIN_PASSWORD=###If using a local Docker OpenSearch instance, this must be set (for versions >= 2.12.0).
 OPENSEARCH_REQUEST_TIMEOUT=### Only used for OpenSearch requests that tend to take longer than the default timeout of 10 seconds, such as bulk or index refresh requests. Defaults to 120 seconds if not set.
 STATUS_UPDATE_INTERVAL=### The ingest process logs the # of records indexed every nth record. Set this env variable to any integer to change the frequency of logging status updates. Can be useful for development/debugging. Defaults to 1000 if not set.
 TIMDEX_OPENSEARCH_ENDPOINT=### If using a local Docker OpenSearch instance, this isn't needed. Otherwise set to OpenSearch instance endpoint without the http scheme (e.g., "search-timdex-env-1234567890.us-east-1.es.amazonaws.com"). Can also be passed directly to the CLI via the `--url` option.
 SENTRY_DSN=### If set to a valid Sentry DSN, enables Sentry exception monitoring This is not needed for local development.
+AUTH_SERVICE_TYPE=### Indicate instanced type for authentication, "es" (cluster) or "aoss" (serverless); defaults to "es"
 ```
 
 ## CLI commands

pyproject.toml

@@ -49,24 +49,24 @@ show-fixes = true
 select = ["ALL", "PT"]
 
 ignore = [
-    # default
-    "COM812",
-    "D107",
-    "N812",
-    "PTH",
-
-    # project-specific
     "C90",
+    "COM812",
     "D100",
     "D101",
     "D102",
     "D103",
     "D104",
+    "D107",
+    "EM101",
+    "EM102",
     "G004",
+    "N812",
     "PLR0912",
     "PLR0913",
     "PLR0915",
+    "PTH",
     "S321",
+    "TRY003"
 ]
 
 # allow autofix behavior for specified rules

tests/test_config.py

@@ -1,4 +1,5 @@
 import logging
+from unittest import mock
 
 from tim.config import (
     OPENSEARCH_BULK_CONFIG_DEFAULTS,
@@ -31,10 +32,12 @@ def test_configure_logger_verbose():
 
 
 def test_configure_opensearch_bulk_settings_from_env(monkeypatch):
+    monkeypatch.setenv("OPENSEARCH_BULK_CHUNK_SIZE", "999")
     monkeypatch.setenv("OPENSEARCH_BULK_MAX_CHUNK_BYTES", "10")
     monkeypatch.setenv("OPENSEARCH_BULK_MAX_RETRIES", "2")
     monkeypatch.setenv("OPENSEARCH_REQUEST_TIMEOUT", "20")
     assert configure_opensearch_bulk_settings() == {
+        "OPENSEARCH_BULK_CHUNK_SIZE": 999,
         "OPENSEARCH_BULK_MAX_CHUNK_BYTES": 10,
         "OPENSEARCH_BULK_MAX_RETRIES": 2,
         "OPENSEARCH_REQUEST_TIMEOUT": 20,
@@ -60,7 +63,12 @@ def test_configure_sentry_env_variable_is_none(monkeypatch):
     assert result == "No Sentry DSN found, exceptions will not be sent to Sentry"
 
 
-def test_configure_sentry_env_variable_is_dsn(monkeypatch):
+@mock.patch("tim.config.sentry_sdk.init")
+def test_configure_sentry_env_variable_is_dsn(mock_sentry_init, monkeypatch):
     monkeypatch.setenv("SENTRY_DSN", "https://1234567890@00000.ingest.sentry.io/123456")
     result = configure_sentry()
+    mock_sentry_init.assert_called_once_with(
+        "https://1234567890@00000.ingest.sentry.io/123456",
+        environment="test",
+    )
     assert result == "Sentry DSN found, exceptions will be sent to Sentry with env=test"

tests/test_opensearch.py

@@ -37,6 +37,21 @@ def test_configure_opensearch_client_for_aws(mocked_boto3_session):
     )
 
 
+@mock.patch("tim.opensearch.OpenSearch")
+@mock.patch("tim.opensearch.AWSV4SignerAuth")
+@mock.patch("tim.opensearch.boto3.Session")
+def test_configure_opensearch_client_for_aoss_sets_auth_service(
+    mocked_session, mocked_auth, mocked_opensearch, monkeypatch
+):
+    credentials = mock.Mock()
+    mocked_session.return_value.get_credentials.return_value = credentials
+    monkeypatch.setenv("AUTH_SERVICE_TYPE", "aoss")
+
+    tim_os.configure_opensearch_client("fake-dev.us-east-1.es.amazonaws.com")
+
+    mocked_auth.assert_called_once_with(credentials, "us-east-1", service="aoss")
+
+
 @my_vcr.use_cassette("ping_localhost.yaml")
 def test_get_formatted_info(test_opensearch_client):
     assert tim_os.get_formatted_info(test_opensearch_client) == (
@@ -133,8 +148,6 @@ def test_get_formatted_indexes(test_opensearch_client):
         "\n  Primary store size: 208b"
         "\n  Total store size: 208b"
         "\n  UUID: 60Gq-vaAScOKGXkG_JAw5A"
-        "\n  Primary Shards: 1"
-        "\n  Replica Shards: 1"
         "\n"
         "\nName: index-with-no-aliases"
         "\n  Aliases: None"
@@ -144,8 +157,6 @@ def test_get_formatted_indexes(test_opensearch_client):
         "\n  Primary store size: 208b"
         "\n  Total store size: 208b"
         "\n  UUID: KqVlOA5lTw-fXZA2TEqi_g"
-        "\n  Primary Shards: 1"
-        "\n  Replica Shards: 1"
         "\n"
         "\nName: index-with-one-alias"
         "\n  Aliases: alias-with-multiple-indexes"
@@ -155,8 +166,6 @@ def test_get_formatted_indexes(test_opensearch_client):
         "\n  Primary store size: 208b"
         "\n  Total store size: 208b"
         "\n  UUID: q-NKXPp3SuWiDKhPkUxP-g"
-        "\n  Primary Shards: 1"
-        "\n  Replica Shards: 1"
         "\n"
     )
 

tim/cli.py

@@ -1,4 +1,3 @@
-# ruff: noqa: TRY003, EM101
 import json
 import logging
 from datetime import timedelta

tim/config.py

@@ -5,8 +5,9 @@
 import sentry_sdk
 
 OPENSEARCH_BULK_CONFIG_DEFAULTS = {
+    "OPENSEARCH_BULK_CHUNK_SIZE": 100,
     "OPENSEARCH_BULK_MAX_CHUNK_BYTES": 20 * 1024 * 1024,
-    "OPENSEARCH_BULK_MAX_RETRIES": 50,
+    "OPENSEARCH_BULK_MAX_RETRIES": 5,
     "OPENSEARCH_REQUEST_TIMEOUT": 120,
 }
 PRIMARY_ALIAS = "all-current"

tim/opensearch.py

@@ -47,7 +47,21 @@ def configure_opensearch_client(url: str) -> OpenSearch:
 
     credentials = boto3.Session().get_credentials()
     region = os.getenv("AWS_REGION", "us-east-1")
-    auth = AWSV4SignerAuth(credentials, region)
+
+    auth_service_type = os.getenv("AUTH_SERVICE_TYPE", "es")
+    valid_auth_service_types = {"aoss", "es"}
+
+    if auth_service_type not in valid_auth_service_types:
+        raise ValueError(
+            f"AUTH_SERVICE_TYPE must be one of {sorted(valid_auth_service_types)}, "
+            f"got {auth_service_type!r}"
+        )
+
+    auth = AWSV4SignerAuth(
+        credentials,
+        region,
+        service=auth_service_type,
+    )
     return OpenSearch(
         hosts=[{"host": url, "port": "443"}],
         http_auth=auth,
@@ -123,8 +137,6 @@ def get_formatted_indexes(client: OpenSearch) -> str:
                 f"  Primary store size: {info['pri.store.size']}\n"
                 f"  Total store size: {info['store.size']}\n"
                 f"  UUID: {info['uuid']}\n"
-                f"  Primary Shards: {int(info['pri']):,}\n"
-                f"  Replica Shards: {int(info['rep']):,}\n"
             )
         return output
     return output + " No indexes present in OpenSearch cluster."
@@ -328,6 +340,9 @@ def bulk_delete(
     responses = streaming_bulk(
         client,
         actions,
+        chunk_size=REQUEST_CONFIG["OPENSEARCH_BULK_CHUNK_SIZE"],
+        max_retries=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_RETRIES"],
+        initial_backoff=3,
         max_chunk_bytes=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_CHUNK_BYTES"],
         raise_on_error=False,
     )
@@ -351,11 +366,6 @@ def bulk_delete(
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records deleted so far!", result["total"])
-    logger.info("Refreshing index.")
-    response = client.indices.refresh(
-        index=index,
-    )
-    logger.debug(response)
     return result
 
 
@@ -382,6 +392,8 @@ def bulk_index(client: OpenSearch, index: str, records: Iterator[dict]) -> dict[
     responses = streaming_bulk(
         client,
         actions,
+        chunk_size=REQUEST_CONFIG["OPENSEARCH_BULK_CHUNK_SIZE"],
+        max_retries=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_RETRIES"],
         max_chunk_bytes=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_CHUNK_BYTES"],
         raise_on_error=False,
     )
@@ -411,11 +423,6 @@ def bulk_index(client: OpenSearch, index: str, records: Iterator[dict]) -> dict[
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records indexed so far!", result["total"])
-    logger.info("Refreshing index.")
-    response = client.indices.refresh(
-        index=index,
-    )
-    logger.debug(response)
     return result
 
 
@@ -437,6 +444,8 @@ def bulk_update(
     responses = streaming_bulk(
         client,
         actions,
+        chunk_size=REQUEST_CONFIG["OPENSEARCH_BULK_CHUNK_SIZE"],
+        max_retries=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_RETRIES"],
         max_chunk_bytes=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_CHUNK_BYTES"],
         raise_on_error=False,
     )
@@ -475,9 +484,4 @@ def bulk_update(
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records updated so far!", result["total"])
-    logger.info("Refreshing index.")
-    response = client.indices.refresh(
-        index=index,
-    )
-    logger.debug(response)
     return result