BSQIK/Search: Indexing

.github/dependabot.yml

@@ -10,6 +10,11 @@ updates:
     schedule:
       interval: "weekly"
 
+  - package-ecosystem: "docker"
+    directory: "/lambdas/indexer"
+    schedule:
+      interval: "daily"
+
   - package-ecosystem: "docker"
     directory: "/lambdas/thumbnail"
     schedule:

.github/workflows/deploy-lambdas.yaml

@@ -14,12 +14,13 @@ jobs:
       matrix:
         path:
           - access_counts
-          - indexer
+          - es_ingest
           - pkgevents
           - pkgpush
           - preview
           - s3hash
           - status_reports
+          - manifest_indexer
           - tabular_preview
           - transcode
     runs-on: ubuntu-latest
@@ -29,6 +30,11 @@ jobs:
       contents: read
     steps:
       - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+      - name: Export requirements.txt from uv.lock
+        if: ${{ hashFiles(format('lambdas/{0}/uv.lock', matrix.path)) != '' }}
+        run: |
+          uv export --locked --no-emit-project --no-hashes --directory lambdas/${{ matrix.path }} -o requirements.txt
       - name: Build zip
         run: |
           BUILDER_IMAGE=quiltdata/lambda:build-$(cat lambdas/${{ matrix.path }}/.python-version)
@@ -64,6 +70,7 @@ jobs:
     strategy:
       matrix:
         path:
+          - indexer
           - thumbnail
     runs-on: ubuntu-latest
     # These permissions are needed to interact with GitHub's OIDC Token endpoint.

.github/workflows/py-ci.yml

@@ -161,7 +161,9 @@ jobs:
       matrix:
         path:
           - access_counts
+          - es_ingest
           - indexer
+          - manifest_indexer
           - pkgevents
           - pkgpush
           - preview
@@ -181,6 +183,12 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version-file: lambdas/${{ matrix.path }}/.python-version
+      - uses: astral-sh/setup-uv@v6
+      - name: Export requirements.txt from uv.lock
+        if: ${{ hashFiles(format('lambdas/{0}/uv.lock', matrix.path)) != '' }}
+        run: |
+          uv export --locked --no-emit-project --no-hashes --directory lambdas/${{ matrix.path }} -o requirements.txt
+          uv export --locked --no-emit-project --no-hashes --directory lambdas/${{ matrix.path }} --only-group test -o test-requirements.txt
       - name: Install dependencies
         run: |
           if [ ${{ matrix.path }} == "thumbnail" ]
@@ -203,7 +211,7 @@ jobs:
           python -m pip install -r lambdas/${{ matrix.path }}/test-requirements.txt
       - name: Pytest
         run: |
-          pytest --cov=lambdas lambdas/${{ matrix.path }}
+          python -m pytest --cov=lambdas lambdas/${{ matrix.path }}
       - uses: codecov/codecov-action@v5
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

lambdas/es_ingest/.python-version

@@ -0,0 +1 @@
+3.11

lambdas/es_ingest/CHANGELOG.md

@@ -0,0 +1,20 @@
+<!-- markdownlint-disable line-length -->
+# Changelog
+
+Changes are listed in reverse chronological order (newer entries at the top).
+The entry format is
+
+```markdown
+- [Verb] Change description ([#<PR-number>](https://github.com/quiltdata/quilt/pull/<PR-number>))
+```
+
+where verb is one of
+
+- Removed
+- Added
+- Fixed
+- Changed
+
+## Changes
+
+- [Added] Bootstrap the change log ([#4422](https://github.com/quiltdata/quilt/pull/4422))

lambdas/es_ingest/pyproject.toml

@@ -0,0 +1,30 @@
+[project]
+name = "t4_lambda_es_ingest"
+version = "0.1.0"
+authors = [
+    { name = "Sergey Fedoseev", email = "[email protected]" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "quilt-shared[boto,es]",
+    "t4-lambda-shared",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv.sources]
+t4-lambda-shared = { url = "https://github.com/quiltdata/quilt/archive/d496dffbfb4b7a2ae05f6c1f7f0cb7d5d43bc984.zip", subdirectory = "lambdas/shared" }
+quilt-shared = { url = "https://github.com/quiltdata/quilt/archive/df53c9ce125ea051e0d1ac41d58796336e202256.zip", subdirectory = "py-shared" }
+
+[dependency-groups]
+test = [
+    "pytest~=8.4",
+    "pytest-cov~=6.2",
+    "pytest-env~=1.1",
+    "pytest-mock~=3.14",
+]
+
+[tool.uv]
+default-groups = ["test"]

lambdas/es_ingest/pytest.ini

@@ -0,0 +1,6 @@
+[pytest]
+env =
+    ES_ENDPOINT = http://localhost:9200
+    AWS_ACCESS_KEY_ID = test_key
+    AWS_SECRET_ACCESS_KEY = test_secret
+    AWS_DEFAULT_REGION = us-east-1

lambdas/es_ingest/src/t4_lambda_es_ingest/__init__.py

@@ -0,0 +1,90 @@
+import json
+import os
+import random
+import time
+
+import boto3
+import elasticsearch
+
+from quilt_shared.es import make_elastic
+from t4_lambda_shared.utils import get_quilt_logger
+
+s3_client = boto3.client("s3")
+es = make_elastic(os.environ["ES_ENDPOINT"])
+logger = get_quilt_logger()
+
+
+EXPECTED_ES_RESPONSE_TIME = 10  # seconds
+
+
+class BulkDocumentError(Exception):
+    pass
+
+
+class TooManyRequestsError(Exception):
+    pass
+
+
+def sleep_until_timeout(context):
+    """Sleep until the lambda timeout"""
+    remaining = context.get_remaining_time_in_millis() / 1000 - 1
+    logger.warning("Sleeping for %s seconds just before lambda timeout", remaining)
+    # good night, sweet prince
+    time.sleep(remaining)
+
+
+def bulk(context, es, data: bytes):
+    t0 = time.time()
+    try:
+        resp = es.bulk(
+            data,
+            filter_path="errors",
+            # wait as much as possible because it's better die trying than just die
+            # leave a second to avoid lambda timeout
+            request_timeout=context.get_remaining_time_in_millis() / 1000 - 1,
+        )
+    except elasticsearch.exceptions.TransportError as e:
+        if e.status_code == 429:
+            # at this point ES seems to be *very* overloaded, so we just sleep until lambda timeout
+            logger.warning("Got a 429 Too Many Requests error, sleeping until lambda timeout")
+            sleep_until_timeout(context)
+            raise TooManyRequestsError
+        raise
+
+    t1 = time.time()
+    delta = t1 - t0
+    logger.info("Bulk request took %s seconds", delta)
+    overtime = delta - EXPECTED_ES_RESPONSE_TIME
+    if overtime > 0:
+        # if the request took so long ES seems to be overloaded, so it's better to sleep
+        # now to avoid 429 Too Many Requests error later causing lambda failure and retry
+        time_to_sleep = min(
+            random.uniform(overtime / 2, overtime) + 15,
+            context.get_remaining_time_in_millis() / 1000 - 1,
+        )
+        logger.warning("Sleeping for %s seconds to avoid ES overload", time_to_sleep)
+        time.sleep(time_to_sleep)
+    if resp["errors"]:
+        # TODO: log errors from items.*.error?
+        # TODO: ignore index_not_found_exception for delete operations?
+        raise BulkDocumentError
+
+
+def handler(event, context):
+    logger.debug("Invoked with event: %s", event)
+    assert len(event["Records"]) == 1, "Expected exactly on SQS message"
+    (event,) = event["Records"]
+    event = json.loads(event["body"])
+    assert len(event["Records"]) == 1, "Expected exactly one S3 event record"
+    (event,) = event["Records"]
+
+    bucket = event["s3"]["bucket"]["name"]
+    key = event["s3"]["object"]["key"]
+    version_id = event["s3"]["object"].get("versionId")
+    params = {"Bucket": bucket, "Key": key}
+    if version_id:
+        params["VersionId"] = version_id
+
+    data = s3_client.get_object(**params)["Body"].read()
+    bulk(context, es, data)
+    s3_client.delete_object(**params)

lambdas/es_ingest/tests/test_main.py

@@ -0,0 +1,91 @@
+import json
+
+import elasticsearch
+import pytest
+from botocore.stub import Stubber
+
+import t4_lambda_es_ingest
+
+
+def test_bulk_error(mocker):
+    mock_bulk = mocker.patch("elasticsearch.Elasticsearch.bulk", return_value={"errors": True})
+    mock_context = mocker.MagicMock()
+    with pytest.raises(t4_lambda_es_ingest.BulkDocumentError):
+        t4_lambda_es_ingest.bulk(mock_context, t4_lambda_es_ingest.es, b"data")
+
+    mock_bulk.assert_called_once_with(
+        b"data",
+        filter_path=mocker.ANY,
+        request_timeout=mocker.ANY,
+    )
+
+
+def test_bulk_too_many_requests(mocker):
+    mocker.patch("elasticsearch.exceptions.TransportError.status_code", 429)
+    mock_bulk = mocker.patch("elasticsearch.Elasticsearch.bulk", side_effect=elasticsearch.exceptions.TransportError)
+    mock_context = mocker.MagicMock()
+    mock_sleep_until_timeout = mocker.patch("t4_lambda_es_ingest.sleep_until_timeout")
+
+    with pytest.raises(t4_lambda_es_ingest.TooManyRequestsError):
+        t4_lambda_es_ingest.bulk(mock_context, t4_lambda_es_ingest.es, b"data")
+
+    mock_bulk.assert_called_once_with(
+        b"data",
+        filter_path=mocker.ANY,
+        request_timeout=mocker.ANY,
+    )
+    mock_sleep_until_timeout.assert_called_once_with(mock_context)
+
+
+@pytest.mark.parametrize("version_id", ["test-version-id", None])
+def test_handler(mocker, version_id):
+    mock_context = mocker.MagicMock()
+    s3_record = {
+        "s3": {
+            "bucket": {"name": "test-bucket"},
+            "object": {
+                "key": "test-key",
+            },
+        }
+    }
+    if version_id:
+        s3_record["s3"]["object"]["versionId"] = version_id
+
+    mock_event = {
+        "Records": [
+            {"body": json.dumps({"Records": [s3_record]})},
+        ]
+    }
+
+    with Stubber(t4_lambda_es_ingest.s3_client) as stubber:
+        mock_bulk = mocker.patch("t4_lambda_es_ingest.bulk")
+        get_object_params = {"Bucket": "test-bucket", "Key": "test-key"}
+        if version_id:
+            get_object_params["VersionId"] = version_id
+
+        stubber.add_response(
+            "get_object",
+            {
+                "Body": mocker.MagicMock(read=lambda: b"test data"),
+                "LastModified": "2023-10-01T00:00:00Z",
+            },
+            get_object_params,
+        )
+        stubber.add_response(
+            "delete_object",
+            {},
+            (
+                {"Bucket": "test-bucket", "Key": "test-key"}
+                if version_id is None
+                else {
+                    "Bucket": "test-bucket",
+                    "Key": "test-key",
+                    "VersionId": version_id,
+                }
+            ),
+        )
+
+        t4_lambda_es_ingest.handler(mock_event, mock_context)
+
+        stubber.assert_no_pending_responses()
+        mock_bulk.assert_called_once_with(mock_context, t4_lambda_es_ingest.es, b"test data")