diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index 0034ac81956..00e495bf99f 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -37,7 +37,7 @@ jobs:
- "4571:4571"
- "8080:8080"
env:
- SERVICES: kinesis,s3
+ SERVICES: kinesis,s3,sqs
options: >-
--health-cmd "curl -k https://localhost:4566"
--health-interval 10s
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index a5e6b2bdeeb..e6af8704d66 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -29,7 +29,7 @@ When you submit a pull request to the project, the CI system runs several verifi
You will be notified by email from the CI system if any issues are discovered, but if you want to run these checks locally before submitting PR or in order to verify changes you can use the following commands in the root directory:
1. To verify that all tests are passing, run `make test-all`.
2. To fix code style and format as well as catch common mistakes run `make fix`. Alternatively, run `make -k test-all docker-compose-down` to tear down the Docker services after running all the tests.
-3. To build docs run `make build-docs`.
+3. To build docs run `make build-rustdoc`.
# Development
@@ -58,7 +58,7 @@ Run `make test-all` to run all tests.
* `make fmt` - runs formatter, this command requires the nightly toolchain to be installed by running `rustup toolchain install nightly`.
* `make fix` - runs formatter and clippy checks.
* `make typos` - runs the spellcheck tool over the codebase. (Install by running `cargo install typos-cli`)
-* `make docs` - builds docs.
+* `make doc` - builds docs.
* `make docker-compose-up` - starts Docker services.
* `make docker-compose-down` - stops Docker services.
* `make docker-compose-logs` - shows Docker logs.
diff --git a/distribution/lambda/README.md b/distribution/lambda/README.md
index 48db36d878c..4ed5143831b 100644
--- a/distribution/lambda/README.md
+++ b/distribution/lambda/README.md
@@ -95,6 +95,12 @@ simplify the setup and avoid unstable deployments.
[1]: https://rust-lang-nursery.github.io/rust-cookbook/development_tools/debugging/config_log.html
+> [!TIP]
+> The Indexer Lambda's logging is quite verbose. To reduce the associated
+> CloudWatch costs, you can disable some lower level logs by setting the
+> `RUST_LOG` environment variable to `info,quickwit_actors=warn`, or disable
+> INFO logs altogether by setting `RUST_LOG=warn`.
+
Indexer only:
| Variable | Description | Default |
|---|---|---|
@@ -151,7 +157,13 @@ You can query and visualize the Quickwit Searcher Lambda from Grafana by using t
#### Configure Grafana data source
-You need to provide the following information.
+If you don't have a Grafana instance running yet, you can start one with the Quickwit plugin installed using Docker:
+
+```bash
+docker run -e GF_INSTALL_PLUGINS="quickwit-quickwit-datasource" -p 3000:3000 grafana/grafana
+```
+
+In the `Connections > Data sources` page, add a new Quickwit data source and configure the following settings:
|Variable|Description|Example|
|--|--|--|
@@ -159,4 +171,4 @@ You need to provide the following information.
|Custom HTTP Headers| If you configure API Gateway to require an API key, set `x-api-key` HTTP Header | Header: `x-api-key`
Value: API key value|
|Index ID| Same as `QW_LAMBDA_INDEX_ID` | hdfs-logs |
-After entering these values, click "Save & test" and you can now query your Quickwit Lambda from Grafana!
+After entering these values, click "Save & test". You can now query your Quickwit Lambda from Grafana!
diff --git a/distribution/lambda/cdk/cli.py b/distribution/lambda/cdk/cli.py
index c18fd14f289..ecb3ffdb155 100644
--- a/distribution/lambda/cdk/cli.py
+++ b/distribution/lambda/cdk/cli.py
@@ -320,14 +320,16 @@ def _clean_s3_bucket(bucket_name: str, prefix: str = ""):
print(f"Cleaning up bucket {bucket_name}/{prefix}...")
s3 = session.resource("s3")
bucket = s3.Bucket(bucket_name)
- bucket.objects.filter(Prefix=prefix).delete()
+ try:
+ bucket.objects.filter(Prefix=prefix).delete()
+ except s3.meta.client.exceptions.NoSuchBucket:
+ print(f"Bucket {bucket_name} not found, skipping cleanup")
def empty_hdfs_bucket():
bucket_name = _get_cloudformation_output_value(
app.HDFS_STACK_NAME, hdfs_stack.INDEX_STORE_BUCKET_NAME_EXPORT_NAME
)
-
_clean_s3_bucket(bucket_name)
diff --git a/distribution/lambda/cdk/stacks/examples/mock_data_stack.py b/distribution/lambda/cdk/stacks/examples/mock_data_stack.py
index 027b8afb98c..8a4a5c9290b 100644
--- a/distribution/lambda/cdk/stacks/examples/mock_data_stack.py
+++ b/distribution/lambda/cdk/stacks/examples/mock_data_stack.py
@@ -165,7 +165,12 @@ def __init__(
index_id=index_id,
index_config_bucket=index_config.s3_bucket_name,
index_config_key=index_config.s3_object_key,
- indexer_environment=lambda_env,
+ indexer_environment={
+ # the actor system is very verbose when the source is shutting
+ # down (each Lambda invocation)
+ "RUST_LOG": "info,quickwit_actors=warn",
+ **lambda_env,
+ },
searcher_environment=lambda_env,
indexer_package_location=indexer_package_location,
searcher_package_location=searcher_package_location,
diff --git a/docker-compose.yml b/docker-compose.yml
index 58b9d8b99f7..0667d9ac434 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -27,7 +27,7 @@ networks:
services:
localstack:
- image: localstack/localstack:${LOCALSTACK_VERSION:-2.3.2}
+ image: localstack/localstack:${LOCALSTACK_VERSION:-3.5.0}
container_name: localstack
ports:
- "${MAP_HOST_LOCALSTACK:-127.0.0.1}:4566:4566"
@@ -37,7 +37,7 @@ services:
- all
- localstack
environment:
- SERVICES: kinesis,s3
+ SERVICES: kinesis,s3,sqs
PERSISTENCE: 1
volumes:
- .localstack:/etc/localstack/init/ready.d
diff --git a/docs/assets/sqs-file-source.tf b/docs/assets/sqs-file-source.tf
new file mode 100644
index 00000000000..ffd348c1193
--- /dev/null
+++ b/docs/assets/sqs-file-source.tf
@@ -0,0 +1,134 @@
+terraform {
+ required_version = "1.7.5"
+ required_providers {
+ aws = {
+ source = "hashicorp/aws"
+ version = "~> 5.39.1"
+ }
+ }
+}
+
+provider "aws" {
+ region = "us-east-1"
+ default_tags {
+ tags = {
+ provisioner = "terraform"
+ author = "Quickwit"
+ }
+ }
+}
+
+locals {
+ sqs_notification_queue_name = "qw-tuto-s3-event-notifications"
+ source_bucket_name = "qw-tuto-source-bucket"
+}
+
+resource "aws_s3_bucket" "file_source" {
+ bucket_prefix = local.source_bucket_name
+ force_destroy = true
+}
+
+data "aws_iam_policy_document" "sqs_notification" {
+ statement {
+ effect = "Allow"
+
+ principals {
+ type = "*"
+ identifiers = ["*"]
+ }
+
+ actions = ["sqs:SendMessage"]
+ resources = ["arn:aws:sqs:*:*:${local.sqs_notification_queue_name}"]
+
+ condition {
+ test = "ArnEquals"
+ variable = "aws:SourceArn"
+ values = [aws_s3_bucket.file_source.arn]
+ }
+ }
+}
+
+
+resource "aws_sqs_queue" "s3_events" {
+ name = local.sqs_notification_queue_name
+ policy = data.aws_iam_policy_document.sqs_notification.json
+
+ redrive_policy = jsonencode({
+ deadLetterTargetArn = aws_sqs_queue.s3_events_deadletter.arn
+ maxReceiveCount = 5
+ })
+}
+
+resource "aws_sqs_queue" "s3_events_deadletter" {
+ name = "${locals.sqs_notification_queue_name}-deadletter"
+}
+
+resource "aws_sqs_queue_redrive_allow_policy" "s3_events_deadletter" {
+ queue_url = aws_sqs_queue.s3_events_deadletter.id
+
+ redrive_allow_policy = jsonencode({
+ redrivePermission = "byQueue",
+ sourceQueueArns = [aws_sqs_queue.s3_events.arn]
+ })
+}
+
+resource "aws_s3_bucket_notification" "bucket_notification" {
+ bucket = aws_s3_bucket.file_source.id
+
+ queue {
+ queue_arn = aws_sqs_queue.s3_events.arn
+ events = ["s3:ObjectCreated:*"]
+ }
+}
+
+data "aws_iam_policy_document" "quickwit_node" {
+ statement {
+ effect = "Allow"
+ actions = [
+ "sqs:ReceiveMessage",
+ "sqs:DeleteMessage",
+ "sqs:ChangeMessageVisibility",
+ "sqs:GetQueueAttributes",
+ ]
+ resources = [aws_sqs_queue.s3_events.arn]
+ }
+ statement {
+ effect = "Allow"
+ actions = ["s3:GetObject"]
+ resources = ["${aws_s3_bucket.file_source.arn}/*"]
+ }
+}
+
+resource "aws_iam_user" "quickwit_node" {
+ name = "quickwit-filesource-tutorial"
+ path = "/system/"
+}
+
+resource "aws_iam_user_policy" "quickwit_node" {
+ name = "quickwit-filesource-tutorial"
+ user = aws_iam_user.quickwit_node.name
+ policy = data.aws_iam_policy_document.quickwit_node.json
+}
+
+resource "aws_iam_access_key" "quickwit_node" {
+ user = aws_iam_user.quickwit_node.name
+}
+
+output "source_bucket_name" {
+ value = aws_s3_bucket.file_source.bucket
+
+}
+
+output "notification_queue_url" {
+ value = aws_sqs_queue.s3_events.id
+}
+
+output "quickwit_node_access_key_id" {
+ value = aws_iam_access_key.quickwit_node.id
+ sensitive = true
+}
+
+output "quickwit_node_secret_access_key" {
+ value = aws_iam_access_key.quickwit_node.secret
+ sensitive = true
+}
diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md
index 479c97e2365..83bdace6f96 100644
--- a/docs/configuration/source-config.md
+++ b/docs/configuration/source-config.md
@@ -29,15 +29,62 @@ The source type designates the kind of source being configured. As of version 0.
The source parameters indicate how to connect to a data store and are specific to the source type.
-### File source (CLI only)
+### File source
-A file source reads data from a local file. The file must consist of JSON objects separated by a newline (NDJSON).
-As of version 0.5, a file source can only be ingested with the [CLI command](/docs/reference/cli.md#tool-local-ingest). Compressed files (bz2, gzip, ...) and remote files (Amazon S3, HTTP, ...) are not supported.
+A file source reads data from files containing JSON objects separated by newlines (NDJSON). Gzip compression is supported provided that the file name ends with the `.gz` suffix.
+
+#### Ingest a single file (CLI only)
+
+To ingest a specific file, run the indexing directly in an adhoc CLI process with:
+
+```bash
+./quickwit tool local-ingest --index --input-path
+```
+
+Both local and object files are supported, provided that the environment is configured with the appropriate permissions. A tutorial is available [here](/docs/ingest-data/ingest-local-file.md).
+
+#### Notification based file ingestion (beta)
+
+Quickwit can automatically ingest all new files that are uploaded to an S3 bucket. This requires creating and configuring an [SQS notification queue](https://docs.aws.amazon.com/AmazonS3/latest/userguide/ways-to-add-notification-config-to-bucket.html). A complete example can be found [in this tutorial](/docs/ingest-data/sqs-files.md).
+
+
+The `notifications` parameter takes an array of notification settings. Currently one notifier can be configured per source and only the SQS notification `type` is supported.
+
+Required fields for the SQS `notifications` parameter items:
+- `type`: `sqs`
+- `queue_url`: complete URL of the SQS queue (e.g `https://sqs.us-east-1.amazonaws.com/123456789012/queue-name`)
+- `message_type`: format of the message payload, either
+ - `s3_notification`: an [S3 event notification](https://docs.aws.amazon.com/AmazonS3/latest/userguide/EventNotifications.html)
+ - `raw_uri`: a message containing just the file object URI (e.g. `s3://mybucket/mykey`)
+
+*Adding a file source with SQS notifications to an index with the [CLI](../reference/cli.md#source)*
```bash
-./quickwit tool local-ingest --input-path
+cat << EOF > source-config.yaml
+version: 0.8
+source_id: my-sqs-file-source
+source_type: file
+num_pipelines: 2
+params:
+ notifications:
+ - type: sqs
+ queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name
+ message_type: s3_notification
+EOF
+./quickwit source create --index my-index --source-config source-config.yaml
```
+:::note
+
+- Quickwit does not automatically delete the source files after a successful ingestion. You can use [S3 object expiration](https://docs.aws.amazon.com/AmazonS3/latest/userguide/lifecycle-expire-general-considerations.html) to configure how long they should be retained in the bucket.
+- Configure the notification to only forward events of type `s3:ObjectCreated:*`. Other events are acknowledged by the source without further processing and an warning is logged.
+- We strongly recommend using a [dead letter queue](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-dead-letter-queues.html) to receive all messages that couldn't be processed by the file source. A `maxReceiveCount` of 5 is a good default value. Here are some common situations where the notification message ends up in the dead letter queue:
+ - the notification message could not be parsed (e.g it is not a valid S3 notification)
+ - the file was not found
+ - the file is corrupted (e.g unexpected compression)
+
+:::
+
### Ingest API source
An ingest API source reads data from the [Ingest API](/docs/reference/rest-api.md#ingest-data-into-an-index). This source is automatically created at the index creation and cannot be deleted nor disabled.
diff --git a/docs/deployment/kubernetes/gke.md b/docs/deployment/kubernetes/gke.md
index d2b43cc5cc3..7c821779aa9 100644
--- a/docs/deployment/kubernetes/gke.md
+++ b/docs/deployment/kubernetes/gke.md
@@ -65,6 +65,10 @@ image:
pullPolicy: Always
tag: edge
+serviceAccount:
+ create: false
+ name: quickwit-sa
+
config:
default_index_root_uri: gs://{BUCKET}/qw-indexes
metastore_uri: gs://{BUCKET}/qw-indexes
diff --git a/docs/get-started/tutorials/tutorial-hdfs-logs.md b/docs/get-started/tutorials/tutorial-hdfs-logs.md
index dd544e1c652..23a941081ee 100644
--- a/docs/get-started/tutorials/tutorial-hdfs-logs.md
+++ b/docs/get-started/tutorials/tutorial-hdfs-logs.md
@@ -80,7 +80,7 @@ curl -o hdfs_logs_index_config.yaml https://raw.githubusercontent.com/quickwit-o
The index config defines five fields: `timestamp`, `tenant_id`, `severity_text`, `body`, and one JSON field
for the nested values `resource.service`, we could use an object field here and maintain a fixed schema, but for convenience we're going to use a JSON field.
It also sets the `default_search_fields`, the `tag_fields`, and the `timestamp_field`.
-The `timestamp_field` and `tag_fields` are used by Quickwit for [splits pruning](../../overview/architecture) at query time to boost search speed.
+The `timestamp_field` and `tag_fields` are used by Quickwit for [splits pruning](../../overview/concepts/querying.md#time-sharding) at query time to boost search speed.
Check out the [index config docs](../../configuration/index-config) for more details.
```yaml title="hdfs-logs-index.yaml"
diff --git a/docs/ingest-data/ingest-local-file.md b/docs/ingest-data/ingest-local-file.md
index 2a5b1bced03..6eb37e7c3eb 100644
--- a/docs/ingest-data/ingest-local-file.md
+++ b/docs/ingest-data/ingest-local-file.md
@@ -72,6 +72,12 @@ Clearing local cache directory...
✔ Documents successfully indexed.
```
+:::tip
+
+Object store URIs like `s3://mybucket/mykey.json` are also supported as `--input-path`, provided that your environment is configured with the appropriate permissions.
+
+:::
+
## Tear down resources (optional)
That's it! You can now tear down the resources you created. You can do so by running the following command:
diff --git a/docs/ingest-data/sqs-files.md b/docs/ingest-data/sqs-files.md
new file mode 100644
index 00000000000..ebca49629d7
--- /dev/null
+++ b/docs/ingest-data/sqs-files.md
@@ -0,0 +1,248 @@
+---
+title: S3 with SQS notifications
+description: A short tutorial describing how to set up Quickwit to ingest data from S3 files using an SQS notifier
+tags: [s3, sqs, integration]
+icon_url: /img/tutorials/file-ndjson.svg
+sidebar_position: 5
+---
+
+In this tutorial, we describe how to set up Quickwit to ingest data from S3
+with bucket notification events flowing through SQS. We will first create the
+AWS resources (S3 bucket, SQS queue, notifications) using terraform. We will
+then configure the Quickwit index and file source. Finally we will send some
+data to the source bucket and verify that it gets indexed.
+
+## AWS resources
+
+The complete terraform script can be downloaded [here](../assets/sqs-file-source.tf).
+
+First, create the bucket that will receive the source data files (NDJSON format):
+
+```
+resource "aws_s3_bucket" "file_source" {
+ bucket_prefix = "qw-tuto-source-bucket"
+}
+```
+
+Then setup the SQS queue that will carry the notifications when files are added
+to the bucket. The queue is configured with a policy that allows the source
+bucket to write the S3 notification messages to it. Also create a dead letter
+queue (DLQ) to receive the messages that couldn't be processed by the file
+source (e.g corrupted files). Messages are moved to the DLQ after 5 indexing
+attempts.
+
+```
+locals {
+ sqs_notification_queue_name = "qw-tuto-s3-event-notifications"
+}
+
+data "aws_iam_policy_document" "sqs_notification" {
+ statement {
+ effect = "Allow"
+
+ principals {
+ type = "*"
+ identifiers = ["*"]
+ }
+
+ actions = ["sqs:SendMessage"]
+ resources = ["arn:aws:sqs:*:*:${local.sqs_notification_queue_name}"]
+
+ condition {
+ test = "ArnEquals"
+ variable = "aws:SourceArn"
+ values = [aws_s3_bucket.file_source.arn]
+ }
+ }
+}
+
+resource "aws_sqs_queue" "s3_events_deadletter" {
+ name = "${locals.sqs_notification_queue_name}-deadletter"
+}
+
+resource "aws_sqs_queue" "s3_events" {
+ name = local.sqs_notification_queue_name
+ policy = data.aws_iam_policy_document.sqs_notification.json
+
+ redrive_policy = jsonencode({
+ deadLetterTargetArn = aws_sqs_queue.s3_events_deadletter.arn
+ maxReceiveCount = 5
+ })
+}
+
+resource "aws_sqs_queue_redrive_allow_policy" "s3_events_deadletter" {
+ queue_url = aws_sqs_queue.s3_events_deadletter.id
+
+ redrive_allow_policy = jsonencode({
+ redrivePermission = "byQueue",
+ sourceQueueArns = [aws_sqs_queue.s3_events.arn]
+ })
+}
+```
+
+Configure the bucket notification that writes messages to SQS each time a new
+file is created in the source bucket:
+
+```
+resource "aws_s3_bucket_notification" "bucket_notification" {
+ bucket = aws_s3_bucket.file_source.id
+
+ queue {
+ queue_arn = aws_sqs_queue.s3_events.arn
+ events = ["s3:ObjectCreated:*"]
+ }
+}
+```
+
+:::note
+
+Only events of type `s3:ObjectCreated:*` are supported. Other types (e.g.
+`ObjectRemoved`) are acknowledged and a warning is logged.
+
+:::
+
+The source needs to have access to both the notification queue and the source
+bucket. The following policy document contains the minimum permissions required
+by the source:
+
+```
+data "aws_iam_policy_document" "quickwit_node" {
+ statement {
+ effect = "Allow"
+ actions = [
+ "sqs:ReceiveMessage",
+ "sqs:DeleteMessage",
+ "sqs:ChangeMessageVisibility",
+ "sqs:GetQueueAttributes",
+ ]
+ resources = [aws_sqs_queue.s3_events.arn]
+ }
+ statement {
+ effect = "Allow"
+ actions = ["s3:GetObject"]
+ resources = ["${aws_s3_bucket.file_source.arn}/*"]
+ }
+}
+```
+
+Create the IAM user and credentials that will be used to
+associate this policy to your local Quickwit instance:
+
+```
+resource "aws_iam_user" "quickwit_node" {
+ name = "quickwit-filesource-tutorial"
+ path = "/system/"
+}
+
+resource "aws_iam_user_policy" "quickwit_node" {
+ name = "quickwit-filesource-tutorial"
+ user = aws_iam_user.quickwit_node.name
+ policy = data.aws_iam_policy_document.quickwit_node.json
+}
+
+resource "aws_iam_access_key" "quickwit_node" {
+ user = aws_iam_user.quickwit_node.name
+}
+```
+
+
+:::warning
+
+We don't recommend using IAM user credentials for running Quickwit nodes in
+production. This is just a simplified setup for the sake of the tutorial. When
+running on EC2/ECS, attach the policy document to an IAM roles instead.
+
+:::
+
+Download the [complete terraform script](../assets/sqs-file-source.tf) and
+deploy it using `terraform init` and `terraform apply`. After a successful
+execution, the outputs required to configure Quickwit will be listed. You can
+display the values of the sensitive outputs (key id and secret key) with:
+
+
+```bash
+terraform output quickwit_node_access_key_id
+terraform output quickwit_node_secret_access_key
+```
+
+## Run Quickwit
+
+[Install Quickwit locally](/docs/get-started/installation), then in your install
+directory, run Quickwit with the necessary access rights by replacing the
+`` and `` with the
+matching Terraform output values:
+
+```bash
+AWS_ACCESS_KEY_ID= \
+AWS_SECRET_ACCESS_KEY= \
+AWS_REGION=us-east-1 \
+./quickwit run
+```
+
+## Configure the index and the source
+
+In another terminal, in the Quickwit install directory, create an index:
+
+```bash
+cat << EOF > tutorial-sqs-file-index.yaml
+version: 0.7
+index_id: tutorial-sqs-file
+doc_mapping:
+ mode: dynamic
+indexing_settings:
+ commit_timeout_secs: 30
+EOF
+
+./quickwit index create --index-config tutorial-sqs-file-index.yaml
+```
+
+Replacing `` with the corresponding Terraform output
+value, create a file source for that index:
+
+```bash
+cat << EOF > tutorial-sqs-file-source.yaml
+version: 0.8
+source_id: sqs-filesource
+source_type: file
+num_pipelines: 2
+params:
+ notifications:
+ - type: sqs
+ queue_url:
+ message_type: s3_notification
+EOF
+
+./quickwit source create --index tutorial-sqs-file --source-config tutorial-sqs-file-source.yaml
+```
+
+:::tip
+
+The `num_pipeline` configuration controls how many consumers will poll from the queue in parallel. Choose the number according to the indexer compute resources you want to dedicate to this source. As a rule of thumb, configure 1 pipeline for every 2 cores.
+
+:::
+
+## Ingest data
+
+We can now ingest data into Quickwit by uploading files to S3. If you have the
+AWS CLI installed, run the following command, replacing ``
+with the associated Terraform output:
+
+```bash
+curl https://quickwit-datasets-public.s3.amazonaws.com/hdfs-logs-multitenants-10000.json | \
+ aws s3 cp - s3:///hdfs-logs-multitenants-10000.json
+```
+
+If you prefer not to use the AWS CLI, you can also download the file and upload
+it manually to the source bucket using the AWS console.
+
+Wait approximately 1 minute and the data should appear in the index:
+
+```bash
+./quickwit index describe --index tutorial-sqs-file
+```
+
+## Tear down the resources
+
+The AWS resources instantiated in this tutorial don't incur any fixed costs, but
+we still recommend deleting them when you are done. In the directory with the
+Terraform script, run `terraform destroy`.
diff --git a/docs/overview/concepts/indexing.md b/docs/overview/concepts/indexing.md
index 4feba085b41..6dab81f392a 100644
--- a/docs/overview/concepts/indexing.md
+++ b/docs/overview/concepts/indexing.md
@@ -30,16 +30,7 @@ The disk space allocated to the split store is controlled by the config paramete
## Data sources
-A data source designates the location and set of parameters that allow to connect to and ingest data from an external data store, which can be a file, a stream, or a database. Often, Quickwit simply refers to data sources as "sources". The indexing engine supports file-based and stream-based sources. Quickwit can insert data into an index from one or multiple sources, defined in the index config.
-
-
-### File sources
-
-File sources are sources that read data from a file stored on the local file system.
-
-### Streaming sources
-
-Streaming sources are sources that read data from a streaming service such as Apache Kafka. As of version 0.2, Quickwit only supports Apache Kafka. Future versions of Quickwit will support additional streaming services such as Amazon Kinesis.
+A data source designates the location and set of parameters that allow to connect to and ingest data from an external data store, which can be a file, a stream, or a database. Often, Quickwit simply refers to data sources as "sources". The indexing engine supports local adhoc file ingests using [the CLI](/docs/reference/cli#tool-local-ingest) and streaming sources (e.g. the Kafka source). Quickwit can insert data into an index from one or multiple sources. More details can be found [in the source configuration page](https://quickwit.io/docs/configuration/source-config).
## Checkpoint
diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock
index 0db9bbed0de..a8d34e1406a 100644
--- a/quickwit/Cargo.lock
+++ b/quickwit/Cargo.lock
@@ -458,6 +458,28 @@ dependencies = [
"url",
]
+[[package]]
+name = "aws-sdk-sqs"
+version = "1.36.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3587fbaf540d65337c2356ebf3f78fba160025b3d69634175f1ea3a7895738e9"
+dependencies = [
+ "aws-credential-types",
+ "aws-runtime",
+ "aws-smithy-async",
+ "aws-smithy-http",
+ "aws-smithy-json",
+ "aws-smithy-runtime",
+ "aws-smithy-runtime-api",
+ "aws-smithy-types",
+ "aws-types",
+ "bytes",
+ "http 0.2.12",
+ "once_cell",
+ "regex-lite",
+ "tracing",
+]
+
[[package]]
name = "aws-sdk-sso"
version = "1.36.0"
@@ -1131,9 +1153,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]]
name = "bytes"
-version = "1.6.0"
+version = "1.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9"
+checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50"
dependencies = [
"serde",
]
@@ -2026,12 +2048,6 @@ version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
-[[package]]
-name = "dotenv"
-version = "0.15.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f"
-
[[package]]
name = "dotenvy"
version = "0.15.7"
@@ -3517,9 +3533,9 @@ dependencies = [
[[package]]
name = "lambda_runtime"
-version = "0.11.3"
+version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9be8f0e7a5db270feb93a7a3593c22a4c5fb8e8f260f5f490e0c3a5ffeb009db"
+checksum = "ed49669d6430292aead991e19bf13153135a884f916e68f32997c951af637ebe"
dependencies = [
"async-stream",
"base64 0.22.1",
@@ -4717,7 +4733,7 @@ checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
[[package]]
name = "ownedbytes"
version = "0.7.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"stable_deref_trait",
]
@@ -5627,6 +5643,7 @@ dependencies = [
"aws-config",
"aws-sdk-kinesis",
"aws-sdk-s3",
+ "aws-sdk-sqs",
"aws-smithy-async",
"aws-smithy-runtime",
"aws-types",
@@ -5952,6 +5969,7 @@ dependencies = [
"async-compression",
"async-trait",
"aws-sdk-kinesis",
+ "aws-sdk-sqs",
"bytes",
"bytesize",
"criterion",
@@ -5988,6 +6006,7 @@ dependencies = [
"quickwit-storage",
"rand 0.8.5",
"rdkafka",
+ "regex",
"reqwest",
"serde",
"serde_json",
@@ -6045,6 +6064,7 @@ name = "quickwit-integration-tests"
version = "0.8.0"
dependencies = [
"anyhow",
+ "aws-sdk-sqs",
"futures-util",
"hyper 0.14.29",
"itertools 0.13.0",
@@ -6052,6 +6072,7 @@ dependencies = [
"quickwit-cli",
"quickwit-common",
"quickwit-config",
+ "quickwit-indexing",
"quickwit-metastore",
"quickwit-proto",
"quickwit-rest-client",
@@ -6063,6 +6084,7 @@ dependencies = [
"tokio",
"tonic",
"tracing",
+ "tracing-subscriber",
]
[[package]]
@@ -6140,7 +6162,7 @@ dependencies = [
"flate2",
"http 0.2.12",
"lambda_http",
- "lambda_runtime 0.11.3",
+ "lambda_runtime 0.13.0",
"mime_guess",
"once_cell",
"opentelemetry",
@@ -6189,7 +6211,7 @@ dependencies = [
"async-trait",
"bytes",
"bytesize",
- "dotenv",
+ "dotenvy",
"futures",
"http 0.2.12",
"itertools 0.13.0",
@@ -8137,7 +8159,7 @@ dependencies = [
[[package]]
name = "tantivy"
version = "0.23.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"aho-corasick",
"arc-swap",
@@ -8190,7 +8212,7 @@ dependencies = [
[[package]]
name = "tantivy-bitpacker"
version = "0.6.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"bitpacking",
]
@@ -8198,7 +8220,7 @@ dependencies = [
[[package]]
name = "tantivy-columnar"
version = "0.3.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"downcast-rs",
"fastdivide",
@@ -8213,7 +8235,7 @@ dependencies = [
[[package]]
name = "tantivy-common"
version = "0.7.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"async-trait",
"byteorder",
@@ -8236,7 +8258,7 @@ dependencies = [
[[package]]
name = "tantivy-query-grammar"
version = "0.22.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"nom",
]
@@ -8244,7 +8266,7 @@ dependencies = [
[[package]]
name = "tantivy-sstable"
version = "0.3.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"tantivy-bitpacker",
"tantivy-common",
@@ -8255,7 +8277,7 @@ dependencies = [
[[package]]
name = "tantivy-stacker"
version = "0.3.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"murmurhash32",
"rand_distr",
@@ -8265,7 +8287,7 @@ dependencies = [
[[package]]
name = "tantivy-tokenizer-api"
version = "0.3.0"
-source = "git+https://github.com/quickwit-oss/tantivy/?rev=13e9885#13e9885dfda8cebf4bfef72f53bf811da8549445"
+source = "git+https://github.com/quickwit-oss/tantivy/?rev=c71ec80#c71ec8086d6563c4bb7e573182a26b280a3ac519"
dependencies = [
"serde",
]
@@ -8471,9 +8493,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokio"
-version = "1.39.1"
+version = "1.39.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d040ac2b29ab03b09d4129c2f5bbd012a3ac2f79d38ff506a4bf8dd34b0eac8a"
+checksum = "daa4fb1bc778bd6f04cbfc4bb2d06a7396a8f299dc33ea1900cedaa316f467b1"
dependencies = [
"backtrace",
"bytes",
diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml
index 025f8e193ac..51b3df6d541 100644
--- a/quickwit/Cargo.toml
+++ b/quickwit/Cargo.toml
@@ -101,7 +101,7 @@ console-subscriber = "0.1.8"
criterion = { version = "0.5", features = ["async_tokio"] }
cron = "0.12.0"
dialoguer = "0.10.3"
-dotenv = "0.15"
+dotenvy = "0.15"
dyn-clone = "1.0.10"
enum-iterator = "1.5"
env_logger = "0.10"
@@ -279,6 +279,7 @@ aws-config = "1.5.4"
aws-credential-types = { version = "1.2", features = ["hardcoded-credentials"] }
aws-sdk-kinesis = "1.36"
aws-sdk-s3 = "1.42"
+aws-sdk-sqs = "1.36"
aws-smithy-async = "1.2"
aws-smithy-runtime = "1.6.2"
aws-smithy-types = { version = "1.2", features = ["byte-stream-poll-next"] }
@@ -324,7 +325,7 @@ quickwit-serve = { path = "quickwit-serve" }
quickwit-storage = { path = "quickwit-storage" }
quickwit-telemetry = { path = "quickwit-telemetry" }
-tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "13e9885", default-features = false, features = [
+tantivy = { git = "https://github.com/quickwit-oss/tantivy/", rev = "c71ec80", default-features = false, features = [
"lz4-compression",
"mmap",
"quickwit",
diff --git a/quickwit/deny.toml b/quickwit/deny.toml
index 0c6c498786b..139785b1c9a 100644
--- a/quickwit/deny.toml
+++ b/quickwit/deny.toml
@@ -9,6 +9,7 @@
# The values provided in this template are the default values that will be used
# when any section or field is not specified in your own configuration
+[graph]
# If 1 or more target triples (and optionally, target_features) are specified,
# only the specified targets will be checked when running `cargo deny check`.
# This means, if a particular package is only ever used as a target specific
@@ -31,42 +32,22 @@ targets = [
# More documentation for the advisories section can be found here:
# https://embarkstudios.github.io/cargo-deny/checks/advisories/cfg.html
[advisories]
+version = 2
# The path where the advisory database is cloned/fetched into
db-path = "~/.cargo/advisory-db"
# The url(s) of the advisory databases to use
db-urls = ["https://github.com/rustsec/advisory-db"]
-# The lint level for security vulnerabilities
-vulnerability = "deny"
-# The lint level for unmaintained crates
-unmaintained = "warn"
-# The lint level for crates that have been yanked from their source registry
-yanked = "warn"
-# The lint level for crates with security notices. Note that as of
-# 2019-12-17 there are no security notice advisories in
-# https://github.com/rustsec/advisory-db
-notice = "warn"
# A list of advisory IDs to ignore. Note that ignored advisories will still
# output a note when they are encountered.
ignore = [
- # TODO Remove me after rsa gets patched and released.
- "RUSTSEC-2023-0071"
+ "RUSTSEC-2021-0153", # `encoding` is unmaintained, it's used in lindera
]
-# Threshold for security vulnerabilities, any vulnerability with a CVSS score
-# lower than the range specified will be ignored. Note that ignored advisories
-# will still output a note when they are encountered.
-# * None - CVSS Score 0.0
-# * Low - CVSS Score 0.1 - 3.9
-# * Medium - CVSS Score 4.0 - 6.9
-# * High - CVSS Score 7.0 - 8.9
-# * Critical - CVSS Score 9.0 - 10.0
-#severity-threshold =
# This section is considered when running `cargo deny check licenses`
# More documentation for the licenses section can be found here:
# https://embarkstudios.github.io/cargo-deny/checks/licenses/cfg.html
[licenses]
-# The lint level for crates which do not have a detectable license
-unlicensed = "deny"
+version = 2
# List of explicitly allowed licenses
# See https://spdx.org/licenses/ for list of possible licenses
# [possible values: any SPDX 3.11 short identifier (+ optional exception)].
@@ -86,26 +67,6 @@ allow = [
"Zlib",
"zlib-acknowledgement",
]
-# List of explicitly disallowed licenses
-# See https://spdx.org/licenses/ for list of possible licenses
-# [possible values: any SPDX 3.11 short identifier (+ optional exception)].
-deny = [
- #"Nokia",
-]
-# Lint level for licenses considered copyleft
-copyleft = "warn"
-# Blanket approval or denial for OSI-approved or FSF Free/Libre licenses
-# * both - The license will be approved if it is both OSI-approved *AND* FSF
-# * either - The license will be approved if it is either OSI-approved *OR* FSF
-# * osi-only - The license will be approved if is OSI-approved *AND NOT* FSF
-# * fsf-only - The license will be approved if is FSF *AND NOT* OSI-approved
-# * neither - This predicate is ignored and the default lint level is used
-allow-osi-fsf-free = "neither"
-# Lint level used when no other predicates are matched
-# 1. License isn't in the allow or deny lists
-# 2. License isn't copyleft
-# 3. License isn't OSI/FSF, or allow-osi-fsf-free = "neither"
-default = "deny"
# The confidence threshold for detecting a license from license text.
# The higher the value, the more closely the license text must be to the
# canonical license text of a valid SPDX license file.
@@ -114,20 +75,30 @@ confidence-threshold = 0.8
# Allow 1 or more licenses on a per-crate basis, so that particular licenses
# aren't accepted for every possible crate as with the normal allow list
exceptions = [
+ { name = "quickwit-actors", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-aws", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-cli", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-cluster", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-codegen", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-codegen-example", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-common", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-config", allow = ["AGPL-3.0"], version = "*" },
- { name = "quickwit-index-management", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-control-plane", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-datetime", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-directories", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-doc-mapper", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-indexing", allow = ["AGPL-3.0"], version = "*" },
- { name = "quickwit-ingest-api", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-index-management", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-ingest", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-integration-tests", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-jaeger", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-janitor", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-lambda", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-macros", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-metastore", allow = ["AGPL-3.0"], version = "*" },
- { name = "quickwit-metastore-utils", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-opentelemetry", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-proto", allow = ["AGPL-3.0"], version = "*" },
+ { name = "quickwit-query", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-rest-client", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-search", allow = ["AGPL-3.0"], version = "*" },
{ name = "quickwit-serve", allow = ["AGPL-3.0"], version = "*" },
@@ -224,8 +195,8 @@ allow-git = []
[sources.allow-org]
# 1 or more github.com organizations to allow git sources for
-github = [""]
+github = ["quickwit-oss"]
# 1 or more gitlab.com organizations to allow git sources for
-gitlab = [""]
+gitlab = []
# 1 or more bitbucket.org organizations to allow git sources for
-bitbucket = [""]
+bitbucket = []
diff --git a/quickwit/quickwit-actors/src/actor_context.rs b/quickwit/quickwit-actors/src/actor_context.rs
index af7a8a3f7c9..0dbe1194bba 100644
--- a/quickwit/quickwit-actors/src/actor_context.rs
+++ b/quickwit/quickwit-actors/src/actor_context.rs
@@ -339,8 +339,11 @@ impl ActorContext {
self.self_mailbox.try_send_message(msg)
}
- /// Schedules a message that will be sent to the high-priority
- /// queue of the actor Mailbox once `after_duration` has elapsed.
+ /// Schedules a message that will be sent to the high-priority queue of the
+ /// actor Mailbox once `after_duration` has elapsed.
+ ///
+ /// Note that this holds a reference to the actor mailbox until the message
+ /// is actually sent.
pub fn schedule_self_msg(&self, after_duration: Duration, message: M)
where
A: DeferableReplyHandler,
diff --git a/quickwit/quickwit-aws/Cargo.toml b/quickwit/quickwit-aws/Cargo.toml
index f67d36155b4..19bf0ceb6e0 100644
--- a/quickwit/quickwit-aws/Cargo.toml
+++ b/quickwit/quickwit-aws/Cargo.toml
@@ -14,6 +14,7 @@ license.workspace = true
aws-config = { workspace = true }
aws-sdk-kinesis = { workspace = true, optional = true }
aws-sdk-s3 = { workspace = true }
+aws-sdk-sqs = { workspace = true, optional = true }
aws-smithy-async = { workspace = true }
aws-smithy-runtime = { workspace = true }
aws-types = { workspace = true }
@@ -27,3 +28,4 @@ quickwit-common = { workspace = true }
[features]
kinesis = ["aws-sdk-kinesis"]
+sqs = ["aws-sdk-sqs"]
diff --git a/quickwit/quickwit-aws/src/error.rs b/quickwit/quickwit-aws/src/error.rs
index e7c6dfdd077..ba2e620a27e 100644
--- a/quickwit/quickwit-aws/src/error.rs
+++ b/quickwit/quickwit-aws/src/error.rs
@@ -19,13 +19,6 @@
#![allow(clippy::match_like_matches_macro)]
-#[cfg(feature = "kinesis")]
-use aws_sdk_kinesis::operation::{
- create_stream::CreateStreamError, delete_stream::DeleteStreamError,
- describe_stream::DescribeStreamError, get_records::GetRecordsError,
- get_shard_iterator::GetShardIteratorError, list_shards::ListShardsError,
- list_streams::ListStreamsError, merge_shards::MergeShardsError, split_shard::SplitShardError,
-};
use aws_sdk_s3::error::SdkError;
use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError;
use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError;
@@ -109,89 +102,124 @@ impl AwsRetryable for HeadObjectError {
}
#[cfg(feature = "kinesis")]
-impl AwsRetryable for GetRecordsError {
- fn is_retryable(&self) -> bool {
- match self {
- GetRecordsError::KmsThrottlingException(_) => true,
- GetRecordsError::ProvisionedThroughputExceededException(_) => true,
- _ => false,
+mod kinesis {
+ use aws_sdk_kinesis::operation::create_stream::CreateStreamError;
+ use aws_sdk_kinesis::operation::delete_stream::DeleteStreamError;
+ use aws_sdk_kinesis::operation::describe_stream::DescribeStreamError;
+ use aws_sdk_kinesis::operation::get_records::GetRecordsError;
+ use aws_sdk_kinesis::operation::get_shard_iterator::GetShardIteratorError;
+ use aws_sdk_kinesis::operation::list_shards::ListShardsError;
+ use aws_sdk_kinesis::operation::list_streams::ListStreamsError;
+ use aws_sdk_kinesis::operation::merge_shards::MergeShardsError;
+ use aws_sdk_kinesis::operation::split_shard::SplitShardError;
+
+ use super::*;
+
+ impl AwsRetryable for GetRecordsError {
+ fn is_retryable(&self) -> bool {
+ match self {
+ GetRecordsError::KmsThrottlingException(_) => true,
+ GetRecordsError::ProvisionedThroughputExceededException(_) => true,
+ _ => false,
+ }
}
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for GetShardIteratorError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- GetShardIteratorError::ProvisionedThroughputExceededException(_)
- )
+ impl AwsRetryable for GetShardIteratorError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ GetShardIteratorError::ProvisionedThroughputExceededException(_)
+ )
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for ListShardsError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- ListShardsError::ResourceInUseException(_) | ListShardsError::LimitExceededException(_)
- )
+ impl AwsRetryable for ListShardsError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ ListShardsError::ResourceInUseException(_)
+ | ListShardsError::LimitExceededException(_)
+ )
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for CreateStreamError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- CreateStreamError::ResourceInUseException(_)
- | CreateStreamError::LimitExceededException(_)
- )
+ impl AwsRetryable for CreateStreamError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ CreateStreamError::ResourceInUseException(_)
+ | CreateStreamError::LimitExceededException(_)
+ )
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for DeleteStreamError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- DeleteStreamError::ResourceInUseException(_)
- | DeleteStreamError::LimitExceededException(_)
- )
+ impl AwsRetryable for DeleteStreamError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ DeleteStreamError::ResourceInUseException(_)
+ | DeleteStreamError::LimitExceededException(_)
+ )
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for DescribeStreamError {
- fn is_retryable(&self) -> bool {
- matches!(self, DescribeStreamError::LimitExceededException(_))
+ impl AwsRetryable for DescribeStreamError {
+ fn is_retryable(&self) -> bool {
+ matches!(self, DescribeStreamError::LimitExceededException(_))
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for ListStreamsError {
- fn is_retryable(&self) -> bool {
- matches!(self, ListStreamsError::LimitExceededException(_))
+ impl AwsRetryable for ListStreamsError {
+ fn is_retryable(&self) -> bool {
+ matches!(self, ListStreamsError::LimitExceededException(_))
+ }
}
-}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for MergeShardsError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- MergeShardsError::ResourceInUseException(_)
- | MergeShardsError::LimitExceededException(_)
- )
+ impl AwsRetryable for MergeShardsError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ MergeShardsError::ResourceInUseException(_)
+ | MergeShardsError::LimitExceededException(_)
+ )
+ }
+ }
+
+ impl AwsRetryable for SplitShardError {
+ fn is_retryable(&self) -> bool {
+ matches!(
+ self,
+ SplitShardError::ResourceInUseException(_)
+ | SplitShardError::LimitExceededException(_)
+ )
+ }
}
}
-#[cfg(feature = "kinesis")]
-impl AwsRetryable for SplitShardError {
- fn is_retryable(&self) -> bool {
- matches!(
- self,
- SplitShardError::ResourceInUseException(_) | SplitShardError::LimitExceededException(_)
- )
+#[cfg(feature = "sqs")]
+mod sqs {
+ use aws_sdk_sqs::operation::change_message_visibility::ChangeMessageVisibilityError;
+ use aws_sdk_sqs::operation::delete_message_batch::DeleteMessageBatchError;
+ use aws_sdk_sqs::operation::receive_message::ReceiveMessageError;
+
+ use super::*;
+
+ impl AwsRetryable for ReceiveMessageError {
+ fn is_retryable(&self) -> bool {
+ false
+ }
+ }
+
+ impl AwsRetryable for DeleteMessageBatchError {
+ fn is_retryable(&self) -> bool {
+ false
+ }
+ }
+
+ impl AwsRetryable for ChangeMessageVisibilityError {
+ fn is_retryable(&self) -> bool {
+ false
+ }
}
}
diff --git a/quickwit/quickwit-cli/Cargo.toml b/quickwit/quickwit-cli/Cargo.toml
index 36cfc371aa7..d403eef2922 100644
--- a/quickwit/quickwit-cli/Cargo.toml
+++ b/quickwit/quickwit-cli/Cargo.toml
@@ -92,6 +92,7 @@ release-feature-set = [
"quickwit-indexing/kafka",
"quickwit-indexing/kinesis",
"quickwit-indexing/pulsar",
+ "quickwit-indexing/sqs",
"quickwit-indexing/vrl",
"quickwit-storage/azure",
"quickwit-storage/gcs",
@@ -104,6 +105,7 @@ release-feature-vendored-set = [
"pprof",
"quickwit-indexing/kinesis",
"quickwit-indexing/pulsar",
+ "quickwit-indexing/sqs",
"quickwit-indexing/vrl",
"quickwit-indexing/vendored-kafka",
"quickwit-storage/azure",
@@ -116,6 +118,7 @@ release-macos-feature-vendored-set = [
"openssl-support",
"quickwit-indexing/kinesis",
"quickwit-indexing/pulsar",
+ "quickwit-indexing/sqs",
"quickwit-indexing/vrl",
"quickwit-indexing/vendored-kafka-macos",
"quickwit-storage/azure",
diff --git a/quickwit/quickwit-cli/src/source.rs b/quickwit/quickwit-cli/src/source.rs
index 1a1948fdd99..ee90689ca94 100644
--- a/quickwit/quickwit-cli/src/source.rs
+++ b/quickwit/quickwit-cli/src/source.rs
@@ -744,7 +744,7 @@ mod tests {
source_id: "foo-source".to_string(),
num_pipelines: NonZeroUsize::new(1).unwrap(),
enabled: true,
- source_params: SourceParams::file("path/to/file"),
+ source_params: SourceParams::file_from_str("path/to/file").unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
}];
@@ -753,9 +753,10 @@ mod tests {
source_type: "file".to_string(),
enabled: "true".to_string(),
}];
+ let expected_uri = Uri::from_str("path/to/file").unwrap();
let expected_params = vec![ParamsRow {
key: "filepath".to_string(),
- value: JsonValue::String("path/to/file".to_string()),
+ value: JsonValue::String(expected_uri.to_string()),
}];
let expected_checkpoint = vec![
CheckpointRow {
@@ -820,12 +821,12 @@ mod tests {
let expected_sources = [
SourceRow {
source_id: "bar-source".to_string(),
- source_type: "file".to_string(),
+ source_type: "stdin".to_string(),
enabled: "true".to_string(),
},
SourceRow {
source_id: "foo-source".to_string(),
- source_type: "file".to_string(),
+ source_type: "stdin".to_string(),
enabled: "true".to_string(),
},
];
diff --git a/quickwit/quickwit-cli/src/tool.rs b/quickwit/quickwit-cli/src/tool.rs
index b936cbc4897..c7ab1911205 100644
--- a/quickwit/quickwit-cli/src/tool.rs
+++ b/quickwit/quickwit-cli/src/tool.rs
@@ -173,7 +173,7 @@ pub fn build_tool_command() -> Command {
pub struct LocalIngestDocsArgs {
pub config_uri: Uri,
pub index_id: IndexId,
- pub input_path_opt: Option,
+ pub input_path_opt: Option,
pub input_format: SourceInputFormat,
pub overwrite: bool,
pub vrl_script: Option,
@@ -251,9 +251,7 @@ impl ToolCliCommand {
.remove_one::("index")
.expect("`index` should be a required arg.");
let input_path_opt = if let Some(input_path) = matches.remove_one::("input-path") {
- Uri::from_str(&input_path)?
- .filepath()
- .map(|path| path.to_path_buf())
+ Some(Uri::from_str(&input_path)?)
} else {
None
};
@@ -410,8 +408,8 @@ pub async fn local_ingest_docs_cli(args: LocalIngestDocsArgs) -> anyhow::Result<
get_resolvers(&config.storage_configs, &config.metastore_configs);
let mut metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
- let source_params = if let Some(filepath) = args.input_path_opt.as_ref() {
- SourceParams::file(filepath)
+ let source_params = if let Some(uri) = args.input_path_opt.as_ref() {
+ SourceParams::file_from_uri(uri.clone())
} else {
SourceParams::stdin()
};
diff --git a/quickwit/quickwit-cli/tests/cli.rs b/quickwit/quickwit-cli/tests/cli.rs
index 205fd778a85..524098537b6 100644
--- a/quickwit/quickwit-cli/tests/cli.rs
+++ b/quickwit/quickwit-cli/tests/cli.rs
@@ -26,7 +26,7 @@ use std::path::Path;
use anyhow::Result;
use clap::error::ErrorKind;
-use helpers::{TestEnv, TestStorageType};
+use helpers::{uri_from_path, TestEnv, TestStorageType};
use quickwit_cli::checklist::ChecklistError;
use quickwit_cli::cli::build_cli;
use quickwit_cli::index::{
@@ -38,6 +38,7 @@ use quickwit_cli::tool::{
};
use quickwit_common::fs::get_cache_directory_path;
use quickwit_common::rand::append_random_suffix;
+use quickwit_common::uri::Uri;
use quickwit_config::{RetentionPolicy, SourceInputFormat, CLI_SOURCE_ID};
use quickwit_metastore::{
ListSplitsRequestExt, MetastoreResolver, MetastoreServiceExt, MetastoreServiceStreamSplitsExt,
@@ -62,11 +63,11 @@ async fn create_logs_index(test_env: &TestEnv) -> anyhow::Result<()> {
create_index_cli(args).await
}
-async fn local_ingest_docs(input_path: &Path, test_env: &TestEnv) -> anyhow::Result<()> {
+async fn local_ingest_docs(uri: Uri, test_env: &TestEnv) -> anyhow::Result<()> {
let args = LocalIngestDocsArgs {
config_uri: test_env.resource_files.config.clone(),
index_id: test_env.index_id.clone(),
- input_path_opt: Some(input_path.to_path_buf()),
+ input_path_opt: Some(uri),
input_format: SourceInputFormat::Json,
overwrite: false,
clear_cache: true,
@@ -75,6 +76,10 @@ async fn local_ingest_docs(input_path: &Path, test_env: &TestEnv) -> anyhow::Res
local_ingest_docs_cli(args).await
}
+async fn local_ingest_log_docs(test_env: &TestEnv) -> anyhow::Result<()> {
+ local_ingest_docs(test_env.resource_files.log_docs.clone(), test_env).await
+}
+
#[test]
fn test_cmd_help() {
let cmd = build_cli();
@@ -253,14 +258,17 @@ async fn test_ingest_docs_cli() {
// Ensure cache directory is empty.
let cache_directory_path = get_cache_directory_path(&test_env.data_dir_path);
-
assert!(cache_directory_path.read_dir().unwrap().next().is_none());
+ let does_not_exist_uri = uri_from_path(&test_env.data_dir_path)
+ .join("file-does-not-exist.json")
+ .unwrap();
+
// Ingest a non-existing file should fail.
let args = LocalIngestDocsArgs {
config_uri: test_env.resource_files.config,
index_id: test_env.index_id,
- input_path_opt: Some(test_env.data_dir_path.join("file-does-not-exist.json")),
+ input_path_opt: Some(does_not_exist_uri),
input_format: SourceInputFormat::Json,
overwrite: false,
clear_cache: true,
@@ -333,9 +341,7 @@ async fn test_cmd_search_aggregation() {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let aggregation: Value = json!(
{
@@ -433,9 +439,7 @@ async fn test_cmd_search_with_snippets() -> Result<()> {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
// search with snippets
let args = SearchIndexArgs {
@@ -488,9 +492,7 @@ async fn test_search_index_cli() {
sort_by_score: false,
};
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let args = create_search_args("level:info");
@@ -601,9 +603,7 @@ async fn test_delete_index_cli_dry_run() {
.unwrap();
assert!(metastore.index_exists(&index_id).await.unwrap());
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
// On non-empty index
let args = create_delete_args(true);
@@ -627,9 +627,7 @@ async fn test_delete_index_cli() {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let args = DeleteIndexArgs {
client_args: test_env.default_client_args(),
@@ -653,9 +651,7 @@ async fn test_garbage_collect_cli_no_grace() {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
let index_uid = test_env.index_metadata().await.unwrap().index_uid;
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let metastore = MetastoreResolver::unconfigured()
.resolve(&test_env.metastore_uri)
@@ -763,9 +759,7 @@ async fn test_garbage_collect_index_cli() {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
let index_uid = test_env.index_metadata().await.unwrap().index_uid;
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let refresh_metastore = |metastore| async {
// In this test we rely on the file backed metastore and
@@ -915,9 +909,7 @@ async fn test_all_local_index() {
.unwrap();
assert!(metadata_file_exists);
- local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env)
- .await
- .unwrap();
+ local_ingest_log_docs(&test_env).await.unwrap();
let query_response = reqwest::get(format!(
"http://127.0.0.1:{}/api/v1/{}/search?query=level:info",
@@ -971,16 +963,21 @@ async fn test_all_with_s3_localstack_cli() {
test_env.start_server().await.unwrap();
create_logs_index(&test_env).await.unwrap();
- let s3_path = upload_test_file(
+ let s3_uri = upload_test_file(
test_env.storage_resolver.clone(),
- test_env.resource_files.log_docs.clone(),
+ test_env
+ .resource_files
+ .log_docs
+ .filepath()
+ .unwrap()
+ .to_path_buf(),
"quickwit-integration-tests",
"sources/",
&append_random_suffix("test-all--cli-s3-localstack"),
)
.await;
- local_ingest_docs(&s3_path, &test_env).await.unwrap();
+ local_ingest_docs(s3_uri, &test_env).await.unwrap();
// Cli search
let args = SearchIndexArgs {
diff --git a/quickwit/quickwit-cli/tests/helpers.rs b/quickwit/quickwit-cli/tests/helpers.rs
index 0a52f6b9792..67839f7368b 100644
--- a/quickwit/quickwit-cli/tests/helpers.rs
+++ b/quickwit/quickwit-cli/tests/helpers.rs
@@ -17,9 +17,8 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
-use std::borrow::Borrow;
use std::fs;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::Arc;
@@ -114,8 +113,8 @@ pub struct TestResourceFiles {
pub index_config: Uri,
pub index_config_without_uri: Uri,
pub index_config_with_retention: Uri,
- pub log_docs: PathBuf,
- pub wikipedia_docs: PathBuf,
+ pub log_docs: Uri,
+ pub wikipedia_docs: Uri,
}
/// A struct to hold few info about the test environment.
@@ -192,8 +191,8 @@ pub enum TestStorageType {
LocalFileSystem,
}
-fn uri_from_path(path: PathBuf) -> Uri {
- Uri::from_str(&format!("file://{}", path.display())).unwrap()
+pub fn uri_from_path(path: &Path) -> Uri {
+ Uri::from_str(path.to_str().unwrap()).unwrap()
}
/// Creates all necessary artifacts in a test environment.
@@ -265,12 +264,12 @@ pub async fn create_test_env(
.context("failed to parse cluster endpoint")?;
let resource_files = TestResourceFiles {
- config: uri_from_path(node_config_path),
- index_config: uri_from_path(index_config_path),
- index_config_without_uri: uri_from_path(index_config_without_uri_path),
- index_config_with_retention: uri_from_path(index_config_with_retention_path),
- log_docs: log_docs_path,
- wikipedia_docs: wikipedia_docs_path,
+ config: uri_from_path(&node_config_path),
+ index_config: uri_from_path(&index_config_path),
+ index_config_without_uri: uri_from_path(&index_config_without_uri_path),
+ index_config_with_retention: uri_from_path(&index_config_with_retention_path),
+ log_docs: uri_from_path(&log_docs_path),
+ wikipedia_docs: uri_from_path(&wikipedia_docs_path),
};
Ok(TestEnv {
@@ -297,15 +296,14 @@ pub async fn upload_test_file(
bucket: &str,
prefix: &str,
filename: &str,
-) -> PathBuf {
+) -> Uri {
let test_data = tokio::fs::read(local_src_path).await.unwrap();
- let mut src_location: PathBuf = [r"s3://", bucket, prefix].iter().collect();
- let storage_uri = Uri::from_str(src_location.to_string_lossy().borrow()).unwrap();
+ let src_location = format!("s3://{}/{}", bucket, prefix);
+ let storage_uri = Uri::from_str(&src_location).unwrap();
let storage = storage_resolver.resolve(&storage_uri).await.unwrap();
storage
.put(&PathBuf::from(filename), Box::new(test_data))
.await
.unwrap();
- src_location.push(filename);
- src_location
+ storage_uri.join(filename).unwrap()
}
diff --git a/quickwit/quickwit-common/src/fs.rs b/quickwit/quickwit-common/src/fs.rs
index adcb432e1b1..1aaa43d8286 100644
--- a/quickwit/quickwit-common/src/fs.rs
+++ b/quickwit/quickwit-common/src/fs.rs
@@ -34,7 +34,7 @@ pub async fn empty_dir>(path: P) -> anyhow::Result<()> {
Ok(())
}
-/// Helper function to get the cache path.
+/// Helper function to get the indexer split cache path.
pub fn get_cache_directory_path(data_dir_path: &Path) -> PathBuf {
data_dir_path.join("indexer-split-cache").join("splits")
}
diff --git a/quickwit/quickwit-config/src/lib.rs b/quickwit/quickwit-config/src/lib.rs
index e41c46dce5b..5e256793fcd 100644
--- a/quickwit/quickwit-config/src/lib.rs
+++ b/quickwit/quickwit-config/src/lib.rs
@@ -55,11 +55,13 @@ pub use quickwit_doc_mapper::DocMapping;
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::Value as JsonValue;
+use source_config::FileSourceParamsForSerde;
pub use source_config::{
- load_source_config_from_user_config, FileSourceParams, KafkaSourceParams, KinesisSourceParams,
- PubSubSourceParams, PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig,
- SourceInputFormat, SourceParams, TransformConfig, VecSourceParams, VoidSourceParams,
- CLI_SOURCE_ID, INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID,
+ load_source_config_from_user_config, FileSourceMessageType, FileSourceNotification,
+ FileSourceParams, FileSourceSqs, KafkaSourceParams, KinesisSourceParams, PubSubSourceParams,
+ PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig, SourceInputFormat,
+ SourceParams, TransformConfig, VecSourceParams, VoidSourceParams, CLI_SOURCE_ID,
+ INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID,
};
use tracing::warn;
@@ -112,7 +114,10 @@ pub fn disable_ingest_v1() -> bool {
IndexTemplateV0_8,
SourceInputFormat,
SourceParams,
- FileSourceParams,
+ FileSourceMessageType,
+ FileSourceNotification,
+ FileSourceParamsForSerde,
+ FileSourceSqs,
PubSubSourceParams,
KafkaSourceParams,
KinesisSourceParams,
diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs
index bc1c0cf3168..b9fcaa15018 100644
--- a/quickwit/quickwit-config/src/source_config/mod.rs
+++ b/quickwit/quickwit-config/src/source_config/mod.rs
@@ -19,8 +19,8 @@
pub(crate) mod serialize;
+use std::borrow::Cow;
use std::num::NonZeroUsize;
-use std::path::{Path, PathBuf};
use std::str::FromStr;
use bytes::Bytes;
@@ -82,6 +82,7 @@ impl SourceConfig {
SourceParams::Kinesis(_) => SourceType::Kinesis,
SourceParams::PubSub(_) => SourceType::PubSub,
SourceParams::Pulsar(_) => SourceType::Pulsar,
+ SourceParams::Stdin => SourceType::Stdin,
SourceParams::Vec(_) => SourceType::Vec,
SourceParams::Void(_) => SourceType::Void,
}
@@ -98,6 +99,7 @@ impl SourceConfig {
SourceParams::Kafka(params) => serde_json::to_value(params),
SourceParams::Kinesis(params) => serde_json::to_value(params),
SourceParams::Pulsar(params) => serde_json::to_value(params),
+ SourceParams::Stdin => serde_json::to_value(()),
SourceParams::Vec(params) => serde_json::to_value(params),
SourceParams::Void(params) => serde_json::to_value(params),
}
@@ -214,6 +216,7 @@ impl FromStr for SourceInputFormat {
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)]
#[serde(tag = "source_type", content = "params", rename_all = "snake_case")]
pub enum SourceParams {
+ #[schema(value_type = FileSourceParamsForSerde)]
File(FileSourceParams),
Ingest,
#[serde(rename = "ingest-api")]
@@ -225,17 +228,22 @@ pub enum SourceParams {
#[serde(rename = "pubsub")]
PubSub(PubSubSourceParams),
Pulsar(PulsarSourceParams),
+ Stdin,
Vec(VecSourceParams),
Void(VoidSourceParams),
}
impl SourceParams {
- pub fn file>(filepath: P) -> Self {
- Self::File(FileSourceParams::file(filepath))
+ pub fn file_from_uri(uri: Uri) -> Self {
+ Self::File(FileSourceParams::Filepath(uri))
+ }
+
+ pub fn file_from_str>(filepath: P) -> anyhow::Result {
+ Uri::from_str(filepath.as_ref()).map(Self::file_from_uri)
}
pub fn stdin() -> Self {
- Self::File(FileSourceParams::stdin())
+ Self::Stdin
}
pub fn void() -> Self {
@@ -243,41 +251,92 @@ impl SourceParams {
}
}
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum FileSourceMessageType {
+ /// See
+ S3Notification,
+ /// A string with the URI of the file (e.g `s3://bucket/key`)
+ RawUri,
+}
+
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)]
+pub struct FileSourceSqs {
+ pub queue_url: String,
+ pub message_type: FileSourceMessageType,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum FileSourceNotification {
+ Sqs(FileSourceSqs),
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, utoipa::ToSchema)]
#[serde(deny_unknown_fields)]
-pub struct FileSourceParams {
- /// Path of the file to read. Assume stdin if None.
- #[schema(value_type = String)]
- #[serde(skip_serializing_if = "Option::is_none")]
- #[serde(default)]
- #[serde(deserialize_with = "absolute_filepath_from_str")]
- pub filepath: Option, //< If None read from stdin.
+pub(super) struct FileSourceParamsForSerde {
+ #[serde(default, skip_serializing_if = "Vec::is_empty")]
+ notifications: Vec,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ filepath: Option,
}
-/// Deserializing as an URI first to validate the input.
-///
-/// TODO: we might want to replace `PathBuf` with `Uri` directly in
-/// `FileSourceParams`
-fn absolute_filepath_from_str<'de, D>(deserializer: D) -> Result, D::Error>
-where D: Deserializer<'de> {
- let filepath_opt: Option = Deserialize::deserialize(deserializer)?;
- if let Some(filepath) = filepath_opt {
- let uri = Uri::from_str(&filepath).map_err(D::Error::custom)?;
- Ok(Some(PathBuf::from(uri.as_str())))
- } else {
- Ok(None)
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(
+ try_from = "FileSourceParamsForSerde",
+ into = "FileSourceParamsForSerde"
+)]
+pub enum FileSourceParams {
+ Notifications(FileSourceNotification),
+ Filepath(Uri),
+}
+
+impl TryFrom for FileSourceParams {
+ type Error = Cow<'static, str>;
+
+ fn try_from(mut value: FileSourceParamsForSerde) -> Result {
+ if value.filepath.is_some() && !value.notifications.is_empty() {
+ return Err(
+ "File source parameters `notifications` and `filepath` are mutually exclusive"
+ .into(),
+ );
+ }
+ if let Some(filepath) = value.filepath {
+ let uri = Uri::from_str(&filepath).map_err(|err| err.to_string())?;
+ Ok(FileSourceParams::Filepath(uri))
+ } else if value.notifications.len() == 1 {
+ Ok(FileSourceParams::Notifications(
+ value.notifications.remove(0),
+ ))
+ } else if value.notifications.len() > 1 {
+ return Err("Only one notification can be specified for now".into());
+ } else {
+ return Err(
+ "Either `notifications` or `filepath` must be specified as file source parameters"
+ .into(),
+ );
+ }
}
}
-impl FileSourceParams {
- pub fn file>(filepath: P) -> Self {
- FileSourceParams {
- filepath: Some(filepath.as_ref().to_path_buf()),
+impl From for FileSourceParamsForSerde {
+ fn from(value: FileSourceParams) -> Self {
+ match value {
+ FileSourceParams::Filepath(uri) => Self {
+ filepath: Some(uri.to_string()),
+ notifications: vec![],
+ },
+ FileSourceParams::Notifications(notification) => Self {
+ filepath: None,
+ notifications: vec![notification],
+ },
}
}
+}
- pub fn stdin() -> Self {
- FileSourceParams { filepath: None }
+impl FileSourceParams {
+ pub fn from_filepath>(filepath: P) -> anyhow::Result {
+ Uri::from_str(filepath.as_ref()).map(Self::Filepath)
}
}
@@ -802,18 +861,88 @@ mod tests {
}
#[test]
- fn test_file_source_params_serialization() {
+ fn test_file_source_params_serde() {
{
let yaml = r#"
filepath: source-path.json
"#;
- let file_params = serde_yaml::from_str::(yaml).unwrap();
+ let file_params_deserialized = serde_yaml::from_str::(yaml).unwrap();
let uri = Uri::from_str("source-path.json").unwrap();
+ assert_eq!(file_params_deserialized, FileSourceParams::Filepath(uri));
+ let file_params_reserialized = serde_json::to_value(file_params_deserialized).unwrap();
+ file_params_reserialized
+ .get("filepath")
+ .unwrap()
+ .as_str()
+ .unwrap()
+ .contains("source-path.json");
+ }
+ {
+ let yaml = r#"
+ notifications:
+ - type: sqs
+ queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name
+ message_type: s3_notification
+ "#;
+ let file_params_deserialized = serde_yaml::from_str::(yaml).unwrap();
+ assert_eq!(
+ file_params_deserialized,
+ FileSourceParams::Notifications(FileSourceNotification::Sqs(FileSourceSqs {
+ queue_url: "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name"
+ .to_string(),
+ message_type: FileSourceMessageType::S3Notification,
+ })),
+ );
+ let file_params_reserialized = serde_json::to_value(&file_params_deserialized).unwrap();
assert_eq!(
- file_params.filepath.unwrap().as_path(),
- Path::new(uri.as_str())
+ file_params_reserialized,
+ json!({"notifications": [{"type": "sqs", "queue_url": "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name", "message_type": "s3_notification"}]})
);
}
+ {
+ let yaml = r#"
+ filepath: source-path.json
+ notifications:
+ - type: sqs
+ queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name
+ message_type: s3_notification
+ "#;
+ let error = serde_yaml::from_str::(yaml).unwrap_err();
+ assert_eq!(
+ error.to_string(),
+ "File source parameters `notifications` and `filepath` are mutually exclusive"
+ );
+ }
+ {
+ let yaml = r#"
+ notifications:
+ - type: sqs
+ queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue1
+ message_type: s3_notification
+ - type: sqs
+ queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue2
+ message_type: s3_notification
+ "#;
+ let error = serde_yaml::from_str::(yaml).unwrap_err();
+ assert_eq!(
+ error.to_string(),
+ "Only one notification can be specified for now"
+ );
+ }
+ {
+ let json = r#"
+ {
+ "notifications": [
+ {
+ "queue_url": "https://sqs.us-east-1.amazonaws.com/123456789012/queue",
+ "message_type": "s3_notification"
+ }
+ ]
+ }
+ "#;
+ let error = serde_json::from_str::(json).unwrap_err();
+ assert!(error.to_string().contains("missing field `type`"));
+ }
}
#[test]
@@ -1199,7 +1328,9 @@ mod tests {
"desired_num_pipelines": 1,
"max_num_pipelines_per_indexer": 1,
"source_type": "file",
- "params": {"filepath": "/test_non_json_corpus.txt"},
+ "params": {
+ "filepath": "s3://mybucket/test_non_json_corpus.txt"
+ },
"input_format": "plain_text"
}"#;
let source_config =
diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs
index 00b76923775..a99f9371eed 100644
--- a/quickwit/quickwit-config/src/source_config/serialize.rs
+++ b/quickwit/quickwit-config/src/source_config/serialize.rs
@@ -24,7 +24,10 @@ use quickwit_proto::types::SourceId;
use serde::{Deserialize, Serialize};
use super::{TransformConfig, RESERVED_SOURCE_IDS};
-use crate::{validate_identifier, ConfigFormat, SourceConfig, SourceInputFormat, SourceParams};
+use crate::{
+ validate_identifier, ConfigFormat, FileSourceParams, SourceConfig, SourceInputFormat,
+ SourceParams,
+};
type SourceConfigForSerialization = SourceConfigV0_8;
@@ -65,12 +68,11 @@ impl SourceConfigForSerialization {
/// Checks the validity of the `SourceConfig` as a "deserializable source".
///
/// Two remarks:
- /// - This does not check connectivity. (See `check_connectivity(..)`)
- /// This just validate configuration, without performing any IO.
- /// - This is only here to validate user input.
- /// When ingesting from stdin, we programmatically create an invalid `SourceConfig`.
- ///
- /// TODO refactor #1065
+ /// - This does not check connectivity, it just validate configuration,
+ /// without performing any IO. See `check_connectivity(..)`.
+ /// - This is used each time the `SourceConfig` is deserialized (at creation but also during
+ /// communications with the metastore). When ingesting from stdin, we programmatically create
+ /// an invalid `SourceConfig` and only use it locally.
fn validate_and_build(self) -> anyhow::Result {
if !RESERVED_SOURCE_IDS.contains(&self.source_id.as_str()) {
validate_identifier("source", &self.source_id)?;
@@ -78,16 +80,16 @@ impl SourceConfigForSerialization {
let num_pipelines = NonZeroUsize::new(self.num_pipelines)
.ok_or_else(|| anyhow::anyhow!("`desired_num_pipelines` must be strictly positive"))?;
match &self.source_params {
- // We want to forbid source_config with no filepath
- SourceParams::File(file_params) => {
- if file_params.filepath.is_none() {
- bail!(
- "source `{}` of type `file` must contain a filepath",
- self.source_id
- )
- }
+ SourceParams::Stdin => {
+ bail!(
+ "stdin can only be used as source through the CLI command `quickwit tool \
+ local-ingest`"
+ );
}
- SourceParams::Kafka(_) | SourceParams::Kinesis(_) | SourceParams::Pulsar(_) => {
+ SourceParams::File(_)
+ | SourceParams::Kafka(_)
+ | SourceParams::Kinesis(_)
+ | SourceParams::Pulsar(_) => {
// TODO consider any validation opportunity
}
SourceParams::PubSub(_)
@@ -98,7 +100,9 @@ impl SourceConfigForSerialization {
| SourceParams::Void(_) => {}
}
match &self.source_params {
- SourceParams::PubSub(_) | SourceParams::Kafka(_) => {}
+ SourceParams::PubSub(_)
+ | SourceParams::Kafka(_)
+ | SourceParams::File(FileSourceParams::Notifications(_)) => {}
_ => {
if self.num_pipelines > 1 {
bail!("Quickwit currently supports multiple pipelines only for GCP PubSub or Kafka sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types");
diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs
index ff1bd5658ad..c396e1ac23a 100644
--- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs
+++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs
@@ -30,11 +30,11 @@ use fnv::{FnvHashMap, FnvHashSet};
use itertools::Itertools;
use once_cell::sync::OnceCell;
use quickwit_common::pretty::PrettySample;
+use quickwit_config::{FileSourceParams, SourceParams};
use quickwit_proto::indexing::{
ApplyIndexingPlanRequest, CpuCapacity, IndexingService, IndexingTask, PIPELINE_FULL_CAPACITY,
PIPELINE_THROUGHPUT,
};
-use quickwit_proto::metastore::SourceType;
use quickwit_proto::types::NodeId;
use scheduling::{SourceToSchedule, SourceToScheduleType};
use serde::Serialize;
@@ -168,22 +168,22 @@ fn get_sources_to_schedule(model: &ControlPlaneModel) -> Vec {
if !source_config.enabled {
continue;
}
- match source_config.source_type() {
- SourceType::Cli
- | SourceType::File
- | SourceType::Vec
- | SourceType::Void
- | SourceType::Unspecified => {
- // We don't need to schedule those.
+ match source_config.source_params {
+ SourceParams::File(FileSourceParams::Filepath(_))
+ | SourceParams::IngestCli
+ | SourceParams::Stdin
+ | SourceParams::Void(_)
+ | SourceParams::Vec(_) => { // We don't need to schedule those.
}
- SourceType::IngestV1 => {
+
+ SourceParams::IngestApi => {
// TODO ingest v1 is scheduled differently
sources.push(SourceToSchedule {
source_uid,
source_type: SourceToScheduleType::IngestV1,
});
}
- SourceType::IngestV2 => {
+ SourceParams::Ingest => {
// Expect: the source should exist since we just read it from `get_source_configs`.
// Note that we keep all shards, including Closed shards:
// A closed shards still needs to be indexed.
@@ -208,11 +208,11 @@ fn get_sources_to_schedule(model: &ControlPlaneModel) -> Vec {
},
});
}
- SourceType::Kafka
- | SourceType::Kinesis
- | SourceType::PubSub
- | SourceType::Nats
- | SourceType::Pulsar => {
+ SourceParams::Kafka(_)
+ | SourceParams::Kinesis(_)
+ | SourceParams::PubSub(_)
+ | SourceParams::Pulsar(_)
+ | SourceParams::File(FileSourceParams::Notifications(_)) => {
sources.push(SourceToSchedule {
source_uid,
source_type: SourceToScheduleType::NonSharded {
diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs
index 29111f28b23..8636cca7379 100644
--- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs
+++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs
@@ -820,6 +820,8 @@ impl IngestController {
leader_id: shard.leader_id.clone(),
follower_id: shard.follower_id.clone(),
doc_mapping_uid: shard.doc_mapping_uid,
+ // Shards are acquired by the ingest sources
+ publish_token: None,
}
})
.collect();
diff --git a/quickwit/quickwit-doc-mapper/src/query_builder.rs b/quickwit/quickwit-doc-mapper/src/query_builder.rs
index 3054ac0661f..ca74afa2462 100644
--- a/quickwit/quickwit-doc-mapper/src/query_builder.rs
+++ b/quickwit/quickwit-doc-mapper/src/query_builder.rs
@@ -244,7 +244,11 @@ impl<'a, 'b: 'a> QueryAstVisitor<'a> for ExtractPrefixTermRanges<'b> {
&mut self,
phrase_prefix: &'a PhrasePrefixQuery,
) -> Result<(), Self::Err> {
- let (_, terms) = phrase_prefix.get_terms(self.schema, self.tokenizer_manager)?;
+ let terms = match phrase_prefix.get_terms(self.schema, self.tokenizer_manager) {
+ Ok((_, terms)) => terms,
+ Err(InvalidQuery::SchemaError(_)) => return Ok(()), /* the query will be nullified when casting to a tantivy ast */
+ Err(e) => return Err(e),
+ };
if let Some((_, term)) = terms.last() {
self.add_prefix_term(term.clone(), phrase_prefix.max_expansions, terms.len() > 1);
}
@@ -270,12 +274,9 @@ fn extract_prefix_term_ranges(
#[cfg(test)]
mod test {
- use quickwit_datetime::{parse_date_time_str, DateTimeInputFormat};
use quickwit_query::create_default_quickwit_tokenizer_manager;
use quickwit_query::query_ast::query_ast_from_user_text;
- use tantivy::columnar::MonotonicallyMappableToU64;
use tantivy::schema::{DateOptions, DateTimePrecision, Schema, FAST, INDEXED, STORED, TEXT};
- use tantivy::DateTime;
use super::build_query;
use crate::{DYNAMIC_FIELD_NAME, SOURCE_FIELD_NAME};
@@ -497,54 +498,46 @@ mod test {
#[test]
fn test_datetime_range_query() {
- let input_formats = [DateTimeInputFormat::Rfc3339];
{
// Check range on datetime in millisecond, precision has no impact as it is in
// milliseconds.
let start_date_time_str = "2023-01-10T08:38:51.150Z";
- let start_date_time = parse_date_time_str(start_date_time_str, &input_formats).unwrap();
- let start_date_time_u64 = start_date_time.to_u64();
let end_date_time_str = "2023-01-10T08:38:51.160Z";
- let end_date_time: DateTime =
- parse_date_time_str(end_date_time_str, &input_formats).unwrap();
- let end_date_time_u64 = end_date_time.to_u64();
- let expectation_with_lower_and_upper_bounds = format!(
- r#"FastFieldRangeWeight {{ field: "dt", lower_bound: Included({start_date_time_u64}), upper_bound: Included({end_date_time_u64}), column_type_opt: Some(DateTime) }}"#,
+ check_build_query_static_mode(
+ &format!("dt:[{start_date_time_str} TO {end_date_time_str}]"),
+ Vec::new(),
+ TestExpectation::Ok("2023-01-10T08:38:51.15Z"),
);
check_build_query_static_mode(
&format!("dt:[{start_date_time_str} TO {end_date_time_str}]"),
Vec::new(),
- TestExpectation::Ok(&expectation_with_lower_and_upper_bounds),
+ TestExpectation::Ok("RangeQuery"),
+ );
+ check_build_query_static_mode(
+ &format!("dt:<{end_date_time_str}"),
+ Vec::new(),
+ TestExpectation::Ok("lower_bound: Unbounded"),
);
- let expectation_with_upper_bound = format!(
- r#"FastFieldRangeWeight {{ field: "dt", lower_bound: Unbounded, upper_bound: Excluded({end_date_time_u64}), column_type_opt: Some(DateTime) }}"#,
+ check_build_query_static_mode(
+ &format!("dt:<{end_date_time_str}"),
+ Vec::new(),
+ TestExpectation::Ok("upper_bound: Excluded"),
);
check_build_query_static_mode(
&format!("dt:<{end_date_time_str}"),
Vec::new(),
- TestExpectation::Ok(&expectation_with_upper_bound),
+ TestExpectation::Ok("2023-01-10T08:38:51.16Z"),
);
}
// Check range on datetime in microseconds and truncation to milliseconds.
{
let start_date_time_str = "2023-01-10T08:38:51.000150Z";
- let start_date_time = parse_date_time_str(start_date_time_str, &input_formats)
- .unwrap()
- .truncate(DateTimePrecision::Milliseconds);
- let start_date_time_u64 = start_date_time.to_u64();
let end_date_time_str = "2023-01-10T08:38:51.000151Z";
- let end_date_time: DateTime = parse_date_time_str(end_date_time_str, &input_formats)
- .unwrap()
- .truncate(DateTimePrecision::Milliseconds);
- let end_date_time_u64 = end_date_time.to_u64();
- let expectation_with_lower_and_upper_bounds = format!(
- r#"FastFieldRangeWeight {{ field: "dt", lower_bound: Included({start_date_time_u64}), upper_bound: Included({end_date_time_u64}), column_type_opt: Some(DateTime) }}"#,
- );
check_build_query_static_mode(
&format!("dt:[{start_date_time_str} TO {end_date_time_str}]"),
Vec::new(),
- TestExpectation::Ok(&expectation_with_lower_and_upper_bounds),
+ TestExpectation::Ok("2023-01-10T08:38:51Z"),
);
}
}
@@ -555,17 +548,17 @@ mod test {
"ip:[127.0.0.1 TO 127.1.1.1]",
Vec::new(),
TestExpectation::Ok(
- "RangeQuery { field: \"ip\", value_type: IpAddr, lower_bound: Included([0, 0, 0, \
- 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 0, 0, 1]), upper_bound: Included([0, 0, 0, \
- 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 1, 1, 1])",
+ "RangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=6, \
+ type=IpAddr, ::ffff:127.0.0.1)), upper_bound: Included(Term(field=6, \
+ type=IpAddr, ::ffff:127.1.1.1)) } }",
),
);
check_build_query_static_mode(
"ip:>127.0.0.1",
Vec::new(),
TestExpectation::Ok(
- "RangeQuery { field: \"ip\", value_type: IpAddr, lower_bound: Excluded([0, 0, 0, \
- 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 0, 0, 1]), upper_bound: Unbounded",
+ "RangeQuery { bounds: BoundsRange { lower_bound: Excluded(Term(field=6, \
+ type=IpAddr, ::ffff:127.0.0.1)), upper_bound: Unbounded } }",
),
);
}
@@ -576,14 +569,14 @@ mod test {
"f64_fast:[7.7 TO 77.7]",
Vec::new(),
TestExpectation::Ok(
- r#"FastFieldRangeWeight { field: "f64_fast", lower_bound: Included(13843727484564851917), upper_bound: Included(13858540105214250189), column_type_opt: Some(F64) }"#,
+ r#"RangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=12, type=F64, 7.7)), upper_bound: Included(Term(field=12, type=F64, 77.7)) } }"#,
),
);
check_build_query_static_mode(
"f64_fast:>7",
Vec::new(),
TestExpectation::Ok(
- r#"FastFieldRangeWeight { field: "f64_fast", lower_bound: Excluded(13842939354630062080), upper_bound: Unbounded, column_type_opt: Some(F64)"#,
+ r#"RangeQuery { bounds: BoundsRange { lower_bound: Excluded(Term(field=12, type=F64, 7.0)), upper_bound: Unbounded } }"#,
),
);
}
@@ -593,12 +586,12 @@ mod test {
check_build_query_static_mode(
"i64_fast:[-7 TO 77]",
Vec::new(),
- TestExpectation::Ok(r#"FastFieldRangeWeight { field: "i64_fast","#),
+ TestExpectation::Ok(r#"field=11"#),
);
check_build_query_static_mode(
"i64_fast:>7",
Vec::new(),
- TestExpectation::Ok(r#"FastFieldRangeWeight { field: "i64_fast","#),
+ TestExpectation::Ok(r#"field=11"#),
);
}
@@ -607,12 +600,12 @@ mod test {
check_build_query_static_mode(
"u64_fast:[7 TO 77]",
Vec::new(),
- TestExpectation::Ok(r#"FastFieldRangeWeight { field: "u64_fast","#),
+ TestExpectation::Ok(r#"field=10,"#),
);
check_build_query_static_mode(
"u64_fast:>7",
Vec::new(),
- TestExpectation::Ok(r#"FastFieldRangeWeight { field: "u64_fast","#),
+ TestExpectation::Ok(r#"field=10,"#),
);
}
@@ -622,9 +615,9 @@ mod test {
"ips:[127.0.0.1 TO 127.1.1.1]",
Vec::new(),
TestExpectation::Ok(
- "RangeQuery { field: \"ips\", value_type: IpAddr, lower_bound: Included([0, 0, 0, \
- 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 0, 0, 1]), upper_bound: Included([0, 0, 0, \
- 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 1, 1, 1])",
+ "RangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=7, \
+ type=IpAddr, ::ffff:127.0.0.1)), upper_bound: Included(Term(field=7, \
+ type=IpAddr, ::ffff:127.1.1.1)) } }",
),
);
}
diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml
index d7008a4579c..7b005301ea9 100644
--- a/quickwit/quickwit-indexing/Cargo.toml
+++ b/quickwit/quickwit-indexing/Cargo.toml
@@ -11,12 +11,12 @@ authors.workspace = true
license.workspace = true
[dependencies]
-aws-sdk-kinesis = { workspace = true, optional = true }
-
anyhow = { workspace = true }
arc-swap = { workspace = true }
async-compression = { workspace = true }
async-trait = { workspace = true }
+aws-sdk-kinesis = { workspace = true, optional = true }
+aws-sdk-sqs = { workspace = true, optional = true }
bytes = { workspace = true }
bytesize = { workspace = true }
fail = { workspace = true }
@@ -34,6 +34,7 @@ oneshot = { workspace = true }
openssl = { workspace = true, optional = true }
pulsar = { workspace = true, optional = true }
quickwit-query = { workspace = true }
+regex = { workspace = true }
rdkafka = { workspace = true, optional = true }
serde = { workspace = true }
serde_json = { workspace = true }
@@ -77,6 +78,13 @@ kinesis = [
kinesis-localstack-tests = []
pulsar = ["dep:pulsar"]
pulsar-broker-tests = []
+queue-sources = []
+sqs = [
+ "aws-sdk-sqs",
+ "queue-sources",
+ "quickwit-aws/sqs",
+]
+sqs-localstack-tests = []
vendored-kafka = [
"kafka",
"libz-sys/static",
diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs
index 0c54c44fd5d..4087f2ed230 100644
--- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs
+++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs
@@ -436,6 +436,7 @@ impl IndexingPipeline {
queues_dir_path: self.params.queues_dir_path.clone(),
storage_resolver: self.params.source_storage_resolver.clone(),
event_broker: self.params.event_broker.clone(),
+ indexing_setting: self.params.indexing_settings.clone(),
};
let source = ctx
.protect_future(quickwit_supported_sources().load_source(source_runtime))
@@ -598,7 +599,7 @@ mod tests {
use quickwit_actors::{Command, Universe};
use quickwit_common::ServiceStream;
- use quickwit_config::{IndexingSettings, SourceInputFormat, SourceParams, VoidSourceParams};
+ use quickwit_config::{IndexingSettings, SourceInputFormat, SourceParams};
use quickwit_doc_mapper::{default_doc_mapper_for_test, DefaultDocMapper};
use quickwit_metastore::checkpoint::IndexCheckpointDelta;
use quickwit_metastore::{IndexMetadata, IndexMetadataResponseExt, PublishSplitsRequestExt};
@@ -639,7 +640,7 @@ mod tests {
source_id: "test-source".to_string(),
num_pipelines: NonZeroUsize::MIN,
enabled: true,
- source_params: SourceParams::file(PathBuf::from(test_file)),
+ source_params: SourceParams::file_from_str(test_file).unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -689,7 +690,7 @@ mod tests {
&& publish_splits_request.staged_split_ids.len() == 1
&& publish_splits_request.replaced_split_ids.is_empty()
&& format!("{:?}", checkpoint_delta.source_delta)
- .ends_with(":(00000000000000000000..00000000000000001030])")
+ .ends_with(":(00000000000000000000..~00000000000000001030])")
})
.returning(|_| Ok(EmptyResponse {}));
@@ -758,7 +759,7 @@ mod tests {
source_id: "test-source".to_string(),
num_pipelines: NonZeroUsize::MIN,
enabled: true,
- source_params: SourceParams::file(PathBuf::from(test_file)),
+ source_params: SourceParams::file_from_str(test_file).unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -801,7 +802,7 @@ mod tests {
&& publish_splits_request.replaced_split_ids.is_empty()
&& checkpoint_delta.source_id == "test-source"
&& format!("{:?}", checkpoint_delta.source_delta)
- .ends_with(":(00000000000000000000..00000000000000001030])")
+ .ends_with(":(00000000000000000000..~00000000000000001030])")
})
.returning(|_| Ok(EmptyResponse {}));
@@ -862,7 +863,7 @@ mod tests {
source_id: "test-source".to_string(),
num_pipelines: NonZeroUsize::MIN,
enabled: true,
- source_params: SourceParams::Void(VoidSourceParams),
+ source_params: SourceParams::void(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -965,7 +966,7 @@ mod tests {
source_id: "test-source".to_string(),
num_pipelines: NonZeroUsize::MIN,
enabled: true,
- source_params: SourceParams::file(PathBuf::from(test_file)),
+ source_params: SourceParams::file_from_str(test_file).unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -1008,7 +1009,7 @@ mod tests {
&& publish_splits_request.replaced_split_ids.is_empty()
&& checkpoint_delta.source_id == "test-source"
&& format!("{:?}", checkpoint_delta.source_delta)
- .ends_with(":(00000000000000000000..00000000000000001030])")
+ .ends_with(":(00000000000000000000..~00000000000000001030])")
})
.returning(|_| Ok(EmptyResponse {}));
let universe = Universe::new();
@@ -1056,6 +1057,7 @@ mod tests {
let (_pipeline_mailbox, pipeline_handler) = universe.spawn_builder().spawn(pipeline);
let (pipeline_exit_status, pipeline_statistics) = pipeline_handler.join().await;
assert!(pipeline_exit_status.is_success());
+ // flaky. Sometimes generations is 2.
assert_eq!(pipeline_statistics.generation, 1);
assert_eq!(pipeline_statistics.num_spawn_attempts, 1);
assert_eq!(pipeline_statistics.num_published_splits, 0);
diff --git a/quickwit/quickwit-indexing/src/actors/publisher.rs b/quickwit/quickwit-indexing/src/actors/publisher.rs
index 1f24ef51477..4e999e7f7a2 100644
--- a/quickwit/quickwit-indexing/src/actors/publisher.rs
+++ b/quickwit/quickwit-indexing/src/actors/publisher.rs
@@ -147,7 +147,7 @@ impl Handler for Publisher {
);
return Ok(());
}
- info!(new_splits=?split_ids, checkpoint_delta=?checkpoint_delta_opt, "publish-new-splits");
+ info!("publish-new-splits");
if let Some(source_mailbox) = self.source_mailbox_opt.as_ref() {
if let Some(checkpoint) = checkpoint_delta_opt {
// We voluntarily do not log anything here.
diff --git a/quickwit/quickwit-indexing/src/source/doc_file_reader.rs b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs
new file mode 100644
index 00000000000..f8fc3c0b8d3
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs
@@ -0,0 +1,561 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::io;
+use std::path::Path;
+
+use anyhow::Context;
+use async_compression::tokio::bufread::GzipDecoder;
+use bytes::Bytes;
+use quickwit_common::uri::Uri;
+use quickwit_common::Progress;
+use quickwit_metastore::checkpoint::PartitionId;
+use quickwit_proto::metastore::SourceType;
+use quickwit_proto::types::Position;
+use quickwit_storage::StorageResolver;
+use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader};
+
+use super::{BatchBuilder, BATCH_NUM_BYTES_LIMIT};
+
+pub struct FileRecord {
+ pub next_offset: u64,
+ pub doc: Bytes,
+ pub is_last: bool,
+}
+
+/// A helper wrapper that lets you skip bytes in compressed files where you
+/// cannot seek (e.g. gzip files).
+struct SkipReader {
+ reader: BufReader>,
+ num_bytes_to_skip: usize,
+}
+
+impl SkipReader {
+ fn new(reader: Box, num_bytes_to_skip: usize) -> Self {
+ Self {
+ reader: BufReader::new(reader),
+ num_bytes_to_skip,
+ }
+ }
+
+ async fn skip(&mut self) -> io::Result<()> {
+ // allocate on the heap to avoid stack overflows
+ let mut buf = vec![0u8; 64_000];
+ while self.num_bytes_to_skip > 0 {
+ let num_bytes_to_read = self.num_bytes_to_skip.min(buf.len());
+ let num_bytes_read = self
+ .reader
+ .read_exact(&mut buf[..num_bytes_to_read])
+ .await?;
+ self.num_bytes_to_skip -= num_bytes_read;
+ }
+ Ok(())
+ }
+
+ /// Reads a line and peeks into the readers buffer. Returns the number of
+ /// bytes read and true the end of the file is reached.
+ async fn read_line_and_peek<'a>(&mut self, buf: &'a mut String) -> io::Result<(usize, bool)> {
+ if self.num_bytes_to_skip > 0 {
+ self.skip().await?;
+ }
+ let line_size = self.reader.read_line(buf).await?;
+ if line_size == 0 {
+ return Ok((0, true));
+ }
+ let next_bytes = self.reader.fill_buf().await?;
+ Ok((line_size, next_bytes.is_empty()))
+ }
+}
+
+pub struct DocFileReader {
+ reader: SkipReader,
+ next_offset: u64,
+}
+
+impl DocFileReader {
+ pub fn empty() -> Self {
+ DocFileReader {
+ reader: SkipReader::new(Box::new(tokio::io::empty()), 0),
+ next_offset: 0,
+ }
+ }
+
+ pub async fn from_uri(
+ storage_resolver: &StorageResolver,
+ uri: &Uri,
+ offset: usize,
+ ) -> anyhow::Result {
+ let (dir_uri, file_name) = dir_and_filename(uri)?;
+ let storage = storage_resolver.resolve(&dir_uri).await?;
+ let file_size = storage.file_num_bytes(file_name).await?.try_into().unwrap();
+ if file_size == 0 {
+ return Ok(DocFileReader::empty());
+ }
+ // If it's a gzip file, we can't seek to a specific offset. `SkipReader`
+ // starts from the beginning of the file, decompresses and skips the
+ // first `offset` bytes.
+ let reader = if uri.extension() == Some("gz") {
+ let stream = storage.get_slice_stream(file_name, 0..file_size).await?;
+ let decompressed_stream = Box::new(GzipDecoder::new(BufReader::new(stream)));
+ DocFileReader {
+ reader: SkipReader::new(decompressed_stream, offset),
+ next_offset: offset as u64,
+ }
+ } else {
+ let stream = storage
+ .get_slice_stream(file_name, offset..file_size)
+ .await?;
+ DocFileReader {
+ reader: SkipReader::new(stream, 0),
+ next_offset: offset as u64,
+ }
+ };
+ Ok(reader)
+ }
+
+ /// Reads the next record from the underlying file. Returns `None` when EOF
+ /// is reached.
+ pub async fn next_record(&mut self) -> anyhow::Result> {
+ let mut buf = String::new();
+ // TODO retry if stream is broken (#5243)
+ let (bytes_read, is_last) = self.reader.read_line_and_peek(&mut buf).await?;
+ if bytes_read == 0 {
+ Ok(None)
+ } else {
+ self.next_offset += bytes_read as u64;
+ Ok(Some(FileRecord {
+ next_offset: self.next_offset,
+ doc: Bytes::from(buf),
+ is_last,
+ }))
+ }
+ }
+}
+
+pub struct ObjectUriBatchReader {
+ partition_id: PartitionId,
+ reader: DocFileReader,
+ current_offset: usize,
+ is_eof: bool,
+}
+
+impl ObjectUriBatchReader {
+ pub async fn try_new(
+ storage_resolver: &StorageResolver,
+ partition_id: PartitionId,
+ uri: &Uri,
+ position: Position,
+ ) -> anyhow::Result {
+ let current_offset = match position {
+ Position::Beginning => 0,
+ Position::Offset(offset) => offset
+ .as_usize()
+ .context("file offset should be stored as usize")?,
+ Position::Eof(_) => {
+ return Ok(ObjectUriBatchReader {
+ partition_id,
+ reader: DocFileReader::empty(),
+ current_offset: 0,
+ is_eof: true,
+ })
+ }
+ };
+ let reader = DocFileReader::from_uri(storage_resolver, uri, current_offset).await?;
+ Ok(ObjectUriBatchReader {
+ partition_id,
+ reader,
+ current_offset,
+ is_eof: false,
+ })
+ }
+
+ pub async fn read_batch(
+ &mut self,
+ source_progress: &Progress,
+ source_type: SourceType,
+ ) -> anyhow::Result {
+ let limit_num_bytes = self.current_offset + BATCH_NUM_BYTES_LIMIT as usize;
+ let mut new_offset = self.current_offset;
+ let mut batch_builder = BatchBuilder::new(source_type);
+ while new_offset < limit_num_bytes {
+ if let Some(record) = source_progress
+ .protect_future(self.reader.next_record())
+ .await?
+ {
+ new_offset = record.next_offset as usize;
+ batch_builder.add_doc(record.doc);
+ if record.is_last {
+ self.is_eof = true;
+ break;
+ }
+ } else {
+ self.is_eof = true;
+ break;
+ }
+ }
+ let to_position = if self.is_eof {
+ Position::eof(new_offset)
+ } else {
+ Position::offset(new_offset)
+ };
+ batch_builder.checkpoint_delta.record_partition_delta(
+ self.partition_id.clone(),
+ Position::offset(self.current_offset),
+ to_position,
+ )?;
+ self.current_offset = new_offset;
+ Ok(batch_builder)
+ }
+
+ pub fn is_eof(&self) -> bool {
+ self.is_eof
+ }
+}
+
+pub(crate) fn dir_and_filename(filepath: &Uri) -> anyhow::Result<(Uri, &Path)> {
+ let dir_uri: Uri = filepath
+ .parent()
+ .context("Parent directory could not be resolved")?;
+ let file_name = filepath
+ .file_name()
+ .context("Path does not appear to be a file")?;
+ Ok((dir_uri, file_name))
+}
+
+#[cfg(test)]
+pub mod file_test_helpers {
+ use std::io::Write;
+
+ use async_compression::tokio::write::GzipEncoder;
+ use tempfile::NamedTempFile;
+
+ pub const DUMMY_DOC: &[u8] = r#"{"body": "hello happy tax payer!"}"#.as_bytes();
+
+ async fn gzip_bytes(bytes: &[u8]) -> Vec {
+ let mut gzip_documents = Vec::new();
+ let mut encoder = GzipEncoder::new(&mut gzip_documents);
+ tokio::io::AsyncWriteExt::write_all(&mut encoder, bytes)
+ .await
+ .unwrap();
+ // flush is not sufficient here and reading the file will raise a unexpected end of file
+ // error.
+ tokio::io::AsyncWriteExt::shutdown(&mut encoder)
+ .await
+ .unwrap();
+ gzip_documents
+ }
+
+ async fn write_to_tmp(data: Vec, gzip: bool) -> NamedTempFile {
+ let mut temp_file: tempfile::NamedTempFile = if gzip {
+ tempfile::Builder::new().suffix(".gz").tempfile().unwrap()
+ } else {
+ tempfile::NamedTempFile::new().unwrap()
+ };
+ if gzip {
+ let gzip_documents = gzip_bytes(&data).await;
+ temp_file.write_all(&gzip_documents).unwrap();
+ } else {
+ temp_file.write_all(&data).unwrap();
+ }
+ temp_file.flush().unwrap();
+ temp_file
+ }
+
+ pub async fn generate_dummy_doc_file(gzip: bool, lines: usize) -> (NamedTempFile, usize) {
+ let mut documents_bytes = Vec::with_capacity(DUMMY_DOC.len() * lines);
+ for _ in 0..lines {
+ documents_bytes.write_all(DUMMY_DOC).unwrap();
+ documents_bytes.write_all("\n".as_bytes()).unwrap();
+ }
+ let size = documents_bytes.len();
+ let file = write_to_tmp(documents_bytes, gzip).await;
+ (file, size)
+ }
+
+ /// Generates a file with increasing padded numbers. Each line is 8 bytes
+ /// including the newline char.
+ ///
+ /// 0000000\n0000001\n0000002\n...
+ pub async fn generate_index_doc_file(gzip: bool, lines: usize) -> NamedTempFile {
+ assert!(lines < 9999999, "each line is 7 digits + newline");
+ let mut documents_bytes = Vec::new();
+ for i in 0..lines {
+ documents_bytes
+ .write_all(format!("{:0>7}\n", i).as_bytes())
+ .unwrap();
+ }
+ write_to_tmp(documents_bytes, gzip).await
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::io::Cursor;
+ use std::str::FromStr;
+
+ use file_test_helpers::generate_index_doc_file;
+ use quickwit_metastore::checkpoint::SourceCheckpointDelta;
+
+ use super::*;
+
+ #[tokio::test]
+ async fn test_skip_reader() {
+ {
+ // Skip 0 bytes.
+ let mut reader = SkipReader::new(Box::new("hello".as_bytes()), 0);
+ let mut buf = String::new();
+ let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap();
+ assert_eq!(buf, "hello");
+ assert!(eof);
+ assert_eq!(bytes_read, 5)
+ }
+ {
+ // Skip 2 bytes.
+ let mut reader = SkipReader::new(Box::new("hello".as_bytes()), 2);
+ let mut buf = String::new();
+ let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap();
+ assert_eq!(buf, "llo");
+ assert!(eof);
+ assert_eq!(bytes_read, 3)
+ }
+ {
+ let input = "hello";
+ let cursor = Cursor::new(input);
+ let mut reader = SkipReader::new(Box::new(cursor), 5);
+ let mut buf = String::new();
+ let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap();
+ assert!(eof);
+ assert_eq!(bytes_read, 0)
+ }
+ {
+ let input = "hello";
+ let cursor = Cursor::new(input);
+ let mut reader = SkipReader::new(Box::new(cursor), 10);
+ let mut buf = String::new();
+ assert!(reader.read_line_and_peek(&mut buf).await.is_err());
+ }
+ {
+ let input = "hello world".repeat(10000);
+ let cursor = Cursor::new(input.clone());
+ let mut reader = SkipReader::new(Box::new(cursor), 64000);
+ let mut buf = String::new();
+ reader.read_line_and_peek(&mut buf).await.unwrap();
+ assert_eq!(buf, input[64000..]);
+ }
+ {
+ let input = "hello world".repeat(10000);
+ let cursor = Cursor::new(input.clone());
+ let mut reader = SkipReader::new(Box::new(cursor), 64001);
+ let mut buf = String::new();
+ reader.read_line_and_peek(&mut buf).await.unwrap();
+ assert_eq!(buf, input[64001..]);
+ }
+ }
+
+ async fn aux_test_full_read_record(file: impl AsRef, expected_lines: usize) {
+ let storage_resolver = StorageResolver::for_test();
+ let uri = Uri::from_str(file.as_ref()).unwrap();
+ let mut doc_reader = DocFileReader::from_uri(&storage_resolver, &uri, 0)
+ .await
+ .unwrap();
+ let mut parsed_lines = 0;
+ while doc_reader.next_record().await.unwrap().is_some() {
+ parsed_lines += 1;
+ }
+ assert_eq!(parsed_lines, expected_lines);
+ }
+
+ #[tokio::test]
+ async fn test_full_read_record() {
+ aux_test_full_read_record("data/test_corpus.json", 4).await;
+ }
+
+ #[tokio::test]
+ async fn test_full_read_record_gz() {
+ aux_test_full_read_record("data/test_corpus.json.gz", 4).await;
+ }
+
+ #[tokio::test]
+ async fn test_empty_file() {
+ let empty_file = tempfile::NamedTempFile::new().unwrap();
+ let empty_file_uri = empty_file.path().to_str().unwrap();
+ aux_test_full_read_record(empty_file_uri, 0).await;
+ }
+
+ async fn aux_test_resumed_read_record(
+ file: impl AsRef,
+ expected_lines: usize,
+ stop_at_line: usize,
+ ) {
+ let storage_resolver = StorageResolver::for_test();
+ let uri = Uri::from_str(file.as_ref()).unwrap();
+ // read the first part of the file
+ let mut first_part_reader = DocFileReader::from_uri(&storage_resolver, &uri, 0)
+ .await
+ .unwrap();
+ let mut resume_offset = 0;
+ let mut parsed_lines = 0;
+ for _ in 0..stop_at_line {
+ let rec = first_part_reader
+ .next_record()
+ .await
+ .unwrap()
+ .expect("EOF happened before stop_at_line");
+ resume_offset = rec.next_offset as usize;
+ assert_eq!(Bytes::from(format!("{:0>7}\n", parsed_lines)), rec.doc);
+ parsed_lines += 1;
+ }
+ // read the second part of the file
+ let mut second_part_reader =
+ DocFileReader::from_uri(&storage_resolver, &uri, resume_offset)
+ .await
+ .unwrap();
+ while let Some(rec) = second_part_reader.next_record().await.unwrap() {
+ assert_eq!(Bytes::from(format!("{:0>7}\n", parsed_lines)), rec.doc);
+ parsed_lines += 1;
+ }
+ assert_eq!(parsed_lines, expected_lines);
+ }
+
+ #[tokio::test]
+ async fn test_resumed_read_record() {
+ let dummy_doc_file = generate_index_doc_file(false, 1000).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 40).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 999).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1000).await;
+ }
+
+ #[tokio::test]
+ async fn test_resumed_read_record_gz() {
+ let dummy_doc_file = generate_index_doc_file(true, 1000).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 40).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 999).await;
+ aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1000).await;
+ }
+
+ async fn aux_test_full_read_batch(
+ file: impl AsRef,
+ expected_lines: usize,
+ expected_batches: usize,
+ file_size: usize,
+ from: Position,
+ ) {
+ let progress = Progress::default();
+ let storage_resolver = StorageResolver::for_test();
+ let uri = Uri::from_str(file.as_ref()).unwrap();
+ let partition = PartitionId::from("test");
+ let mut batch_reader =
+ ObjectUriBatchReader::try_new(&storage_resolver, partition.clone(), &uri, from)
+ .await
+ .unwrap();
+
+ let mut parsed_lines = 0;
+ let mut parsed_batches = 0;
+ let mut checkpoint_delta = SourceCheckpointDelta::default();
+ while !batch_reader.is_eof() {
+ let batch = batch_reader
+ .read_batch(&progress, SourceType::Unspecified)
+ .await
+ .unwrap();
+ parsed_lines += batch.docs.len();
+ parsed_batches += 1;
+ checkpoint_delta.extend(batch.checkpoint_delta).unwrap();
+ }
+ assert_eq!(parsed_lines, expected_lines);
+ assert_eq!(parsed_batches, expected_batches);
+ let position = checkpoint_delta
+ .get_source_checkpoint()
+ .position_for_partition(&partition)
+ .unwrap()
+ .clone();
+ assert_eq!(position, Position::eof(file_size))
+ }
+
+ #[tokio::test]
+ async fn test_read_batch_empty_file() {
+ let empty_file = tempfile::NamedTempFile::new().unwrap();
+ let empty_file_uri = empty_file.path().to_str().unwrap();
+ aux_test_full_read_batch(empty_file_uri, 0, 1, 0, Position::Beginning).await;
+ }
+
+ #[tokio::test]
+ async fn test_full_read_single_batch() {
+ let num_lines = 10;
+ let dummy_doc_file = generate_index_doc_file(false, num_lines).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_full_read_batch(
+ dummy_doc_file_uri,
+ num_lines,
+ 1,
+ num_lines * 8,
+ Position::Beginning,
+ )
+ .await;
+ }
+
+ #[tokio::test]
+ async fn test_full_read_single_batch_max_size() {
+ let num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8;
+ let dummy_doc_file = generate_index_doc_file(false, num_lines).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_full_read_batch(
+ dummy_doc_file_uri,
+ num_lines,
+ 1,
+ num_lines * 8,
+ Position::Beginning,
+ )
+ .await;
+ }
+
+ #[tokio::test]
+ async fn test_full_read_two_batches() {
+ let num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8 + 10;
+ let dummy_doc_file = generate_index_doc_file(false, num_lines).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_full_read_batch(
+ dummy_doc_file_uri,
+ num_lines,
+ 2,
+ num_lines * 8,
+ Position::Beginning,
+ )
+ .await;
+ }
+
+ #[tokio::test]
+ async fn test_resume_read_batches() {
+ let total_num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8 * 3;
+ let resume_after_lines = total_num_lines / 2;
+ let dummy_doc_file = generate_index_doc_file(false, total_num_lines).await;
+ let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap();
+ aux_test_full_read_batch(
+ dummy_doc_file_uri,
+ total_num_lines - resume_after_lines,
+ 2,
+ total_num_lines * 8,
+ Position::offset(resume_after_lines * 8),
+ )
+ .await;
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/file_source.rs b/quickwit/quickwit-indexing/src/source/file_source.rs
index e4674f7a5ce..e3be553ae1b 100644
--- a/quickwit/quickwit-indexing/src/source/file_source.rs
+++ b/quickwit/quickwit-indexing/src/source/file_source.rs
@@ -17,45 +17,36 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
-use std::borrow::Borrow;
-use std::ffi::OsStr;
-use std::path::Path;
+use std::fmt;
use std::time::Duration;
-use std::{fmt, io};
-use anyhow::Context;
-use async_compression::tokio::bufread::GzipDecoder;
use async_trait::async_trait;
-use bytes::Bytes;
use quickwit_actors::{ActorExitStatus, Mailbox};
-use quickwit_common::uri::Uri;
use quickwit_config::FileSourceParams;
-use quickwit_metastore::checkpoint::PartitionId;
+use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint};
use quickwit_proto::metastore::SourceType;
-use quickwit_proto::types::{Position, SourceId};
-use serde::Serialize;
-use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader};
-use tracing::info;
+use quickwit_proto::types::SourceId;
-use super::BatchBuilder;
+use super::doc_file_reader::ObjectUriBatchReader;
+#[cfg(feature = "queue-sources")]
+use super::queue_sources::coordinator::QueueCoordinator;
use crate::actors::DocProcessor;
use crate::source::{Source, SourceContext, SourceRuntime, TypedSourceFactory};
-/// Number of bytes after which a new batch is cut.
-pub(crate) const BATCH_NUM_BYTES_LIMIT: u64 = 500_000u64;
-
-#[derive(Default, Clone, Debug, Eq, PartialEq, Serialize)]
-pub struct FileSourceCounters {
- pub previous_offset: u64,
- pub current_offset: u64,
- pub num_lines_processed: u64,
+enum FileSourceState {
+ #[cfg(feature = "queue-sources")]
+ Notification(QueueCoordinator),
+ Filepath {
+ batch_reader: ObjectUriBatchReader,
+ num_bytes_processed: u64,
+ num_lines_processed: u64,
+ },
}
pub struct FileSource {
source_id: SourceId,
- params: FileSourceParams,
- counters: FileSourceCounters,
- reader: FileSourceReader,
+ state: FileSourceState,
+ source_type: SourceType,
}
impl fmt::Debug for FileSource {
@@ -66,66 +57,90 @@ impl fmt::Debug for FileSource {
#[async_trait]
impl Source for FileSource {
+ #[allow(unused_variables)]
+ async fn initialize(
+ &mut self,
+ doc_processor_mailbox: &Mailbox,
+ ctx: &SourceContext,
+ ) -> Result<(), ActorExitStatus> {
+ match &mut self.state {
+ #[cfg(feature = "queue-sources")]
+ FileSourceState::Notification(coordinator) => {
+ coordinator.initialize(doc_processor_mailbox, ctx).await
+ }
+ FileSourceState::Filepath { .. } => Ok(()),
+ }
+ }
+
+ #[allow(unused_variables)]
async fn emit_batches(
&mut self,
doc_processor_mailbox: &Mailbox,
ctx: &SourceContext,
) -> Result {
- // We collect batches of documents before sending them to the indexer.
- let limit_num_bytes = self.counters.previous_offset + BATCH_NUM_BYTES_LIMIT;
- let mut reached_eof = false;
- let mut batch_builder = BatchBuilder::new(SourceType::File);
-
- while self.counters.current_offset < limit_num_bytes {
- let mut doc_line = String::new();
- // guard the zone in case of slow read, such as reading from someone
- // typing to stdin
- let num_bytes = ctx
- .protect_future(self.reader.read_line(&mut doc_line))
- .await
- .map_err(anyhow::Error::from)?;
- if num_bytes == 0 {
- reached_eof = true;
- break;
+ match &mut self.state {
+ #[cfg(feature = "queue-sources")]
+ FileSourceState::Notification(coordinator) => {
+ coordinator.emit_batches(doc_processor_mailbox, ctx).await?;
}
- batch_builder.add_doc(Bytes::from(doc_line));
- self.counters.current_offset += num_bytes as u64;
- self.counters.num_lines_processed += 1;
- }
- if !batch_builder.docs.is_empty() {
- if let Some(filepath) = &self.params.filepath {
- let filepath_str = filepath
- .to_str()
- .context("path is invalid utf-8")?
- .to_string();
- let partition_id = PartitionId::from(filepath_str);
- batch_builder
- .checkpoint_delta
- .record_partition_delta(
- partition_id,
- Position::offset(self.counters.previous_offset),
- Position::offset(self.counters.current_offset),
- )
- .unwrap();
+ FileSourceState::Filepath {
+ batch_reader,
+ num_bytes_processed,
+ num_lines_processed,
+ } => {
+ let batch_builder = batch_reader
+ .read_batch(ctx.progress(), self.source_type)
+ .await?;
+ *num_bytes_processed += batch_builder.num_bytes;
+ *num_lines_processed += batch_builder.docs.len() as u64;
+ doc_processor_mailbox
+ .send_message(batch_builder.build())
+ .await?;
+ if batch_reader.is_eof() {
+ ctx.send_exit_with_success(doc_processor_mailbox).await?;
+ return Err(ActorExitStatus::Success);
+ }
}
- self.counters.previous_offset = self.counters.current_offset;
- ctx.send_message(doc_processor_mailbox, batch_builder.build())
- .await?;
}
- if reached_eof {
- info!("reached end of file");
- ctx.send_exit_with_success(doc_processor_mailbox).await?;
- return Err(ActorExitStatus::Success);
- }
- Ok(Duration::default())
+ Ok(Duration::ZERO)
}
fn name(&self) -> String {
format!("{:?}", self)
}
+ #[allow(unused_variables)]
+ async fn suggest_truncate(
+ &mut self,
+ checkpoint: SourceCheckpoint,
+ ctx: &SourceContext,
+ ) -> anyhow::Result<()> {
+ match &mut self.state {
+ #[cfg(feature = "queue-sources")]
+ FileSourceState::Notification(coordinator) => {
+ coordinator.suggest_truncate(checkpoint, ctx).await
+ }
+ FileSourceState::Filepath { .. } => Ok(()),
+ }
+ }
+
fn observable_state(&self) -> serde_json::Value {
- serde_json::to_value(&self.counters).unwrap()
+ match &self.state {
+ #[cfg(feature = "queue-sources")]
+ FileSourceState::Notification(coordinator) => {
+ serde_json::to_value(coordinator.observable_state()).unwrap()
+ }
+ FileSourceState::Filepath {
+ num_bytes_processed,
+ num_lines_processed,
+ ..
+ } => {
+ serde_json::json!({
+ "num_bytes_processed": num_bytes_processed,
+ "num_lines_processed": num_lines_processed,
+ })
+ }
+ }
}
}
@@ -140,116 +155,71 @@ impl TypedSourceFactory for FileSourceFactory {
source_runtime: SourceRuntime,
params: FileSourceParams,
) -> anyhow::Result {
- let checkpoint = source_runtime.fetch_checkpoint().await?;
- let mut offset = 0;
-
- let reader: FileSourceReader = if let Some(filepath) = ¶ms.filepath {
- let partition_id = PartitionId::from(filepath.to_string_lossy().borrow());
- offset = checkpoint
- .position_for_partition(&partition_id)
- .map(|position| {
- position
- .as_usize()
- .expect("file offset should be stored as usize")
- })
- .unwrap_or(0);
- let (dir_uri, file_name) = dir_and_filename(filepath)?;
- let storage = source_runtime.storage_resolver.resolve(&dir_uri).await?;
- let file_size = storage.file_num_bytes(file_name).await?.try_into().unwrap();
- // If it's a gzip file, we can't seek to a specific offset, we need to start from the
- // beginning of the file, decompress and skip the first `offset` bytes.
- if filepath.extension() == Some(OsStr::new("gz")) {
- let stream = storage.get_slice_stream(file_name, 0..file_size).await?;
- FileSourceReader::new(Box::new(GzipDecoder::new(BufReader::new(stream))), offset)
- } else {
- let stream = storage
- .get_slice_stream(file_name, offset..file_size)
- .await?;
- FileSourceReader::new(stream, 0)
+ let source_id = source_runtime.source_config.source_id.clone();
+ let source_type = source_runtime.source_config.source_type();
+ let state = match params {
+ FileSourceParams::Filepath(file_uri) => {
+ let partition_id = PartitionId::from(file_uri.as_str());
+ let position = source_runtime
+ .fetch_checkpoint()
+ .await?
+ .position_for_partition(&partition_id)
+ .cloned()
+ .unwrap_or_default();
+ let batch_reader = ObjectUriBatchReader::try_new(
+ &source_runtime.storage_resolver,
+ partition_id,
+ &file_uri,
+ position,
+ )
+ .await?;
+ FileSourceState::Filepath {
+ batch_reader,
+ num_bytes_processed: 0,
+ num_lines_processed: 0,
+ }
+ }
+ #[cfg(feature = "sqs")]
+ FileSourceParams::Notifications(quickwit_config::FileSourceNotification::Sqs(
+ sqs_config,
+ )) => {
+ let coordinator =
+ QueueCoordinator::try_from_sqs_config(sqs_config, source_runtime).await?;
+ FileSourceState::Notification(coordinator)
+ }
+ #[cfg(not(feature = "sqs"))]
+ FileSourceParams::Notifications(quickwit_config::FileSourceNotification::Sqs(_)) => {
+ anyhow::bail!("Quickwit was compiled without the `sqs` feature")
}
- } else {
- // We cannot use the checkpoint.
- FileSourceReader::new(Box::new(tokio::io::stdin()), 0)
- };
- let file_source = FileSource {
- source_id: source_runtime.source_id().to_string(),
- counters: FileSourceCounters {
- previous_offset: offset as u64,
- current_offset: offset as u64,
- num_lines_processed: 0,
- },
- reader,
- params,
};
- Ok(file_source)
- }
-}
-
-struct FileSourceReader {
- reader: BufReader>,
- num_bytes_to_skip: usize,
-}
-impl FileSourceReader {
- fn new(reader: Box, num_bytes_to_skip: usize) -> Self {
- Self {
- reader: BufReader::new(reader),
- num_bytes_to_skip,
- }
+ Ok(FileSource {
+ state,
+ source_id,
+ source_type,
+ })
}
-
- // This function is only called for GZIP file.
- // Because they cannot be sought into, we have to scan them to the right initial position.
- async fn skip(&mut self) -> io::Result<()> {
- // Allocate once a 64kb buffer.
- let mut buf = [0u8; 64000];
- while self.num_bytes_to_skip > 0 {
- let num_bytes_to_read = self.num_bytes_to_skip.min(buf.len());
- let num_bytes_read = self
- .reader
- .read_exact(&mut buf[..num_bytes_to_read])
- .await?;
- self.num_bytes_to_skip -= num_bytes_read;
- }
- Ok(())
- }
-
- async fn read_line<'a>(&mut self, buf: &'a mut String) -> io::Result {
- if self.num_bytes_to_skip > 0 {
- self.skip().await?;
- }
- self.reader.read_line(buf).await
- }
-}
-
-pub(crate) fn dir_and_filename(filepath: &Path) -> anyhow::Result<(Uri, &Path)> {
- let dir_uri: Uri = filepath
- .parent()
- .context("Parent directory could not be resolved")?
- .to_str()
- .context("Path cannot be turned to string")?
- .parse()?;
- let file_name = filepath
- .file_name()
- .context("Path does not appear to be a file")?;
- Ok((dir_uri, file_name.as_ref()))
}
#[cfg(test)]
mod tests {
- use std::io::{Cursor, Write};
use std::num::NonZeroUsize;
+ use std::str::FromStr;
- use async_compression::tokio::write::GzipEncoder;
+ use bytes::Bytes;
use quickwit_actors::{Command, Universe};
+ use quickwit_common::uri::Uri;
use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams};
- use quickwit_metastore::checkpoint::SourceCheckpointDelta;
- use quickwit_proto::types::IndexUid;
+ use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpointDelta};
+ use quickwit_proto::types::{IndexUid, Position};
use super::*;
use crate::models::RawDocBatch;
+ use crate::source::doc_file_reader::file_test_helpers::{
+ generate_dummy_doc_file, generate_index_doc_file, DUMMY_DOC,
+ };
use crate::source::tests::SourceRuntimeBuilder;
- use crate::source::SourceActor;
+ use crate::source::{SourceActor, BATCH_NUM_BYTES_LIMIT};
#[tokio::test]
async fn test_file_source() {
@@ -261,9 +231,9 @@ mod tests {
let universe = Universe::with_accelerated_time();
let (doc_processor_mailbox, indexer_inbox) = universe.create_test_mailbox();
let params = if gzip {
- FileSourceParams::file("data/test_corpus.json.gz")
+ FileSourceParams::from_filepath("data/test_corpus.json.gz").unwrap()
} else {
- FileSourceParams::file("data/test_corpus.json")
+ FileSourceParams::from_filepath("data/test_corpus.json").unwrap()
};
let source_config = SourceConfig {
source_id: "test-file-source".to_string(),
@@ -289,13 +259,13 @@ mod tests {
assert_eq!(
counters,
serde_json::json!({
- "previous_offset": 1030u64,
- "current_offset": 1030u64,
+ "num_bytes_processed": 1030u64,
"num_lines_processed": 4u32
})
);
let batch = indexer_inbox.drain_for_test();
assert_eq!(batch.len(), 2);
+ batch[0].downcast_ref::().unwrap();
assert!(matches!(
batch[1].downcast_ref::().unwrap(),
Command::ExitWithSuccess
@@ -312,33 +282,11 @@ mod tests {
quickwit_common::setup_logging_for_tests();
let universe = Universe::with_accelerated_time();
let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox();
- let mut documents_bytes = Vec::new();
- for _ in 0..20_000 {
- documents_bytes
- .write_all(r#"{"body": "hello happy tax payer!"}"#.as_bytes())
- .unwrap();
- documents_bytes.write_all("\n".as_bytes()).unwrap();
- }
- let mut temp_file: tempfile::NamedTempFile = if gzip {
- tempfile::Builder::new().suffix(".gz").tempfile().unwrap()
- } else {
- tempfile::NamedTempFile::new().unwrap()
- };
- if gzip {
- let gzip_documents = gzip_bytes(&documents_bytes).await;
- temp_file.write_all(&gzip_documents).unwrap();
- } else {
- temp_file.write_all(&documents_bytes).unwrap();
- }
- temp_file.flush().unwrap();
- let params = FileSourceParams::file(temp_file.path());
- let filepath = params
- .filepath
- .as_ref()
- .unwrap()
- .to_string_lossy()
- .to_string();
-
+ let lines = BATCH_NUM_BYTES_LIMIT as usize / DUMMY_DOC.len() + 1;
+ let (temp_file, temp_file_size) = generate_dummy_doc_file(gzip, lines).await;
+ let filepath = temp_file.path().to_str().unwrap();
+ let uri = Uri::from_str(filepath).unwrap();
+ let params = FileSourceParams::Filepath(uri.clone());
let source_config = SourceConfig {
source_id: "test-file-source".to_string(),
num_pipelines: NonZeroUsize::new(1).unwrap(),
@@ -363,9 +311,8 @@ mod tests {
assert_eq!(
counters,
serde_json::json!({
- "previous_offset": 700_000u64,
- "current_offset": 700_000u64,
- "num_lines_processed": 20_000u64
+ "num_lines_processed": lines,
+ "num_bytes_processed": temp_file_size,
})
);
let indexer_msgs = doc_processor_inbox.drain_for_test();
@@ -377,27 +324,19 @@ mod tests {
format!("{:?}", &batch1.checkpoint_delta),
format!(
"∆({}:{})",
- filepath, "(00000000000000000000..00000000000000500010]"
+ uri, "(00000000000000000000..00000000000005242895]"
)
);
assert_eq!(
- &extract_position_delta(&batch1.checkpoint_delta).unwrap(),
- "00000000000000000000..00000000000000500010"
- );
- assert_eq!(
- &extract_position_delta(&batch2.checkpoint_delta).unwrap(),
- "00000000000000500010..00000000000000700000"
+ format!("{:?}", &batch2.checkpoint_delta),
+ format!(
+ "∆({}:{})",
+ uri, "(00000000000005242895..~00000000000005397105]"
+ )
);
assert!(matches!(command, &Command::ExitWithSuccess));
}
- fn extract_position_delta(checkpoint_delta: &SourceCheckpointDelta) -> Option {
- let checkpoint_delta_str = format!("{checkpoint_delta:?}");
- let (_left, right) =
- &checkpoint_delta_str[..checkpoint_delta_str.len() - 2].rsplit_once('(')?;
- Some(right.to_string())
- }
-
#[tokio::test]
async fn test_file_source_resume_from_checkpoint() {
aux_test_file_source_resume_from_checkpoint(false).await;
@@ -408,27 +347,10 @@ mod tests {
quickwit_common::setup_logging_for_tests();
let universe = Universe::with_accelerated_time();
let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox();
- let mut documents_bytes = Vec::new();
- for i in 0..100 {
- documents_bytes
- .write_all(format!("{i}\n").as_bytes())
- .unwrap();
- }
- let mut temp_file: tempfile::NamedTempFile = if gzip {
- tempfile::Builder::new().suffix(".gz").tempfile().unwrap()
- } else {
- tempfile::NamedTempFile::new().unwrap()
- };
- let temp_file_path = temp_file.path().canonicalize().unwrap();
- if gzip {
- let gzipped_documents = gzip_bytes(&documents_bytes).await;
- temp_file.write_all(&gzipped_documents).unwrap();
- } else {
- temp_file.write_all(&documents_bytes).unwrap();
- }
- temp_file.flush().unwrap();
-
- let params = FileSourceParams::file(&temp_file_path);
+ let temp_file = generate_index_doc_file(gzip, 100).await;
+ let temp_file_path = temp_file.path().to_str().unwrap();
+ let uri = Uri::from_str(temp_file_path).unwrap();
+ let params = FileSourceParams::Filepath(uri.clone());
let source_config = SourceConfig {
source_id: "test-file-source".to_string(),
num_pipelines: NonZeroUsize::new(1).unwrap(),
@@ -437,11 +359,11 @@ mod tests {
transform_config: None,
input_format: SourceInputFormat::Json,
};
- let partition_id = PartitionId::from(temp_file_path.to_string_lossy().borrow());
+ let partition_id = PartitionId::from(uri.as_str());
let source_checkpoint_delta = SourceCheckpointDelta::from_partition_delta(
partition_id,
Position::Beginning,
- Position::offset(4u64),
+ Position::offset(16usize),
)
.unwrap();
@@ -465,74 +387,93 @@ mod tests {
assert_eq!(
counters,
serde_json::json!({
- "previous_offset": 290u64,
- "current_offset": 290u64,
- "num_lines_processed": 98u64
+ "num_bytes_processed": (800-16) as u64,
+ "num_lines_processed": (100-2) as u64,
})
);
let indexer_messages: Vec = doc_processor_inbox.drain_for_test_typed();
- assert!(&indexer_messages[0].docs[0].starts_with(b"2\n"));
+ assert_eq!(
+ indexer_messages[0].docs[0],
+ Bytes::from_static(b"0000002\n")
+ );
}
+}
- async fn gzip_bytes(bytes: &[u8]) -> Vec {
- let mut gzip_documents = Vec::new();
- let mut encoder = GzipEncoder::new(&mut gzip_documents);
- tokio::io::AsyncWriteExt::write_all(&mut encoder, bytes)
- .await
- .unwrap();
- // flush is not sufficient here and reading the file will raise a unexpected end of file
- // error.
- tokio::io::AsyncWriteExt::shutdown(&mut encoder)
+#[cfg(all(test, feature = "sqs-localstack-tests"))]
+mod localstack_tests {
+ use std::str::FromStr;
+
+ use quickwit_actors::Universe;
+ use quickwit_common::rand::append_random_suffix;
+ use quickwit_common::uri::Uri;
+ use quickwit_config::{
+ FileSourceMessageType, FileSourceNotification, FileSourceSqs, SourceConfig, SourceParams,
+ };
+ use quickwit_metastore::metastore_for_test;
+
+ use super::*;
+ use crate::models::RawDocBatch;
+ use crate::source::doc_file_reader::file_test_helpers::generate_dummy_doc_file;
+ use crate::source::queue_sources::sqs_queue::test_helpers::{
+ create_queue, get_localstack_sqs_client, send_message,
+ };
+ use crate::source::test_setup_helper::setup_index;
+ use crate::source::tests::SourceRuntimeBuilder;
+ use crate::source::SourceActor;
+
+ #[tokio::test]
+ async fn test_file_source_sqs_notifications() {
+ // queue setup
+ let sqs_client = get_localstack_sqs_client().await.unwrap();
+ let queue_url = create_queue(&sqs_client, "file-source-sqs-notifications").await;
+ let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ send_message(&sqs_client, &queue_url, test_uri.as_str()).await;
+
+ // source setup
+ let source_params =
+ FileSourceParams::Notifications(FileSourceNotification::Sqs(FileSourceSqs {
+ queue_url,
+ message_type: FileSourceMessageType::RawUri,
+ }));
+ let source_config = SourceConfig::for_test(
+ "test-file-source-sqs-notifications",
+ SourceParams::File(source_params.clone()),
+ );
+ let metastore = metastore_for_test();
+ let index_id = append_random_suffix("test-sqs-index");
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
+ let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config)
+ .with_metastore(metastore)
+ .build();
+ let sqs_source = FileSourceFactory::typed_create_source(source_runtime, source_params)
.await
.unwrap();
- gzip_documents
- }
- #[tokio::test]
- async fn test_skip_reader() {
- {
- // Skip 0 bytes.
- let mut reader = FileSourceReader::new(Box::new("hello".as_bytes()), 0);
- let mut buf = String::new();
- reader.read_line(&mut buf).await.unwrap();
- assert_eq!(buf, "hello");
- }
- {
- // Skip 2 bytes.
- let mut reader = FileSourceReader::new(Box::new("hello".as_bytes()), 2);
- let mut buf = String::new();
- reader.read_line(&mut buf).await.unwrap();
- assert_eq!(buf, "llo");
- }
- {
- let input = "hello";
- let cursor = Cursor::new(input);
- let mut reader = FileSourceReader::new(Box::new(cursor), 5);
- let mut buf = String::new();
- assert!(reader.read_line(&mut buf).await.is_ok());
- }
- {
- let input = "hello";
- let cursor = Cursor::new(input);
- let mut reader = FileSourceReader::new(Box::new(cursor), 10);
- let mut buf = String::new();
- assert!(reader.read_line(&mut buf).await.is_err());
- }
- {
- let input = "hello world".repeat(10000);
- let cursor = Cursor::new(input.clone());
- let mut reader = FileSourceReader::new(Box::new(cursor), 64000);
- let mut buf = String::new();
- reader.read_line(&mut buf).await.unwrap();
- assert_eq!(buf, input[64000..]);
- }
+ // actor setup
+ let universe = Universe::with_accelerated_time();
+ let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox();
{
- let input = "hello world".repeat(10000);
- let cursor = Cursor::new(input.clone());
- let mut reader = FileSourceReader::new(Box::new(cursor), 64001);
- let mut buf = String::new();
- reader.read_line(&mut buf).await.unwrap();
- assert_eq!(buf, input[64001..]);
+ let actor = SourceActor {
+ source: Box::new(sqs_source),
+ doc_processor_mailbox: doc_processor_mailbox.clone(),
+ };
+ let (_mailbox, handle) = universe.spawn_builder().spawn(actor);
+
+ // run the source actor for a while
+ tokio::time::timeout(Duration::from_millis(500), handle.join())
+ .await
+ .unwrap_err();
+
+ let next_message = doc_processor_inbox
+ .drain_for_test()
+ .into_iter()
+ .flat_map(|box_any| box_any.downcast::().ok())
+ .map(|box_raw_doc_batch| *box_raw_doc_batch)
+ .next()
+ .unwrap();
+ assert_eq!(next_message.docs.len(), 10);
}
+ universe.assert_quit().await;
}
}
diff --git a/quickwit/quickwit-indexing/src/source/ingest/mod.rs b/quickwit/quickwit-indexing/src/source/ingest/mod.rs
index 1d1855bbaed..76aafab3344 100644
--- a/quickwit/quickwit-indexing/src/source/ingest/mod.rs
+++ b/quickwit/quickwit-indexing/src/source/ingest/mod.rs
@@ -677,7 +677,7 @@ mod tests {
use quickwit_common::metrics::MEMORY_METRICS;
use quickwit_common::stream_utils::InFlightValue;
use quickwit_common::ServiceStream;
- use quickwit_config::{SourceConfig, SourceParams};
+ use quickwit_config::{IndexingSettings, SourceConfig, SourceParams};
use quickwit_proto::indexing::IndexingPipelineId;
use quickwit_proto::ingest::ingester::{
FetchMessage, IngesterServiceClient, MockIngesterService, TruncateShardsResponse,
@@ -944,6 +944,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::no_retries();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1145,6 +1146,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1307,6 +1309,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1372,6 +1375,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1604,6 +1608,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1758,6 +1763,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker,
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
@@ -1889,6 +1895,7 @@ mod tests {
queues_dir_path: PathBuf::from("./queues"),
storage_resolver: StorageResolver::for_test(),
event_broker: event_broker.clone(),
+ indexing_setting: IndexingSettings::default(),
};
let retry_params = RetryParams::for_test();
let mut source = IngestSource::try_new(source_runtime, retry_params)
diff --git a/quickwit/quickwit-indexing/src/source/kafka_source.rs b/quickwit/quickwit-indexing/src/source/kafka_source.rs
index ea4e26e77c6..6aae5b7c5bb 100644
--- a/quickwit/quickwit-indexing/src/source/kafka_source.rs
+++ b/quickwit/quickwit-indexing/src/source/kafka_source.rs
@@ -766,15 +766,9 @@ mod kafka_broker_tests {
use quickwit_actors::{ActorContext, Universe};
use quickwit_common::rand::append_random_suffix;
- use quickwit_config::{IndexConfig, SourceConfig, SourceInputFormat, SourceParams};
- use quickwit_metastore::checkpoint::{IndexCheckpointDelta, SourceCheckpointDelta};
- use quickwit_metastore::{
- metastore_for_test, CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt,
- };
- use quickwit_proto::metastore::{
- CreateIndexRequest, MetastoreService, MetastoreServiceClient, PublishSplitsRequest,
- StageSplitsRequest,
- };
+ use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams};
+ use quickwit_metastore::checkpoint::SourceCheckpointDelta;
+ use quickwit_metastore::metastore_for_test;
use quickwit_proto::types::IndexUid;
use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
use rdkafka::client::DefaultClientContext;
@@ -783,7 +777,7 @@ mod kafka_broker_tests {
use tokio::sync::watch;
use super::*;
- use crate::new_split_id;
+ use crate::source::test_setup_helper::setup_index;
use crate::source::tests::SourceRuntimeBuilder;
use crate::source::{quickwit_supported_sources, RawDocBatch, SourceActor};
@@ -915,71 +909,6 @@ mod kafka_broker_tests {
Ok(merged_batch)
}
- async fn setup_index(
- metastore: MetastoreServiceClient,
- index_id: &str,
- source_config: &SourceConfig,
- partition_deltas: &[(u64, i64, i64)],
- ) -> IndexUid {
- let index_uri = format!("ram:///indexes/{index_id}");
- let index_config = IndexConfig::for_test(index_id, &index_uri);
- let create_index_request = CreateIndexRequest::try_from_index_and_source_configs(
- &index_config,
- &[source_config.clone()],
- )
- .unwrap();
- let index_uid: IndexUid = metastore
- .create_index(create_index_request)
- .await
- .unwrap()
- .index_uid()
- .clone();
-
- if partition_deltas.is_empty() {
- return index_uid;
- }
- let split_id = new_split_id();
- let split_metadata = SplitMetadata::for_test(split_id.clone());
- let stage_splits_request =
- StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata)
- .unwrap();
- metastore.stage_splits(stage_splits_request).await.unwrap();
-
- let mut source_delta = SourceCheckpointDelta::default();
- for (partition_id, from_position, to_position) in partition_deltas.iter().copied() {
- source_delta
- .record_partition_delta(
- partition_id.into(),
- {
- if from_position < 0 {
- Position::Beginning
- } else {
- Position::offset(from_position as u64)
- }
- },
- Position::offset(to_position as u64),
- )
- .unwrap();
- }
- let checkpoint_delta = IndexCheckpointDelta {
- source_id: source_config.source_id.to_string(),
- source_delta,
- };
- let checkpoint_delta_json = serde_json::to_string(&checkpoint_delta).unwrap();
- let publish_splits_request = PublishSplitsRequest {
- index_uid: Some(index_uid.clone()),
- index_checkpoint_delta_json_opt: Some(checkpoint_delta_json),
- staged_split_ids: vec![split_id.clone()],
- replaced_split_ids: Vec::new(),
- publish_token_opt: None,
- };
- metastore
- .publish_splits(publish_splits_request)
- .await
- .unwrap();
- index_uid
- }
-
#[tokio::test]
async fn test_kafka_source_process_message() {
let admin_client = create_admin_client();
@@ -1110,8 +1039,17 @@ mod kafka_broker_tests {
let index_id = append_random_suffix("test-kafka-source--process-assign-partitions--index");
let (_source_id, source_config) = get_source_config(&topic, "earliest");
- let index_uid =
- setup_index(metastore.clone(), &index_id, &source_config, &[(2, -1, 42)]).await;
+ let index_uid = setup_index(
+ metastore.clone(),
+ &index_id,
+ &source_config,
+ &[(
+ PartitionId::from(2u64),
+ Position::Beginning,
+ Position::offset(42u64),
+ )],
+ )
+ .await;
let SourceParams::Kafka(params) = source_config.clone().source_params else {
panic!(
@@ -1243,8 +1181,17 @@ mod kafka_broker_tests {
let metastore = metastore_for_test();
let index_id = append_random_suffix("test-kafka-source--suggest-truncate--index");
let (_source_id, source_config) = get_source_config(&topic, "earliest");
- let index_uid =
- setup_index(metastore.clone(), &index_id, &source_config, &[(2, -1, 42)]).await;
+ let index_uid = setup_index(
+ metastore.clone(),
+ &index_id,
+ &source_config,
+ &[(
+ PartitionId::from(2u64),
+ Position::Beginning,
+ Position::offset(42u64),
+ )],
+ )
+ .await;
let SourceParams::Kafka(params) = source_config.clone().source_params else {
panic!(
@@ -1436,7 +1383,18 @@ mod kafka_broker_tests {
metastore.clone(),
&index_id,
&source_config,
- &[(0, -1, 0), (1, -1, 2)],
+ &[
+ (
+ PartitionId::from(0u64),
+ Position::Beginning,
+ Position::offset(0u64),
+ ),
+ (
+ PartitionId::from(1u64),
+ Position::Beginning,
+ Position::offset(2u64),
+ ),
+ ],
)
.await;
let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config)
diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs
index db2583d0e95..20d1bad6fc4 100644
--- a/quickwit/quickwit-indexing/src/source/mod.rs
+++ b/quickwit/quickwit-indexing/src/source/mod.rs
@@ -57,6 +57,7 @@
//! that file.
//! - the kafka source: the partition id is a kafka topic partition id, and the position is a kafka
//! offset.
+mod doc_file_reader;
mod file_source;
#[cfg(feature = "gcp-pubsub")]
mod gcp_pubsub_source;
@@ -68,7 +69,10 @@ mod kafka_source;
mod kinesis;
#[cfg(feature = "pulsar")]
mod pulsar_source;
+#[cfg(feature = "queue-sources")]
+mod queue_sources;
mod source_factory;
+mod stdin_source;
mod vec_source;
mod void_source;
@@ -89,11 +93,15 @@ pub use kinesis::kinesis_source::{KinesisSource, KinesisSourceFactory};
use once_cell::sync::OnceCell;
#[cfg(feature = "pulsar")]
pub use pulsar_source::{PulsarSource, PulsarSourceFactory};
+#[cfg(feature = "sqs")]
+pub use queue_sources::sqs_queue;
use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler, Mailbox};
use quickwit_common::metrics::{GaugeGuard, MEMORY_METRICS};
use quickwit_common::pubsub::EventBroker;
use quickwit_common::runtimes::RuntimeType;
-use quickwit_config::{SourceConfig, SourceParams};
+use quickwit_config::{
+ FileSourceNotification, FileSourceParams, IndexingSettings, SourceConfig, SourceParams,
+};
use quickwit_ingest::IngesterPool;
use quickwit_metastore::checkpoint::{SourceCheckpoint, SourceCheckpointDelta};
use quickwit_metastore::IndexMetadataResponseExt;
@@ -111,7 +119,7 @@ use tracing::error;
pub use vec_source::{VecSource, VecSourceFactory};
pub use void_source::{VoidSource, VoidSourceFactory};
-use self::file_source::dir_and_filename;
+use self::doc_file_reader::dir_and_filename;
use crate::actors::DocProcessor;
use crate::models::RawDocBatch;
use crate::source::ingest::IngestSourceFactory;
@@ -143,6 +151,7 @@ pub struct SourceRuntime {
pub queues_dir_path: PathBuf,
pub storage_resolver: StorageResolver,
pub event_broker: EventBroker,
+ pub indexing_setting: IndexingSettings,
}
impl SourceRuntime {
@@ -232,7 +241,7 @@ pub trait Source: Send + 'static {
/// In that case, `batch_sink` will block.
///
/// It returns an optional duration specifying how long the batch requester
- /// should wait before pooling gain.
+ /// should wait before polling again.
async fn emit_batches(
&mut self,
doc_processor_mailbox: &Mailbox,
@@ -410,15 +419,26 @@ pub async fn check_source_connectivity(
source_config: &SourceConfig,
) -> anyhow::Result<()> {
match &source_config.source_params {
- SourceParams::File(params) => {
- if let Some(filepath) = ¶ms.filepath {
- let (dir_uri, file_name) = dir_and_filename(filepath)?;
- let storage = storage_resolver.resolve(&dir_uri).await?;
- storage.file_num_bytes(file_name).await?;
- }
+ SourceParams::File(FileSourceParams::Filepath(file_uri)) => {
+ let (dir_uri, file_name) = dir_and_filename(file_uri)?;
+ let storage = storage_resolver.resolve(&dir_uri).await?;
+ storage.file_num_bytes(file_name).await?;
Ok(())
}
#[allow(unused_variables)]
+ SourceParams::File(FileSourceParams::Notifications(FileSourceNotification::Sqs(
+ sqs_config,
+ ))) => {
+ #[cfg(not(feature = "sqs"))]
+ anyhow::bail!("Quickwit was compiled without the `sqs` feature");
+
+ #[cfg(feature = "sqs")]
+ {
+ queue_sources::sqs_queue::check_connectivity(&sqs_config.queue_url).await?;
+ Ok(())
+ }
+ }
+ #[allow(unused_variables)]
SourceParams::Kafka(params) => {
#[cfg(not(feature = "kafka"))]
anyhow::bail!("Quickwit was compiled without the `kafka` feature");
@@ -591,10 +611,11 @@ mod tests {
source_config: self.source_config,
storage_resolver: StorageResolver::for_test(),
event_broker: EventBroker::default(),
+ indexing_setting: IndexingSettings::default(),
}
}
- #[cfg(feature = "kafka")]
+ #[cfg(any(feature = "kafka", feature = "sqs"))]
pub fn with_metastore(mut self, metastore: MetastoreServiceClient) -> Self {
self.metastore_opt = Some(metastore);
self
@@ -676,7 +697,7 @@ mod tests {
source_id: "file".to_string(),
num_pipelines: NonZeroUsize::new(1).unwrap(),
enabled: true,
- source_params: SourceParams::file("file-does-not-exist.json"),
+ source_params: SourceParams::file_from_str("file-does-not-exist.json").unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -691,7 +712,7 @@ mod tests {
source_id: "file".to_string(),
num_pipelines: NonZeroUsize::new(1).unwrap(),
enabled: true,
- source_params: SourceParams::file("data/test_corpus.json"),
+ source_params: SourceParams::file_from_str("data/test_corpus.json").unwrap(),
transform_config: None,
input_format: SourceInputFormat::Json,
};
@@ -704,3 +725,74 @@ mod tests {
Ok(())
}
}
+
+#[cfg(all(
+ test,
+ any(feature = "sqs-localstack-tests", feature = "kafka-broker-tests")
+))]
+mod test_setup_helper {
+
+ use quickwit_config::IndexConfig;
+ use quickwit_metastore::checkpoint::{IndexCheckpointDelta, PartitionId};
+ use quickwit_metastore::{CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt};
+ use quickwit_proto::metastore::{CreateIndexRequest, PublishSplitsRequest, StageSplitsRequest};
+ use quickwit_proto::types::Position;
+
+ use super::*;
+ use crate::new_split_id;
+
+ pub async fn setup_index(
+ metastore: MetastoreServiceClient,
+ index_id: &str,
+ source_config: &SourceConfig,
+ partition_deltas: &[(PartitionId, Position, Position)],
+ ) -> IndexUid {
+ let index_uri = format!("ram:///indexes/{index_id}");
+ let index_config = IndexConfig::for_test(index_id, &index_uri);
+ let create_index_request = CreateIndexRequest::try_from_index_and_source_configs(
+ &index_config,
+ &[source_config.clone()],
+ )
+ .unwrap();
+ let index_uid: IndexUid = metastore
+ .create_index(create_index_request)
+ .await
+ .unwrap()
+ .index_uid()
+ .clone();
+
+ if partition_deltas.is_empty() {
+ return index_uid;
+ }
+ let split_id = new_split_id();
+ let split_metadata = SplitMetadata::for_test(split_id.clone());
+ let stage_splits_request =
+ StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata)
+ .unwrap();
+ metastore.stage_splits(stage_splits_request).await.unwrap();
+
+ let mut source_delta = SourceCheckpointDelta::default();
+ for (partition_id, from_position, to_position) in partition_deltas.iter().cloned() {
+ source_delta
+ .record_partition_delta(partition_id, from_position, to_position)
+ .unwrap();
+ }
+ let checkpoint_delta = IndexCheckpointDelta {
+ source_id: source_config.source_id.to_string(),
+ source_delta,
+ };
+ let checkpoint_delta_json = serde_json::to_string(&checkpoint_delta).unwrap();
+ let publish_splits_request = PublishSplitsRequest {
+ index_uid: Some(index_uid.clone()),
+ index_checkpoint_delta_json_opt: Some(checkpoint_delta_json),
+ staged_split_ids: vec![split_id.clone()],
+ replaced_split_ids: Vec::new(),
+ publish_token_opt: None,
+ };
+ metastore
+ .publish_splits(publish_splits_request)
+ .await
+ .unwrap();
+ index_uid
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/pulsar_source.rs b/quickwit/quickwit-indexing/src/source/pulsar_source.rs
index 6dec087a454..6dcc91abc71 100644
--- a/quickwit/quickwit-indexing/src/source/pulsar_source.rs
+++ b/quickwit/quickwit-indexing/src/source/pulsar_source.rs
@@ -445,22 +445,15 @@ mod pulsar_broker_tests {
use futures::future::join_all;
use quickwit_actors::{ActorHandle, Inbox, Universe, HEARTBEAT};
use quickwit_common::rand::append_random_suffix;
- use quickwit_config::{IndexConfig, SourceConfig, SourceInputFormat, SourceParams};
- use quickwit_metastore::checkpoint::{
- IndexCheckpointDelta, PartitionId, SourceCheckpointDelta,
- };
- use quickwit_metastore::{
- metastore_for_test, CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt,
- };
- use quickwit_proto::metastore::{
- CreateIndexRequest, MetastoreService, MetastoreServiceClient, PublishSplitsRequest,
- StageSplitsRequest,
- };
+ use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams};
+ use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpointDelta};
+ use quickwit_metastore::metastore_for_test;
+ use quickwit_proto::metastore::MetastoreServiceClient;
use reqwest::StatusCode;
use super::*;
- use crate::new_split_id;
use crate::source::pulsar_source::{msg_id_from_position, msg_id_to_position};
+ use crate::source::test_setup_helper::setup_index;
use crate::source::tests::SourceRuntimeBuilder;
use crate::source::{quickwit_supported_sources, RawDocBatch, SuggestTruncate};
@@ -492,63 +485,6 @@ mod pulsar_broker_tests {
}};
}
- async fn setup_index(
- metastore: MetastoreServiceClient,
- index_id: &str,
- source_id: &str,
- partition_deltas: &[(&str, Position, Position)],
- ) -> IndexUid {
- let index_uri = format!("ram:///indexes/{index_id}");
- let index_config = IndexConfig::for_test(index_id, &index_uri);
- let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
- let index_uid: IndexUid = metastore
- .create_index(create_index_request)
- .await
- .unwrap()
- .index_uid()
- .clone();
-
- if partition_deltas.is_empty() {
- return index_uid;
- }
- let split_id = new_split_id();
- let split_metadata = SplitMetadata::for_test(split_id.clone());
- let stage_splits_request =
- StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata)
- .unwrap();
- metastore.stage_splits(stage_splits_request).await.unwrap();
-
- let mut source_delta = SourceCheckpointDelta::default();
- for (partition_id, from_position, to_position) in partition_deltas {
- source_delta
- .record_partition_delta(
- PartitionId::from(&**partition_id),
- from_position.clone(),
- to_position.clone(),
- )
- .unwrap();
- }
- let checkpoint_delta = IndexCheckpointDelta {
- source_id: source_id.to_string(),
- source_delta,
- };
- let publish_splits_request = PublishSplitsRequest {
- index_uid: Some(index_uid.clone()),
- staged_split_ids: vec![split_id.clone()],
- replaced_split_ids: Vec::new(),
- index_checkpoint_delta_json_opt: Some(
- serde_json::to_string(&checkpoint_delta).unwrap(),
- ),
- publish_token_opt: None,
- };
- metastore
- .publish_splits(publish_splits_request)
- .await
- .unwrap();
- index_uid
- }
-
fn get_source_config>(
topics: impl IntoIterator- ,
) -> (String, SourceConfig) {
@@ -895,7 +831,7 @@ mod pulsar_broker_tests {
let index_id = append_random_suffix("test-pulsar-source--topic-ingestion--index");
let (source_id, source_config) = get_source_config([&topic]);
- let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await;
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
let (source_handle, doc_processor_inbox) = create_source(
&universe,
@@ -952,7 +888,7 @@ mod pulsar_broker_tests {
let index_id = append_random_suffix("test-pulsar-source--topic-ingestion--index");
let (source_id, source_config) = get_source_config([&topic1, &topic2]);
- let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await;
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
let (source_handle, doc_processor_inbox) = create_source(
&universe,
@@ -1020,7 +956,7 @@ mod pulsar_broker_tests {
let (source_id, source_config) = get_source_config([&topic]);
create_partitioned_topic(&topic, 2).await;
- let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await;
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
let (source_handle, doc_processor_inbox) = create_source(
&universe,
@@ -1074,7 +1010,7 @@ mod pulsar_broker_tests {
let (source_id, source_config) = get_source_config([&topic]);
create_partitioned_topic(&topic, 2).await;
- let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await;
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
let topic_partition_1 = format!("{topic}-partition-0");
let topic_partition_2 = format!("{topic}-partition-1");
@@ -1158,10 +1094,10 @@ mod pulsar_broker_tests {
let index_id =
append_random_suffix("test-pulsar-source--partitioned-multi-consumer-failure--index");
- let (source_id, source_config) = get_source_config([&topic]);
+ let (_, source_config) = get_source_config([&topic]);
create_partitioned_topic(&topic, 2).await;
- let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await;
+ let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await;
let topic_partition_1 = format!("{topic}-partition-0");
let topic_partition_2 = format!("{topic}-partition-1");
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs b/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs
new file mode 100644
index 00000000000..bd00840f657
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs
@@ -0,0 +1,521 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::fmt;
+use std::sync::Arc;
+use std::time::Duration;
+
+use itertools::Itertools;
+use quickwit_actors::{ActorExitStatus, Mailbox};
+use quickwit_common::rate_limited_error;
+use quickwit_config::{FileSourceMessageType, FileSourceSqs};
+use quickwit_metastore::checkpoint::SourceCheckpoint;
+use quickwit_proto::indexing::IndexingPipelineId;
+use quickwit_proto::metastore::SourceType;
+use quickwit_storage::StorageResolver;
+use serde::Serialize;
+use ulid::Ulid;
+
+use super::helpers::QueueReceiver;
+use super::local_state::QueueLocalState;
+use super::message::{MessageType, PreProcessingError, ReadyMessage};
+use super::shared_state::{checkpoint_messages, QueueSharedState};
+use super::visibility::{spawn_visibility_task, VisibilitySettings};
+use super::Queue;
+use crate::actors::DocProcessor;
+use crate::models::{NewPublishLock, NewPublishToken, PublishLock};
+use crate::source::{SourceContext, SourceRuntime};
+
+/// Maximum duration that the `emit_batches()` callback can wait for
+/// `queue.receive()` calls. If too small, the actor loop will spin
+/// un-necessarily. If too large, the actor loop will be slow to react to new
+/// messages (or shutdown).
+pub const RECEIVE_POLL_TIMEOUT: Duration = Duration::from_millis(500);
+
+#[derive(Default, Serialize)]
+pub struct QueueCoordinatorObservableState {
+ /// Number of bytes processed by the source.
+ pub num_bytes_processed: u64,
+ /// Number of lines processed by the source.
+ pub num_lines_processed: u64,
+ /// Number of messages processed by the source.
+ pub num_messages_processed: u64,
+ /// Number of messages that could not be pre-processed.
+ pub num_messages_failed_preprocessing: u64,
+ /// Number of messages that could not be moved to in-progress.
+ pub num_messages_failed_opening: u64,
+}
+
+/// The `QueueCoordinator` fetches messages from a queue, converts them into
+/// record batches, and tracks the messages' state until their entire content is
+/// published. Its API closely resembles the [`crate::source::Source`] trait,
+/// making the implementation of queue sources straightforward.
+pub struct QueueCoordinator {
+ storage_resolver: StorageResolver,
+ pipeline_id: IndexingPipelineId,
+ source_type: SourceType,
+ queue: Arc,
+ queue_receiver: QueueReceiver,
+ observable_state: QueueCoordinatorObservableState,
+ message_type: MessageType,
+ publish_lock: PublishLock,
+ shared_state: QueueSharedState,
+ local_state: QueueLocalState,
+ publish_token: String,
+ visible_settings: VisibilitySettings,
+}
+
+impl fmt::Debug for QueueCoordinator {
+ fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+ formatter
+ .debug_struct("QueueTracker")
+ .field("index_id", &self.pipeline_id.index_uid.index_id)
+ .field("queue", &self.queue)
+ .finish()
+ }
+}
+
+impl QueueCoordinator {
+ pub fn new(
+ source_runtime: SourceRuntime,
+ queue: Arc,
+ message_type: MessageType,
+ ) -> Self {
+ Self {
+ shared_state: QueueSharedState {
+ metastore: source_runtime.metastore,
+ source_id: source_runtime.pipeline_id.source_id.clone(),
+ index_uid: source_runtime.pipeline_id.index_uid.clone(),
+ },
+ local_state: QueueLocalState::default(),
+ pipeline_id: source_runtime.pipeline_id,
+ source_type: source_runtime.source_config.source_type(),
+ storage_resolver: source_runtime.storage_resolver,
+ queue_receiver: QueueReceiver::new(queue.clone(), RECEIVE_POLL_TIMEOUT),
+ queue,
+ observable_state: QueueCoordinatorObservableState::default(),
+ message_type,
+ publish_lock: PublishLock::default(),
+ publish_token: Ulid::new().to_string(),
+ visible_settings: VisibilitySettings::from_commit_timeout(
+ source_runtime.indexing_setting.commit_timeout_secs,
+ ),
+ }
+ }
+
+ #[cfg(feature = "sqs")]
+ pub async fn try_from_sqs_config(
+ config: FileSourceSqs,
+ source_runtime: SourceRuntime,
+ ) -> anyhow::Result {
+ use super::sqs_queue::SqsQueue;
+ let queue = SqsQueue::try_new(config.queue_url).await?;
+ let message_type = match config.message_type {
+ FileSourceMessageType::S3Notification => MessageType::S3Notification,
+ FileSourceMessageType::RawUri => MessageType::RawUri,
+ };
+ Ok(QueueCoordinator::new(
+ source_runtime,
+ Arc::new(queue),
+ message_type,
+ ))
+ }
+
+ pub async fn initialize(
+ &mut self,
+ doc_processor_mailbox: &Mailbox,
+ ctx: &SourceContext,
+ ) -> Result<(), ActorExitStatus> {
+ let publish_lock = self.publish_lock.clone();
+ ctx.send_message(doc_processor_mailbox, NewPublishLock(publish_lock))
+ .await?;
+ ctx.send_message(
+ doc_processor_mailbox,
+ NewPublishToken(self.publish_token.clone()),
+ )
+ .await?;
+ Ok(())
+ }
+
+ /// Polls messages from the queue and prepares them for processing
+ async fn poll_messages(&mut self, ctx: &SourceContext) -> Result<(), ActorExitStatus> {
+ let raw_messages = self
+ .queue_receiver
+ .receive(1, self.visible_settings.deadline_for_receive)
+ .await?;
+
+ let mut format_errors = Vec::new();
+ let mut discardable_ack_ids = Vec::new();
+ let mut preprocessed_messages = Vec::new();
+ for message in raw_messages {
+ match message.pre_process(self.message_type) {
+ Ok(preprocessed_message) => preprocessed_messages.push(preprocessed_message),
+ Err(PreProcessingError::UnexpectedFormat(err)) => format_errors.push(err),
+ Err(PreProcessingError::Discardable { ack_id }) => discardable_ack_ids.push(ack_id),
+ }
+ }
+ if !format_errors.is_empty() {
+ self.observable_state.num_messages_failed_preprocessing += format_errors.len() as u64;
+ rate_limited_error!(
+ limit_per_min = 10,
+ count = format_errors.len(),
+ last_err = ?format_errors.last().unwrap(),
+ "invalid messages not processed, use a dead letter queue to limit retries"
+ );
+ }
+ if preprocessed_messages.is_empty() {
+ self.queue.acknowledge(&discardable_ack_ids).await?;
+ return Ok(());
+ }
+
+ // in rare situations, there might be duplicates within a batch
+ let deduplicated_messages = preprocessed_messages
+ .into_iter()
+ .unique_by(|x| x.partition_id());
+
+ let mut untracked_locally = Vec::new();
+ let mut already_completed = Vec::new();
+ for message in deduplicated_messages {
+ let partition_id = message.partition_id();
+ if self.local_state.is_completed(&partition_id) {
+ already_completed.push(message);
+ } else if !self.local_state.is_tracked(&partition_id) {
+ untracked_locally.push(message);
+ }
+ }
+
+ let checkpointed_messages =
+ checkpoint_messages(&self.shared_state, &self.publish_token, untracked_locally).await?;
+
+ let mut ready_messages = Vec::new();
+ for (message, position) in checkpointed_messages {
+ if position.is_eof() {
+ self.local_state.mark_completed(message.partition_id());
+ already_completed.push(message);
+ } else {
+ ready_messages.push(ReadyMessage {
+ visibility_handle: spawn_visibility_task(
+ ctx,
+ self.queue.clone(),
+ message.metadata.ack_id.clone(),
+ message.metadata.initial_deadline,
+ self.visible_settings.clone(),
+ ),
+ content: message,
+ position,
+ })
+ }
+ }
+
+ self.local_state.set_ready_for_read(ready_messages);
+
+ // Acknowledge messages that already have been processed
+ let mut ack_ids = already_completed
+ .iter()
+ .map(|msg| msg.metadata.ack_id.clone())
+ .collect::>();
+ ack_ids.append(&mut discardable_ack_ids);
+ self.queue.acknowledge(&ack_ids).await?;
+
+ Ok(())
+ }
+
+ pub async fn emit_batches(
+ &mut self,
+ doc_processor_mailbox: &Mailbox,
+ ctx: &SourceContext,
+ ) -> Result {
+ if let Some(in_progress_ref) = self.local_state.read_in_progress_mut() {
+ // TODO: should we kill the publish lock if the message visibility extension failed?
+ let batch_builder = in_progress_ref
+ .batch_reader
+ .read_batch(ctx.progress(), self.source_type)
+ .await?;
+ self.observable_state.num_lines_processed += batch_builder.docs.len() as u64;
+ self.observable_state.num_bytes_processed += batch_builder.num_bytes;
+ doc_processor_mailbox
+ .send_message(batch_builder.build())
+ .await?;
+ if in_progress_ref.batch_reader.is_eof() {
+ self.local_state
+ .drop_currently_read(self.visible_settings.deadline_for_last_extension)
+ .await?;
+ self.observable_state.num_messages_processed += 1;
+ }
+ } else if let Some(ready_message) = self.local_state.get_ready_for_read() {
+ match ready_message.start_processing(&self.storage_resolver).await {
+ Ok(new_in_progress) => {
+ self.local_state.set_currently_read(new_in_progress)?;
+ }
+ Err(err) => {
+ self.observable_state.num_messages_failed_opening += 1;
+ rate_limited_error!(
+ limit_per_min = 5,
+ err = ?err,
+ "failed to start message processing"
+ );
+ }
+ }
+ } else {
+ self.poll_messages(ctx).await?;
+ }
+
+ Ok(Duration::ZERO)
+ }
+
+ pub async fn suggest_truncate(
+ &mut self,
+ checkpoint: SourceCheckpoint,
+ _ctx: &SourceContext,
+ ) -> anyhow::Result<()> {
+ let committed_partition_ids = checkpoint
+ .iter()
+ .filter(|(_, pos)| pos.is_eof())
+ .map(|(pid, _)| pid)
+ .collect::>();
+ let mut completed = Vec::new();
+ for partition_id in committed_partition_ids {
+ let ack_id_opt = self.local_state.mark_completed(partition_id);
+ if let Some(ack_id) = ack_id_opt {
+ completed.push(ack_id);
+ }
+ }
+ self.queue.acknowledge(&completed).await
+ }
+
+ pub fn observable_state(&self) -> &QueueCoordinatorObservableState {
+ &self.observable_state
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::str::FromStr;
+
+ use quickwit_actors::{ActorContext, Universe};
+ use quickwit_common::uri::Uri;
+ use quickwit_proto::types::{NodeId, PipelineUid, Position};
+ use tokio::sync::watch;
+ use ulid::Ulid;
+
+ use super::*;
+ use crate::models::RawDocBatch;
+ use crate::source::doc_file_reader::file_test_helpers::{generate_dummy_doc_file, DUMMY_DOC};
+ use crate::source::queue_sources::memory_queue::MemoryQueueForTests;
+ use crate::source::queue_sources::message::PreProcessedPayload;
+ use crate::source::queue_sources::shared_state::shared_state_for_tests::shared_state_for_tests;
+ use crate::source::{SourceActor, BATCH_NUM_BYTES_LIMIT};
+
+ fn setup_coordinator(
+ queue: Arc,
+ shared_state: QueueSharedState,
+ ) -> QueueCoordinator {
+ let pipeline_id = IndexingPipelineId {
+ node_id: NodeId::from_str("test-node").unwrap(),
+ index_uid: shared_state.index_uid.clone(),
+ source_id: shared_state.source_id.clone(),
+ pipeline_uid: PipelineUid::random(),
+ };
+
+ QueueCoordinator {
+ local_state: QueueLocalState::default(),
+ shared_state,
+ pipeline_id,
+ observable_state: QueueCoordinatorObservableState::default(),
+ publish_lock: PublishLock::default(),
+ // set a very high chunking timeout to make it possible to count the
+ // number of iterations required to process messages
+ queue_receiver: QueueReceiver::new(queue.clone(), Duration::from_secs(10)),
+ queue,
+ message_type: MessageType::RawUri,
+ source_type: SourceType::Unspecified,
+ storage_resolver: StorageResolver::for_test(),
+ publish_token: Ulid::new().to_string(),
+ visible_settings: VisibilitySettings::from_commit_timeout(5),
+ }
+ }
+
+ async fn process_messages(
+ coordinator: &mut QueueCoordinator,
+ queue: Arc,
+ messages: &[(&Uri, &str)],
+ ) -> Vec {
+ let universe = Universe::with_accelerated_time();
+ let (source_mailbox, _source_inbox) = universe.create_test_mailbox::();
+ let (doc_processor_mailbox, doc_processor_inbox) =
+ universe.create_test_mailbox::();
+ let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null);
+ let ctx: SourceContext =
+ ActorContext::for_test(&universe, source_mailbox, observable_state_tx);
+
+ coordinator
+ .initialize(&doc_processor_mailbox, &ctx)
+ .await
+ .unwrap();
+
+ coordinator
+ .emit_batches(&doc_processor_mailbox, &ctx)
+ .await
+ .unwrap();
+
+ for (uri, ack_id) in messages {
+ queue.send_message(uri.to_string(), ack_id);
+ }
+
+ // Need 3 iterations for each msg to emit the first batch (receive,
+ // start, emit), assuming the `QueueReceiver` doesn't chunk the receive
+ // future.
+ for _ in 0..(messages.len() * 4) {
+ coordinator
+ .emit_batches(&doc_processor_mailbox, &ctx)
+ .await
+ .unwrap();
+ }
+
+ let batches = doc_processor_inbox
+ .drain_for_test()
+ .into_iter()
+ .flat_map(|box_any| box_any.downcast::().ok())
+ .map(|box_raw_doc_batch| *box_raw_doc_batch)
+ .collect::>();
+ universe.assert_quit().await;
+ batches
+ }
+
+ #[tokio::test]
+ async fn test_process_empty_queue() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state);
+ let batches = process_messages(&mut coordinator, queue, &[]).await;
+ assert_eq!(batches.len(), 0);
+ }
+
+ #[tokio::test]
+ async fn test_process_one_small_message() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state.clone());
+ let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id();
+ let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id")]).await;
+ assert_eq!(batches.len(), 1);
+ assert_eq!(batches[0].docs.len(), 10);
+ assert!(coordinator.local_state.is_awaiting_commit(&partition_id));
+ }
+
+ #[tokio::test]
+ async fn test_process_one_big_message() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state);
+ let lines = BATCH_NUM_BYTES_LIMIT as usize / DUMMY_DOC.len() + 1;
+ let (dummy_doc_file, _) = generate_dummy_doc_file(true, lines).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id")]).await;
+ assert_eq!(batches.len(), 2);
+ assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), lines);
+ }
+
+ #[tokio::test]
+ async fn test_process_two_messages_different_compression() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state);
+ let (dummy_doc_file_1, _) = generate_dummy_doc_file(false, 10).await;
+ let test_uri_1 = Uri::from_str(dummy_doc_file_1.path().to_str().unwrap()).unwrap();
+ let (dummy_doc_file_2, _) = generate_dummy_doc_file(true, 10).await;
+ let test_uri_2 = Uri::from_str(dummy_doc_file_2.path().to_str().unwrap()).unwrap();
+ let batches = process_messages(
+ &mut coordinator,
+ queue,
+ &[(&test_uri_1, "ack-id-1"), (&test_uri_2, "ack-id-2")],
+ )
+ .await;
+ // could be generated in 1 or 2 batches, it doesn't matter
+ assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), 20);
+ }
+
+ #[tokio::test]
+ async fn test_process_local_duplicate_message() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state);
+ let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ let batches = process_messages(
+ &mut coordinator,
+ queue,
+ &[(&test_uri, "ack-id-1"), (&test_uri, "ack-id-2")],
+ )
+ .await;
+ assert_eq!(batches.len(), 1);
+ assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), 10);
+ }
+
+ #[tokio::test]
+ async fn test_process_shared_complete_message() {
+ let (dummy_doc_file, file_size) = generate_dummy_doc_file(false, 10).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id();
+
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests(
+ "test-index",
+ &[(
+ partition_id.clone(),
+ ("existing_token".to_string(), Position::eof(file_size)),
+ )],
+ );
+ let mut coordinator = setup_coordinator(queue.clone(), shared_state.clone());
+
+ assert!(!coordinator.local_state.is_tracked(&partition_id));
+ let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id-1")]).await;
+ assert_eq!(batches.len(), 0);
+ assert!(coordinator.local_state.is_completed(&partition_id));
+ }
+
+ #[tokio::test]
+ async fn test_process_multiple_coordinator() {
+ let queue = Arc::new(MemoryQueueForTests::new());
+ let shared_state = shared_state_for_tests("test-index", Default::default());
+ let mut proc_1 = setup_coordinator(queue.clone(), shared_state.clone());
+ let mut proc_2 = setup_coordinator(queue.clone(), shared_state.clone());
+ let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await;
+ let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap();
+ let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id();
+
+ let batches_1 = process_messages(&mut proc_1, queue.clone(), &[(&test_uri, "ack1")]).await;
+ let batches_2 = process_messages(&mut proc_2, queue, &[(&test_uri, "ack2")]).await;
+
+ assert_eq!(batches_1.len(), 1);
+ assert_eq!(batches_1[0].docs.len(), 10);
+ assert!(proc_1.local_state.is_awaiting_commit(&partition_id));
+ // proc_2 doesn't know for sure what is happening with the message
+ // (proc_1 might have crashed), so it just acquires it and takes over
+ // processing
+ //
+ // TODO: this test should fail once we implement the grace
+ // period before a partition can be re-acquired
+ assert_eq!(batches_2.len(), 1);
+ assert_eq!(batches_2[0].docs.len(), 10);
+ assert!(proc_2.local_state.is_awaiting_commit(&partition_id));
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs b/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs
new file mode 100644
index 00000000000..8a215f66710
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs
@@ -0,0 +1,130 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::future::BoxFuture;
+
+use super::message::RawMessage;
+use super::Queue;
+
+type ReceiveResult = anyhow::Result>;
+
+/// A statefull wrapper around a `Queue` that chunks the slow `receive()` call
+/// into shorter iterations. This enables yielding back to the actor system
+/// without compromising on queue poll durations. Without this, an actor that
+/// tries to receive messages from a `Queue` will be blocked for multiple seconds
+/// before being able to process new mailbox messages (or shutting down).
+pub struct QueueReceiver {
+ queue: Arc,
+ receive: Option>,
+ iteration: Duration,
+}
+
+impl QueueReceiver {
+ pub fn new(queue: Arc, iteration: Duration) -> Self {
+ Self {
+ queue,
+ receive: None,
+ iteration,
+ }
+ }
+
+ pub async fn receive(
+ &mut self,
+ max_messages: usize,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result> {
+ if self.receive.is_none() {
+ self.receive = Some(self.queue.clone().receive(max_messages, suggested_deadline));
+ }
+ tokio::select! {
+ res = self.receive.as_mut().unwrap() => {
+ self.receive = None;
+ res
+ }
+ _ = tokio::time::sleep(self.iteration) => {
+ Ok(Vec::new())
+ }
+
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::time::{Duration, Instant};
+
+ use anyhow::bail;
+ use async_trait::async_trait;
+
+ use super::*;
+
+ #[derive(Clone, Debug)]
+ struct SleepyQueue {
+ receive_sleep: Duration,
+ }
+
+ #[async_trait]
+ impl Queue for SleepyQueue {
+ async fn receive(
+ self: Arc,
+ _max_messages: usize,
+ _suggested_deadline: Duration,
+ ) -> anyhow::Result> {
+ tokio::time::sleep(self.receive_sleep).await;
+ bail!("Waking up from my nap")
+ }
+
+ async fn acknowledge(&self, _ack_ids: &[String]) -> anyhow::Result<()> {
+ unimplemented!()
+ }
+
+ async fn modify_deadlines(
+ &self,
+ _ack_id: &str,
+ _suggested_deadline: Duration,
+ ) -> anyhow::Result {
+ unimplemented!()
+ }
+ }
+
+ #[tokio::test]
+ async fn test_queue_receiver_slow_receive() {
+ let queue = Arc::new(SleepyQueue {
+ receive_sleep: Duration::from_millis(100),
+ });
+ let mut receiver = QueueReceiver::new(queue, Duration::from_millis(20));
+ let mut iterations = 0;
+ while receiver.receive(1, Duration::from_secs(1)).await.is_ok() {
+ iterations += 1;
+ }
+ assert!(iterations >= 4);
+ }
+
+ #[tokio::test]
+ async fn test_queue_receiver_fast_receive() {
+ let queue = Arc::new(SleepyQueue {
+ receive_sleep: Duration::from_millis(10),
+ });
+ let mut receiver = QueueReceiver::new(queue, Duration::from_millis(50));
+ assert!(receiver.receive(1, Duration::from_secs(1)).await.is_err());
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs b/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs
new file mode 100644
index 00000000000..b0361b69adf
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs
@@ -0,0 +1,134 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::collections::{BTreeMap, BTreeSet, VecDeque};
+use std::time::Duration;
+
+use anyhow::bail;
+use quickwit_metastore::checkpoint::PartitionId;
+
+use super::message::{InProgressMessage, ReadyMessage};
+
+/// Tracks the state of the queue messages that are known to the owning indexing
+/// pipeline.
+///
+/// Messages first land in the `ready_for_read` queue. They are then moved to
+/// `read_in_progress` to track the reader's progress. Once the reader reaches
+/// EOF, the message is transitioned as `awaiting_commit`. Once the message is
+/// known to be fully indexed and committed (e.g after receiving the
+/// `suggest_truncate` call), it is moved to `completed`.
+#[derive(Default)]
+pub struct QueueLocalState {
+ /// Messages that were received from the queue and are ready to be read
+ ready_for_read: VecDeque,
+ /// Message that is currently being read and sent to the `DocProcessor`
+ read_in_progress: Option,
+ /// Partitions that were read and are still being indexed, with their
+ /// associated ack_id
+ awaiting_commit: BTreeMap,
+ /// Partitions that were fully indexed and committed
+ completed: BTreeSet,
+}
+
+impl QueueLocalState {
+ pub fn is_ready_for_read(&self, partition_id: &PartitionId) -> bool {
+ self.ready_for_read
+ .iter()
+ .any(|msg| &msg.partition_id() == partition_id)
+ }
+
+ pub fn is_read_in_progress(&self, partition_id: &PartitionId) -> bool {
+ self.read_in_progress
+ .as_ref()
+ .map_or(false, |msg| &msg.partition_id == partition_id)
+ }
+
+ pub fn is_awaiting_commit(&self, partition_id: &PartitionId) -> bool {
+ self.awaiting_commit.contains_key(partition_id)
+ }
+
+ pub fn is_completed(&self, partition_id: &PartitionId) -> bool {
+ self.completed.contains(partition_id)
+ }
+
+ pub fn is_tracked(&self, partition_id: &PartitionId) -> bool {
+ self.is_ready_for_read(partition_id)
+ || self.is_read_in_progress(partition_id)
+ || self.is_awaiting_commit(partition_id)
+ || self.is_completed(partition_id)
+ }
+
+ pub fn set_ready_for_read(&mut self, ready_messages: Vec) {
+ for message in ready_messages {
+ self.ready_for_read.push_back(message)
+ }
+ }
+
+ pub fn get_ready_for_read(&mut self) -> Option {
+ while let Some(msg) = self.ready_for_read.pop_front() {
+ // don't return messages for which we didn't manage to extend the
+ // visibility, they will pop up in the queue again anyway
+ if !msg.visibility_handle.extension_failed() {
+ return Some(msg);
+ }
+ }
+ None
+ }
+
+ pub fn read_in_progress_mut(&mut self) -> Option<&mut InProgressMessage> {
+ self.read_in_progress.as_mut()
+ }
+
+ pub async fn drop_currently_read(
+ &mut self,
+ deadline_for_last_extension: Duration,
+ ) -> anyhow::Result<()> {
+ if let Some(in_progress) = self.read_in_progress.take() {
+ self.awaiting_commit.insert(
+ in_progress.partition_id.clone(),
+ in_progress.visibility_handle.ack_id().to_string(),
+ );
+ in_progress
+ .visibility_handle
+ .request_last_extension(deadline_for_last_extension)
+ .await?;
+ }
+ Ok(())
+ }
+
+ /// Tries to set the message that is currently being read. Returns an error
+ /// if there is already a message being read.
+ pub fn set_currently_read(
+ &mut self,
+ in_progress: Option,
+ ) -> anyhow::Result<()> {
+ if self.read_in_progress.is_some() {
+ bail!("trying to replace in progress message");
+ }
+ self.read_in_progress = in_progress;
+ Ok(())
+ }
+
+ /// Returns the ack_id if that message was awaiting_commit
+ pub fn mark_completed(&mut self, partition_id: PartitionId) -> Option {
+ let ack_id_opt = self.awaiting_commit.remove(&partition_id);
+ self.completed.insert(partition_id);
+ ack_id_opt
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs b/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs
new file mode 100644
index 00000000000..627b7587de4
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs
@@ -0,0 +1,233 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::collections::{BTreeMap, VecDeque};
+use std::fmt;
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use anyhow::bail;
+use async_trait::async_trait;
+use quickwit_storage::OwnedBytes;
+use ulid::Ulid;
+
+use super::message::{MessageMetadata, RawMessage};
+use super::Queue;
+
+#[derive(Default)]
+struct InnerState {
+ in_queue: VecDeque,
+ in_flight: BTreeMap,
+ acked: Vec,
+}
+
+impl fmt::Debug for InnerState {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("Queue")
+ .field("in_queue_count", &self.in_queue.len())
+ .field("in_flight_count", &self.in_flight.len())
+ .field("acked_count", &self.acked.len())
+ .finish()
+ }
+}
+
+/// A simple in-memory queue
+#[derive(Clone, Debug)]
+pub struct MemoryQueueForTests {
+ inner_state: Arc>,
+ receive_sleep: Duration,
+}
+
+impl MemoryQueueForTests {
+ pub fn new() -> Self {
+ let inner_state = Arc::new(Mutex::new(InnerState::default()));
+ let inner_weak = Arc::downgrade(&inner_state);
+ tokio::spawn(async move {
+ loop {
+ if let Some(inner_state) = inner_weak.upgrade() {
+ let mut inner_state = inner_state.lock().unwrap();
+ let mut expired = Vec::new();
+ for (ack_id, msg) in inner_state.in_flight.iter() {
+ if msg.metadata.initial_deadline < Instant::now() {
+ expired.push(ack_id.clone());
+ }
+ }
+ for ack_id in expired {
+ let msg = inner_state.in_flight.remove(&ack_id).unwrap();
+ inner_state.in_queue.push_back(msg);
+ }
+ } else {
+ break;
+ }
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ }
+ });
+ MemoryQueueForTests {
+ inner_state: Arc::new(Mutex::new(InnerState::default())),
+ receive_sleep: Duration::from_millis(50),
+ }
+ }
+
+ pub fn send_message(&self, payload: String, ack_id: &str) {
+ let message = RawMessage {
+ payload: OwnedBytes::new(payload.into_bytes()),
+ metadata: MessageMetadata {
+ ack_id: ack_id.to_string(),
+ delivery_attempts: 0,
+ initial_deadline: Instant::now(),
+ message_id: Ulid::new().to_string(),
+ },
+ };
+ self.inner_state.lock().unwrap().in_queue.push_back(message);
+ }
+
+ /// Returns the next visibility deadline for the message if it is in flight
+ pub fn next_visibility_deadline(&self, ack_id: &str) -> Option {
+ let inner_state = self.inner_state.lock().unwrap();
+ inner_state
+ .in_flight
+ .get(ack_id)
+ .map(|msg| msg.metadata.initial_deadline)
+ }
+}
+
+#[async_trait]
+impl Queue for MemoryQueueForTests {
+ async fn receive(
+ self: Arc,
+ max_messages: usize,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result> {
+ {
+ let mut inner_state = self.inner_state.lock().unwrap();
+ let mut response = Vec::new();
+ while let Some(mut msg) = inner_state.in_queue.pop_front() {
+ msg.metadata.delivery_attempts += 1;
+ msg.metadata.initial_deadline = Instant::now() + suggested_deadline;
+ let msg_cloned = RawMessage {
+ payload: msg.payload.clone(),
+ metadata: msg.metadata.clone(),
+ };
+ inner_state
+ .in_flight
+ .insert(msg.metadata.ack_id.clone(), msg_cloned);
+ response.push(msg);
+ if response.len() >= max_messages {
+ break;
+ }
+ }
+ if !response.is_empty() {
+ return Ok(response);
+ }
+ }
+ // `sleep` to avoid using all the CPU when called in a loop
+ tokio::time::sleep(self.receive_sleep).await;
+
+ Ok(vec![])
+ }
+
+ async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()> {
+ let mut inner_state = self.inner_state.lock().unwrap();
+ for ack_id in ack_ids {
+ if let Some(msg) = inner_state.in_flight.remove(ack_id) {
+ inner_state.acked.push(msg);
+ }
+ }
+ Ok(())
+ }
+
+ async fn modify_deadlines(
+ &self,
+ ack_id: &str,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result {
+ let mut inner_state = self.inner_state.lock().unwrap();
+ let in_flight = inner_state.in_flight.get_mut(ack_id);
+ if let Some(msg) = in_flight {
+ msg.metadata.initial_deadline = Instant::now() + suggested_deadline;
+ } else {
+ bail!("ack_id {} not found in in-flight", ack_id);
+ }
+ return Ok(Instant::now() + suggested_deadline);
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn prefilled_queue(nb_message: usize) -> Arc {
+ let memory_queue = MemoryQueueForTests::new();
+ for i in 0..nb_message {
+ let payload = format!("Test message {}", i);
+ let ack_id = i.to_string();
+ memory_queue.send_message(payload.clone(), &ack_id);
+ }
+ Arc::new(memory_queue)
+ }
+
+ #[tokio::test]
+ async fn test_receive_1_by_1() {
+ let memory_queue = prefilled_queue(2);
+ for i in 0..2 {
+ let messages = memory_queue
+ .clone()
+ .receive(1, Duration::from_secs(5))
+ .await
+ .unwrap();
+ assert_eq!(messages.len(), 1);
+ let message = &messages[0];
+ let exp_payload = format!("Test message {}", i);
+ let exp_ack_id = i.to_string();
+ assert_eq!(message.payload.as_ref(), exp_payload.as_bytes());
+ assert_eq!(message.metadata.ack_id, exp_ack_id);
+ }
+ }
+
+ #[tokio::test]
+ async fn test_receive_2_by_2() {
+ let memory_queue = prefilled_queue(2);
+ let messages = memory_queue
+ .receive(2, Duration::from_secs(5))
+ .await
+ .unwrap();
+ assert_eq!(messages.len(), 2);
+ for (i, message) in messages.iter().enumerate() {
+ let exp_payload = format!("Test message {}", i);
+ let exp_ack_id = i.to_string();
+ assert_eq!(message.payload.as_ref(), exp_payload.as_bytes());
+ assert_eq!(message.metadata.ack_id, exp_ack_id);
+ }
+ }
+
+ #[tokio::test]
+ async fn test_receive_early_if_only_1() {
+ let memory_queue = prefilled_queue(1);
+ let messages = memory_queue
+ .receive(2, Duration::from_secs(5))
+ .await
+ .unwrap();
+ assert_eq!(messages.len(), 1);
+ let message = &messages[0];
+ let exp_payload = "Test message 0".to_string();
+ let exp_ack_id = "0";
+ assert_eq!(message.payload.as_ref(), exp_payload.as_bytes());
+ assert_eq!(message.metadata.ack_id, exp_ack_id);
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/message.rs b/quickwit/quickwit-indexing/src/source/queue_sources/message.rs
new file mode 100644
index 00000000000..c0ad5e2b235
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/message.rs
@@ -0,0 +1,352 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use core::fmt;
+use std::io::read_to_string;
+use std::str::FromStr;
+use std::time::Instant;
+
+use anyhow::Context;
+use quickwit_common::rate_limited_warn;
+use quickwit_common::uri::Uri;
+use quickwit_metastore::checkpoint::PartitionId;
+use quickwit_proto::types::Position;
+use quickwit_storage::{OwnedBytes, StorageResolver};
+use serde_json::Value;
+use thiserror::Error;
+use tracing::info;
+
+use super::visibility::VisibilityTaskHandle;
+use crate::source::doc_file_reader::ObjectUriBatchReader;
+
+#[derive(Debug, Clone, Copy)]
+pub enum MessageType {
+ S3Notification,
+ // GcsNotification,
+ RawUri,
+ // RawData,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MessageMetadata {
+ /// The handle that should be used to acknowledge the message or change its visibility deadline
+ pub ack_id: String,
+
+ /// The unique message id assigned by the queue
+ pub message_id: String,
+
+ /// The approximate number of times the message was delivered. 1 means it is
+ /// the first time this message is being delivered.
+ pub delivery_attempts: usize,
+
+ /// The first deadline when the message is received. It can be extended later using the ack_id.
+ pub initial_deadline: Instant,
+}
+
+/// The raw messages as received from the queue abstraction
+pub struct RawMessage {
+ pub metadata: MessageMetadata,
+ pub payload: OwnedBytes,
+}
+
+impl fmt::Debug for RawMessage {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("RawMessage")
+ .field("metadata", &self.metadata)
+ .field("payload", &"")
+ .finish()
+ }
+}
+
+#[derive(Error, Debug)]
+pub enum PreProcessingError {
+ #[error("message can be acknowledged without processing")]
+ Discardable { ack_id: String },
+ #[error("unexpected message format: {0}")]
+ UnexpectedFormat(#[from] anyhow::Error),
+}
+
+impl RawMessage {
+ pub fn pre_process(
+ self,
+ message_type: MessageType,
+ ) -> Result {
+ let payload = match message_type {
+ MessageType::S3Notification => PreProcessedPayload::ObjectUri(
+ uri_from_s3_notification(&self.payload, &self.metadata.ack_id)?,
+ ),
+ MessageType::RawUri => {
+ let payload_str = read_to_string(self.payload).context("failed to read payload")?;
+ PreProcessedPayload::ObjectUri(Uri::from_str(&payload_str)?)
+ }
+ };
+ Ok(PreProcessedMessage {
+ metadata: self.metadata,
+ payload,
+ })
+ }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum PreProcessedPayload {
+ /// The message contains an object URI
+ ObjectUri(Uri),
+ // /// The message contains the raw JSON data
+ // RawData(OwnedBytes),
+}
+
+impl PreProcessedPayload {
+ pub fn partition_id(&self) -> PartitionId {
+ match &self {
+ Self::ObjectUri(uri) => PartitionId::from(uri.as_str()),
+ }
+ }
+}
+
+/// A message that went through the minimal transformation to discover its
+/// partition id. Indeed, the message might be discarded if the partition was
+/// already processed, so it's better to avoid doing unnecessary work at this
+/// stage.
+#[derive(Debug, PartialEq, Eq)]
+pub struct PreProcessedMessage {
+ pub metadata: MessageMetadata,
+ pub payload: PreProcessedPayload,
+}
+
+impl PreProcessedMessage {
+ pub fn partition_id(&self) -> PartitionId {
+ self.payload.partition_id()
+ }
+}
+
+fn uri_from_s3_notification(message: &[u8], ack_id: &str) -> Result {
+ let value: Value = serde_json::from_slice(message).context("invalid JSON message")?;
+ if matches!(value["Event"].as_str(), Some("s3:TestEvent")) {
+ info!("discarding S3 test event");
+ return Err(PreProcessingError::Discardable {
+ ack_id: ack_id.to_string(),
+ });
+ }
+ let event_name = value["Records"][0]["eventName"]
+ .as_str()
+ .context("invalid S3 notification: Records[0].eventName not found")?;
+ if !event_name.starts_with("ObjectCreated:") {
+ rate_limited_warn!(
+ limit_per_min = 5,
+ event = event_name,
+ "only s3:ObjectCreated:* events are supported"
+ );
+ return Err(PreProcessingError::Discardable {
+ ack_id: ack_id.to_string(),
+ });
+ }
+ let key = value["Records"][0]["s3"]["object"]["key"]
+ .as_str()
+ .context("invalid S3 notification: Records[0].s3.object.key not found")?;
+ let bucket = value["Records"][0]["s3"]["bucket"]["name"]
+ .as_str()
+ .context("invalid S3 notification: Records[0].s3.bucket.name not found")?;
+ Uri::from_str(&format!("s3://{}/{}", bucket, key)).map_err(|e| e.into())
+}
+
+/// A message for which we know as much of the global processing status as
+/// possible and that is now ready to be processed.
+pub struct ReadyMessage {
+ pub position: Position,
+ pub content: PreProcessedMessage,
+ pub visibility_handle: VisibilityTaskHandle,
+}
+
+impl ReadyMessage {
+ pub async fn start_processing(
+ self,
+ storage_resolver: &StorageResolver,
+ ) -> anyhow::Result> {
+ let partition_id = self.partition_id();
+ match self.content.payload {
+ PreProcessedPayload::ObjectUri(uri) => {
+ let batch_reader = ObjectUriBatchReader::try_new(
+ storage_resolver,
+ partition_id.clone(),
+ &uri,
+ self.position,
+ )
+ .await?;
+ if batch_reader.is_eof() {
+ Ok(None)
+ } else {
+ Ok(Some(InProgressMessage {
+ batch_reader,
+ partition_id,
+ visibility_handle: self.visibility_handle,
+ }))
+ }
+ }
+ }
+ }
+
+ pub fn partition_id(&self) -> PartitionId {
+ self.content.partition_id()
+ }
+}
+
+/// A message that is actively being read
+pub struct InProgressMessage {
+ pub partition_id: PartitionId,
+ pub visibility_handle: VisibilityTaskHandle,
+ pub batch_reader: ObjectUriBatchReader,
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_uri_from_s3_notification_valid() {
+ let test_message = r#"
+ {
+ "Records": [
+ {
+ "eventVersion": "2.1",
+ "eventSource": "aws:s3",
+ "awsRegion": "us-west-2",
+ "eventTime": "2021-05-22T09:22:41.789Z",
+ "eventName": "ObjectCreated:Put",
+ "userIdentity": {
+ "principalId": "AWS:AIDAJDPLRKLG7UEXAMPLE"
+ },
+ "requestParameters": {
+ "sourceIPAddress": "127.0.0.1"
+ },
+ "responseElements": {
+ "x-amz-request-id": "C3D13FE58DE4C810",
+ "x-amz-id-2": "FMyUVURIx7Zv2cPi/IZb9Fk1/U4QfTaVK5fahHPj/"
+ },
+ "s3": {
+ "s3SchemaVersion": "1.0",
+ "configurationId": "testConfigRule",
+ "bucket": {
+ "name": "mybucket",
+ "ownerIdentity": {
+ "principalId": "A3NL1KOZZKExample"
+ },
+ "arn": "arn:aws:s3:::mybucket"
+ },
+ "object": {
+ "key": "logs.json",
+ "size": 1024,
+ "eTag": "d41d8cd98f00b204e9800998ecf8427e",
+ "versionId": "096fKKXTRTtl3on89fVO.nfljtsv6qko",
+ "sequencer": "0055AED6DCD90281E5"
+ }
+ }
+ }
+ ]
+ }"#;
+ let actual_uri = uri_from_s3_notification(test_message.as_bytes(), "myackid").unwrap();
+ let expected_uri = Uri::from_str("s3://mybucket/logs.json").unwrap();
+ assert_eq!(actual_uri, expected_uri);
+ }
+
+ #[test]
+ fn test_uri_from_s3_notification_invalid() {
+ let invalid_message = r#"{
+ "Records": [
+ {
+ "s3": {
+ "object": {
+ "key": "test_key"
+ }
+ }
+ }
+ ]
+ }"#;
+ let result =
+ uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid");
+ assert!(matches!(
+ result,
+ Err(PreProcessingError::UnexpectedFormat(_))
+ ));
+ }
+
+ #[test]
+ fn test_uri_from_s3_bad_event_type() {
+ let invalid_message = r#"{
+ "Records": [
+ {
+ "eventVersion": "2.1",
+ "eventSource": "aws:s3",
+ "awsRegion": "us-east-1",
+ "eventTime": "2024-07-29T12:47:14.577Z",
+ "eventName": "ObjectRemoved:Delete",
+ "userIdentity": {
+ "principalId": "AWS:ARGHGOHSDGOKGHOGHMCC4:user"
+ },
+ "requestParameters": {
+ "sourceIPAddress": "1.1.1.1"
+ },
+ "responseElements": {
+ "x-amz-request-id": "GHGSH",
+ "x-amz-id-2": "gndflghndflhmnrflsh+gLLKU6X0PvD6ANdVY1+/hspflhjladgfkelagfkndl"
+ },
+ "s3": {
+ "s3SchemaVersion": "1.0",
+ "configurationId": "hello",
+ "bucket": {
+ "name": "mybucket",
+ "ownerIdentity": {
+ "principalId": "KMGP12GHKKH"
+ },
+ "arn": "arn:aws:s3:::mybucket"
+ },
+ "object": {
+ "key": "my_deleted_file",
+ "sequencer": "GKHOFLGKHSALFK0"
+ }
+ }
+ }
+ ]
+ }"#;
+ let result =
+ uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid");
+ assert!(matches!(
+ result,
+ Err(PreProcessingError::Discardable { .. })
+ ));
+ }
+
+ #[test]
+ fn test_uri_from_s3_notification_discardable() {
+ let invalid_message = r#"{
+ "Service":"Amazon S3",
+ "Event":"s3:TestEvent",
+ "Time":"2014-10-13T15:57:02.089Z",
+ "Bucket":"bucketname",
+ "RequestId":"5582815E1AEA5ADF",
+ "HostId":"8cLeGAmw098X5cv4Zkwcmo8vvZa3eH3eKxsPzbB9wrR+YstdA6Knx4Ip8EXAMPLE"
+ }"#;
+ let result =
+ uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid");
+ if let Err(PreProcessingError::Discardable { ack_id }) = result {
+ assert_eq!(ack_id, "myackid");
+ } else {
+ panic!("Expected skippable error");
+ }
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs b/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs
new file mode 100644
index 00000000000..dc94a53d7f8
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs
@@ -0,0 +1,89 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+pub mod coordinator;
+mod helpers;
+mod local_state;
+#[cfg(test)]
+mod memory_queue;
+mod message;
+mod shared_state;
+#[cfg(feature = "sqs")]
+pub mod sqs_queue;
+mod visibility;
+
+use std::fmt;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use async_trait::async_trait;
+use message::RawMessage;
+
+/// The queue abstraction is based on the AWS SQS and Google Pubsub APIs. The
+/// only requirement of the underlying implementation is that messages exposed
+/// to a given consumer are hidden to other consumers for a configurable period
+/// of time. Retries are handled by the implementation because queues might
+/// behave differently (throttling, deduplication...).
+#[async_trait]
+pub trait Queue: fmt::Debug + Send + Sync + 'static {
+ /// Polls the queue to receive messages.
+ ///
+ /// The implementation is in charge of choosing the wait strategy when there
+ /// are no messages in the queue. It will typically use long polling to do
+ /// this efficiently. On the other hand, when there is a message available
+ /// in the queue, it should be returned as quickly as possible, regardless
+ /// of the `max_messages` parameter. The `max_messages` paramater should
+ /// always be clamped by the implementation to not violate the maximum value
+ /// supported by the backing queue (e.g 10 messages for AWS SQS).
+ ///
+ /// As soon as the message is received, the caller is responsible for
+ /// maintaining the message visibility in a timely fashion. Failing to do so
+ /// implies that duplicates will be received by other indexing pipelines,
+ /// thus increasing competition for the commit lock.
+ async fn receive(
+ // `Arc` to make the resulting future `'static` and thus easily
+ // wrappable by the `QueueReceiver`
+ self: Arc,
+ max_messages: usize,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result>;
+
+ /// Tries to acknowledge (delete) the messages.
+ ///
+ /// The call returns `Ok(())` if at the message level:
+ /// - the acknowledgement failed due to a transient failure
+ /// - the message was already acknowledged
+ /// - the message was not acknowledged in time and is back to the queue
+ ///
+ /// If an empty list of ack_ids is provided, the call should be a no-op.
+ async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()>;
+
+ /// Modifies the visibility deadline of the messages.
+ ///
+ /// We try to set the initial visibility large enough to avoid having to
+ /// call this too often. The implementation can retry as long as desired,
+ /// it's the caller's responsibility to cancel the future if the deadline is
+ /// getting to close to the expiration. The returned `Instant` is a
+ /// conservative estimate of the new deadline expiration time.
+ async fn modify_deadlines(
+ &self,
+ ack_id: &str,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result;
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs b/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs
new file mode 100644
index 00000000000..cdd0ade05b3
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs
@@ -0,0 +1,371 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::collections::BTreeMap;
+
+use anyhow::{bail, Context};
+use quickwit_metastore::checkpoint::PartitionId;
+use quickwit_proto::metastore::{
+ AcquireShardsRequest, MetastoreService, MetastoreServiceClient, OpenShardSubrequest,
+ OpenShardsRequest,
+};
+use quickwit_proto::types::{DocMappingUid, IndexUid, Position, ShardId};
+use tracing::info;
+
+use super::message::PreProcessedMessage;
+
+#[derive(Clone)]
+pub struct QueueSharedState {
+ pub metastore: MetastoreServiceClient,
+ pub index_uid: IndexUid,
+ pub source_id: String,
+}
+
+impl QueueSharedState {
+ /// Tries to acquire the ownership for the provided messages from the global
+ /// shared context. For each partition id, if the ownership was successfully
+ /// acquired or the partition was already successfully indexed, the position
+ /// is returned along with the partition id, otherwise the partition id is
+ /// dropped.
+ async fn acquire_partitions(
+ &self,
+ publish_token: &str,
+ partitions: Vec,
+ ) -> anyhow::Result> {
+ let open_shard_subrequests = partitions
+ .iter()
+ .enumerate()
+ .map(|(idx, partition_id)| OpenShardSubrequest {
+ subrequest_id: idx as u32,
+ index_uid: Some(self.index_uid.clone()),
+ source_id: self.source_id.clone(),
+ leader_id: String::new(),
+ follower_id: None,
+ shard_id: Some(ShardId::from(partition_id.as_str())),
+ doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: Some(publish_token.to_string()),
+ })
+ .collect();
+
+ let open_shard_resp = self
+ .metastore
+ .open_shards(OpenShardsRequest {
+ subrequests: open_shard_subrequests,
+ })
+ .await?;
+
+ let mut shards = Vec::new();
+ let mut re_acquired_shards = Vec::new();
+ for sub in open_shard_resp.subresponses {
+ // we could also just cast the shard_id back to a partition_id
+ let partition_id = partitions[sub.subrequest_id as usize].clone();
+ let shard = sub.open_shard();
+ let position = shard.publish_position_inclusive.clone().unwrap_or_default();
+ let is_owned = sub.open_shard().publish_token.as_deref() == Some(publish_token);
+ if position.is_eof() || (is_owned && position.is_beginning()) {
+ shards.push((partition_id, position));
+ } else if !is_owned {
+ // TODO: Add logic to only re-acquire shards that have a token that is not
+ // the local token when they haven't been updated recently
+ info!(previous_token = shard.publish_token, "shard re-acquired");
+ re_acquired_shards.push(shard.shard_id().clone());
+ } else if is_owned && !position.is_beginning() {
+ bail!("Partition is owned by this indexing pipeline but is not at the beginning. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues.")
+ }
+ }
+
+ if re_acquired_shards.is_empty() {
+ return Ok(shards);
+ }
+
+ // re-acquire shards that have a token that is not the local token
+ let acquire_shard_resp = self
+ .metastore
+ .acquire_shards(AcquireShardsRequest {
+ index_uid: Some(self.index_uid.clone()),
+ source_id: self.source_id.clone(),
+ shard_ids: re_acquired_shards,
+ publish_token: publish_token.to_string(),
+ })
+ .await
+ .unwrap();
+ for shard in acquire_shard_resp.acquired_shards {
+ let partition_id = PartitionId::from(shard.shard_id().as_str());
+ let position = shard.publish_position_inclusive.unwrap_or_default();
+ shards.push((partition_id, position));
+ }
+
+ Ok(shards)
+ }
+}
+
+/// Acquires shards from the shared state for the provided list of messages and
+/// maps results to that same list
+pub async fn checkpoint_messages(
+ shared_state: &QueueSharedState,
+ publish_token: &str,
+ messages: Vec,
+) -> anyhow::Result> {
+ let mut message_map =
+ BTreeMap::from_iter(messages.into_iter().map(|msg| (msg.partition_id(), msg)));
+ let partition_ids = message_map.keys().cloned().collect();
+
+ let shards = shared_state
+ .acquire_partitions(publish_token, partition_ids)
+ .await?;
+
+ shards
+ .into_iter()
+ .map(|(partition_id, position)| {
+ let content = message_map.remove(&partition_id).context("Unexpected partition ID. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues.")?;
+ Ok((
+ content,
+ position,
+ ))
+ })
+ .collect::>()
+}
+
+#[cfg(test)]
+pub mod shared_state_for_tests {
+ use std::sync::{Arc, Mutex};
+
+ use quickwit_proto::ingest::{Shard, ShardState};
+ use quickwit_proto::metastore::{
+ AcquireShardsResponse, MockMetastoreService, OpenShardSubresponse, OpenShardsResponse,
+ };
+
+ use super::*;
+
+ pub(super) fn mock_metastore(
+ initial_state: &[(PartitionId, (String, Position))],
+ open_shard_times: Option,
+ acquire_times: Option,
+ ) -> MetastoreServiceClient {
+ let mut mock_metastore = MockMetastoreService::new();
+ let inner_state = Arc::new(Mutex::new(BTreeMap::from_iter(
+ initial_state.iter().cloned(),
+ )));
+ let inner_state_ref = Arc::clone(&inner_state);
+ let open_shards_expectation =
+ mock_metastore
+ .expect_open_shards()
+ .returning(move |request| {
+ let subresponses = request
+ .subrequests
+ .into_iter()
+ .map(|sub_req| {
+ let partition_id: PartitionId = sub_req.shard_id().to_string().into();
+ let (token, position) = inner_state_ref
+ .lock()
+ .unwrap()
+ .get(&partition_id)
+ .cloned()
+ .unwrap_or((sub_req.publish_token.unwrap(), Position::Beginning));
+ inner_state_ref
+ .lock()
+ .unwrap()
+ .insert(partition_id, (token.clone(), position.clone()));
+ OpenShardSubresponse {
+ subrequest_id: sub_req.subrequest_id,
+ open_shard: Some(Shard {
+ shard_id: sub_req.shard_id,
+ source_id: sub_req.source_id,
+ publish_token: Some(token),
+ index_uid: sub_req.index_uid,
+ follower_id: sub_req.follower_id,
+ leader_id: sub_req.leader_id,
+ doc_mapping_uid: sub_req.doc_mapping_uid,
+ publish_position_inclusive: Some(position),
+ shard_state: ShardState::Open as i32,
+ }),
+ }
+ })
+ .collect();
+ Ok(OpenShardsResponse { subresponses })
+ });
+ if let Some(times) = open_shard_times {
+ open_shards_expectation.times(times);
+ }
+ let acquire_shards_expectation = mock_metastore
+ .expect_acquire_shards()
+ // .times(acquire_times)
+ .returning(move |request| {
+ let acquired_shards = request
+ .shard_ids
+ .into_iter()
+ .map(|shard_id| {
+ let partition_id: PartitionId = shard_id.to_string().into();
+ let (existing_token, position) = inner_state
+ .lock()
+ .unwrap()
+ .get(&partition_id)
+ .cloned()
+ .expect("we should never try to acquire a shard that doesn't exist");
+ inner_state.lock().unwrap().insert(
+ partition_id,
+ (request.publish_token.clone(), position.clone()),
+ );
+ assert_ne!(existing_token, request.publish_token);
+ Shard {
+ shard_id: Some(shard_id),
+ source_id: "dummy".to_string(),
+ publish_token: Some(request.publish_token.clone()),
+ index_uid: None,
+ follower_id: None,
+ leader_id: "dummy".to_string(),
+ doc_mapping_uid: None,
+ publish_position_inclusive: Some(position),
+ shard_state: ShardState::Open as i32,
+ }
+ })
+ .collect();
+ Ok(AcquireShardsResponse { acquired_shards })
+ });
+ if let Some(times) = acquire_times {
+ acquire_shards_expectation.times(times);
+ }
+ MetastoreServiceClient::from_mock(mock_metastore)
+ }
+
+ pub fn shared_state_for_tests(
+ index_id: &str,
+ initial_state: &[(PartitionId, (String, Position))],
+ ) -> QueueSharedState {
+ let index_uid = IndexUid::new_with_random_ulid(index_id);
+ let metastore = mock_metastore(initial_state, None, None);
+ QueueSharedState {
+ metastore,
+ index_uid,
+ source_id: "test-queue-src".to_string(),
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::str::FromStr;
+ use std::time::Instant;
+ use std::vec;
+
+ use quickwit_common::uri::Uri;
+ use shared_state_for_tests::mock_metastore;
+
+ use super::*;
+ use crate::source::queue_sources::message::{MessageMetadata, PreProcessedPayload};
+
+ fn test_messages(message_number: usize) -> Vec {
+ (0..message_number)
+ .map(|i| PreProcessedMessage {
+ metadata: MessageMetadata {
+ ack_id: format!("ackid{}", i),
+ delivery_attempts: 0,
+ initial_deadline: Instant::now(),
+ message_id: format!("mid{}", i),
+ },
+ payload: PreProcessedPayload::ObjectUri(
+ Uri::from_str(&format!("s3://bucket/key{}", i)).unwrap(),
+ ),
+ })
+ .collect()
+ }
+
+ #[tokio::test]
+ async fn test_acquire_shards_with_completed() {
+ let index_id = "test-sqs-index";
+ let index_uid = IndexUid::new_with_random_ulid(index_id);
+ let init_state = &[("p1".into(), ("token2".to_string(), Position::eof(100usize)))];
+ let metastore = mock_metastore(init_state, Some(1), Some(0));
+
+ let shared_state = QueueSharedState {
+ metastore,
+ index_uid,
+ source_id: "test-sqs-source".to_string(),
+ };
+
+ let aquired = shared_state
+ .acquire_partitions("token1", vec!["p1".into(), "p2".into()])
+ .await
+ .unwrap();
+ assert!(aquired.contains(&("p1".into(), Position::eof(100usize))));
+ assert!(aquired.contains(&("p2".into(), Position::Beginning)));
+ }
+
+ #[tokio::test]
+ async fn test_re_acquire_shards() {
+ let index_id = "test-sqs-index";
+ let index_uid = IndexUid::new_with_random_ulid(index_id);
+ let init_state = &[(
+ "p1".into(),
+ ("token2".to_string(), Position::offset(100usize)),
+ )];
+ let metastore = mock_metastore(init_state, Some(1), Some(1));
+
+ let shared_state = QueueSharedState {
+ metastore,
+ index_uid,
+ source_id: "test-sqs-source".to_string(),
+ };
+
+ let aquired = shared_state
+ .acquire_partitions("token1", vec!["p1".into(), "p2".into()])
+ .await
+ .unwrap();
+ // TODO: this test should fail once we implement the grace
+ // period before a partition can be re-acquired
+ assert!(aquired.contains(&("p1".into(), Position::offset(100usize))));
+ assert!(aquired.contains(&("p2".into(), Position::Beginning)));
+ }
+
+ #[tokio::test]
+ async fn test_checkpoint_with_completed() {
+ let index_id = "test-sqs-index";
+ let index_uid = IndexUid::new_with_random_ulid(index_id);
+
+ let source_messages = test_messages(2);
+ let completed_partition_id = source_messages[0].partition_id();
+ let new_partition_id = source_messages[1].partition_id();
+
+ let init_state = &[(
+ completed_partition_id.clone(),
+ ("token2".to_string(), Position::eof(100usize)),
+ )];
+ let metastore = mock_metastore(init_state, Some(1), Some(0));
+ let shared_state = QueueSharedState {
+ metastore,
+ index_uid,
+ source_id: "test-sqs-source".to_string(),
+ };
+
+ let checkpointed_msg = checkpoint_messages(&shared_state, "token1", source_messages)
+ .await
+ .unwrap();
+ assert_eq!(checkpointed_msg.len(), 2);
+ let completed_msg = checkpointed_msg
+ .iter()
+ .find(|(msg, _)| msg.partition_id() == completed_partition_id)
+ .unwrap();
+ assert_eq!(completed_msg.1, Position::eof(100usize));
+ let new_msg = checkpointed_msg
+ .iter()
+ .find(|(msg, _)| msg.partition_id() == new_partition_id)
+ .unwrap();
+ assert_eq!(new_msg.1, Position::Beginning);
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs b/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs
new file mode 100644
index 00000000000..49d5923a9e1
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs
@@ -0,0 +1,468 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use anyhow::{bail, Context};
+use async_trait::async_trait;
+use aws_sdk_sqs::config::{BehaviorVersion, Builder, Region, SharedAsyncSleep};
+use aws_sdk_sqs::types::{DeleteMessageBatchRequestEntry, MessageSystemAttributeName};
+use aws_sdk_sqs::{Client, Config};
+use itertools::Itertools;
+use quickwit_aws::retry::{aws_retry, AwsRetryable};
+use quickwit_aws::{get_aws_config, DEFAULT_AWS_REGION};
+use quickwit_common::rate_limited_error;
+use quickwit_common::retry::RetryParams;
+use quickwit_storage::OwnedBytes;
+use regex::Regex;
+
+use super::message::MessageMetadata;
+use super::{Queue, RawMessage};
+
+#[derive(Debug)]
+pub struct SqsQueue {
+ sqs_client: Client,
+ queue_url: String,
+ receive_retries: RetryParams,
+ acknowledge_retries: RetryParams,
+ modify_deadline_retries: RetryParams,
+}
+
+impl SqsQueue {
+ pub async fn try_new(queue_url: String) -> anyhow::Result {
+ let sqs_client = get_sqs_client(&queue_url).await?;
+ Ok(SqsQueue {
+ sqs_client,
+ queue_url,
+ receive_retries: RetryParams::standard(),
+ // Acknowledgment is retried when the message is received again
+ acknowledge_retries: RetryParams::no_retries(),
+ // Retry aggressively to avoid loosing the ownership of the message
+ modify_deadline_retries: RetryParams::aggressive(),
+ })
+ }
+}
+
+#[async_trait]
+impl Queue for SqsQueue {
+ async fn receive(
+ self: Arc,
+ max_messages: usize,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result> {
+ // TODO: We estimate the message deadline using the start of the
+ // ReceiveMessage request. This might be overly pessimistic: the docs
+ // state that it starts when the message is returned.
+ let initial_deadline = Instant::now() + suggested_deadline;
+ let clamped_max_messages = std::cmp::min(max_messages, 10) as i32;
+ let receive_output = aws_retry(&self.receive_retries, || async {
+ self.sqs_client
+ .receive_message()
+ .queue_url(&self.queue_url)
+ .message_system_attribute_names(MessageSystemAttributeName::ApproximateReceiveCount)
+ .wait_time_seconds(20)
+ .set_max_number_of_messages(Some(clamped_max_messages))
+ .visibility_timeout(suggested_deadline.as_secs() as i32)
+ .send()
+ .await
+ })
+ .await?;
+
+ let received_messages = receive_output.messages.unwrap_or_default();
+ let mut resulting_raw_messages = Vec::with_capacity(received_messages.len());
+ for received_message in received_messages {
+ let delivery_attempts: usize = received_message
+ .attributes
+ .as_ref()
+ .and_then(|attrs| attrs.get(&MessageSystemAttributeName::ApproximateReceiveCount))
+ .and_then(|s| s.parse().ok())
+ .unwrap_or(0);
+ let ack_id = received_message
+ .receipt_handle
+ .context("missing receipt_handle in received message")?;
+ let message_id = received_message
+ .message_id
+ .context("missing message_id in received message")?;
+ let raw_message = RawMessage {
+ metadata: MessageMetadata {
+ ack_id,
+ message_id,
+ initial_deadline,
+ delivery_attempts,
+ },
+ payload: OwnedBytes::new(received_message.body.unwrap_or_default().into_bytes()),
+ };
+ resulting_raw_messages.push(raw_message);
+ }
+ Ok(resulting_raw_messages)
+ }
+
+ async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()> {
+ if ack_ids.is_empty() {
+ return Ok(());
+ }
+ let entry_batches: Vec> = ack_ids
+ .iter()
+ .dedup()
+ .enumerate()
+ .map(|(i, id)| {
+ DeleteMessageBatchRequestEntry::builder()
+ .id(i.to_string())
+ .receipt_handle(id.to_string())
+ .build()
+ .unwrap()
+ })
+ .chunks(10)
+ .into_iter()
+ .map(|chunk| chunk.collect())
+ .collect();
+
+ // TODO: parallelization
+ let mut batch_errors = Vec::new();
+ let mut message_errors = Vec::new();
+ for batch in entry_batches {
+ let res = aws_retry(&self.acknowledge_retries, || {
+ self.sqs_client
+ .delete_message_batch()
+ .queue_url(&self.queue_url)
+ .set_entries(Some(batch.clone()))
+ .send()
+ })
+ .await;
+ match res {
+ Ok(res) => {
+ message_errors.extend(res.failed.into_iter());
+ }
+ Err(err) => {
+ batch_errors.push(err);
+ }
+ }
+ }
+ if batch_errors.iter().any(|err| !err.is_retryable()) {
+ let fatal_error = batch_errors
+ .into_iter()
+ .find(|err| !err.is_retryable())
+ .unwrap();
+ bail!(fatal_error);
+ } else if !batch_errors.is_empty() {
+ rate_limited_error!(
+ limit_per_min = 10,
+ count = batch_errors.len(),
+ first_err = ?batch_errors.into_iter().next().unwrap(),
+ "failed to acknowledge some message batches",
+ );
+ }
+ // The documentation is unclear about these partial failures. We assume
+ // it is either:
+ // - a transient failure
+ // - the message is already acknowledged
+ // - the message is expired
+ if !message_errors.is_empty() {
+ rate_limited_error!(
+ limit_per_min = 10,
+ count = message_errors.len(),
+ first_err = ?message_errors.into_iter().next().unwrap(),
+ "failed to acknowledge individual messages",
+ );
+ }
+ Ok(())
+ }
+
+ async fn modify_deadlines(
+ &self,
+ ack_id: &str,
+ suggested_deadline: Duration,
+ ) -> anyhow::Result {
+ let visibility_timeout = std::cmp::min(suggested_deadline.as_secs() as i32, 43200);
+ let new_deadline = Instant::now() + suggested_deadline;
+ aws_retry(&self.modify_deadline_retries, || {
+ self.sqs_client
+ .change_message_visibility()
+ .queue_url(&self.queue_url)
+ .visibility_timeout(visibility_timeout)
+ .receipt_handle(ack_id)
+ .send()
+ })
+ .await?;
+ Ok(new_deadline)
+ }
+}
+
+async fn preconfigured_builder() -> anyhow::Result {
+ let aws_config = get_aws_config().await;
+
+ let mut sqs_config = Config::builder().behavior_version(BehaviorVersion::v2024_03_28());
+ sqs_config.set_retry_config(aws_config.retry_config().cloned());
+ sqs_config.set_credentials_provider(aws_config.credentials_provider());
+ sqs_config.set_http_client(aws_config.http_client());
+ sqs_config.set_timeout_config(aws_config.timeout_config().cloned());
+
+ if let Some(identity_cache) = aws_config.identity_cache() {
+ sqs_config.set_identity_cache(identity_cache);
+ }
+ sqs_config.set_sleep_impl(Some(SharedAsyncSleep::new(
+ quickwit_aws::TokioSleep::default(),
+ )));
+
+ Ok(sqs_config)
+}
+
+fn queue_url_region(queue_url: &str) -> Option {
+ let re = Regex::new(r"^https?://sqs\.(.*?)\.amazonaws\.com").unwrap();
+ let caps = re.captures(queue_url)?;
+ let region_str = caps.get(1)?.as_str();
+ Some(Region::new(region_str.to_string()))
+}
+
+fn queue_url_endpoint(queue_url: &str) -> anyhow::Result {
+ let re = Regex::new(r"(^https?://[^/]+)").unwrap();
+ let caps = re.captures(queue_url).context("Invalid queue URL")?;
+ let endpoint_str = caps.get(1).context("Invalid queue URL")?.as_str();
+ Ok(endpoint_str.to_string())
+}
+
+pub async fn get_sqs_client(queue_url: &str) -> anyhow::Result {
+ let mut sqs_config = preconfigured_builder().await?;
+ // region is required by the SDK to work
+ let inferred_region = queue_url_region(queue_url).unwrap_or(DEFAULT_AWS_REGION);
+ let inferred_endpoint = queue_url_endpoint(queue_url)?;
+ sqs_config.set_region(Some(inferred_region));
+ sqs_config.set_endpoint_url(Some(inferred_endpoint));
+ Ok(Client::from_conf(sqs_config.build()))
+}
+
+/// Checks whether we can establish a connection to the SQS service and we can
+/// access the provided queue_url
+pub(crate) async fn check_connectivity(queue_url: &str) -> anyhow::Result<()> {
+ let client = get_sqs_client(queue_url).await?;
+ client
+ .get_queue_attributes()
+ .queue_url(queue_url)
+ .send()
+ .await?;
+
+ Ok(())
+}
+
+#[cfg(feature = "sqs-localstack-tests")]
+pub mod test_helpers {
+ use aws_sdk_sqs::types::QueueAttributeName;
+ use ulid::Ulid;
+
+ use super::*;
+
+ pub async fn get_localstack_sqs_client() -> anyhow::Result {
+ let mut sqs_config = preconfigured_builder().await?;
+ sqs_config.set_endpoint_url(Some("http://localhost:4566".to_string()));
+ sqs_config.set_region(Some(DEFAULT_AWS_REGION));
+ Ok(Client::from_conf(sqs_config.build()))
+ }
+
+ pub async fn create_queue(sqs_client: &Client, queue_name_prefix: &str) -> String {
+ let queue_name = format!("{}-{}", queue_name_prefix, Ulid::new());
+ sqs_client
+ .create_queue()
+ .queue_name(queue_name)
+ .send()
+ .await
+ .unwrap()
+ .queue_url
+ .unwrap()
+ }
+
+ pub async fn send_message(sqs_client: &Client, queue_url: &str, payload: &str) {
+ sqs_client
+ .send_message()
+ .queue_url(queue_url)
+ .message_body(payload.to_string())
+ .send()
+ .await
+ .unwrap();
+ }
+
+ pub async fn get_queue_attribute(
+ sqs_client: &Client,
+ queue_url: &str,
+ attribute: QueueAttributeName,
+ ) -> String {
+ let queue_attributes = sqs_client
+ .get_queue_attributes()
+ .queue_url(queue_url)
+ .attribute_names(attribute.clone())
+ .send()
+ .await
+ .unwrap();
+ queue_attributes
+ .attributes
+ .unwrap()
+ .get(&attribute)
+ .unwrap()
+ .to_string()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_queue_url_region() {
+ let url = "https://sqs.eu-west-2.amazonaws.com/12345678910/test";
+ let region = queue_url_region(url);
+ assert_eq!(region, Some(Region::from_static("eu-west-2")));
+
+ let url = "https://sqs.ap-south-1.amazonaws.com/12345678910/test";
+ let region = queue_url_region(url);
+ assert_eq!(region, Some(Region::from_static("ap-south-1")));
+
+ let url = "http://localhost:4566/000000000000/test-queue";
+ let region = queue_url_region(url);
+ assert_eq!(region, None);
+ }
+
+ #[test]
+ fn test_queue_url_endpoint() {
+ let url = "https://sqs.eu-west-2.amazonaws.com/12345678910/test";
+ let endpoint = queue_url_endpoint(url).unwrap();
+ assert_eq!(endpoint, "https://sqs.eu-west-2.amazonaws.com");
+
+ let url = "https://sqs.ap-south-1.amazonaws.com/12345678910/test";
+ let endpoint = queue_url_endpoint(url).unwrap();
+ assert_eq!(endpoint, "https://sqs.ap-south-1.amazonaws.com");
+
+ let url = "http://localhost:4566/000000000000/test-queue";
+ let endpoint = queue_url_endpoint(url).unwrap();
+ assert_eq!(endpoint, "http://localhost:4566");
+
+ let url = "http://localhost:4566/000000000000/test-queue";
+ let endpoint = queue_url_endpoint(url).unwrap();
+ assert_eq!(endpoint, "http://localhost:4566");
+ }
+}
+
+#[cfg(all(test, feature = "sqs-localstack-tests"))]
+mod localstack_tests {
+ use aws_sdk_sqs::types::QueueAttributeName;
+
+ use super::*;
+ use crate::source::queue_sources::helpers::QueueReceiver;
+ use crate::source::queue_sources::sqs_queue::test_helpers::{
+ create_queue, get_localstack_sqs_client,
+ };
+
+ #[tokio::test]
+ async fn test_check_connectivity() {
+ let sqs_client = get_localstack_sqs_client().await.unwrap();
+ let queue_url = create_queue(&sqs_client, "check-connectivity").await;
+ check_connectivity(&queue_url).await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn test_receive_existing_msg_quickly() {
+ let client = test_helpers::get_localstack_sqs_client().await.unwrap();
+ let queue_url = test_helpers::create_queue(&client, "test-receive-existing-msg").await;
+ let message = "hello world";
+ test_helpers::send_message(&client, &queue_url, message).await;
+
+ let queue = Arc::new(SqsQueue::try_new(queue_url).await.unwrap());
+ let messages = tokio::time::timeout(
+ Duration::from_millis(500),
+ queue.clone().receive(5, Duration::from_secs(60)),
+ )
+ .await
+ .unwrap()
+ .unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].payload.as_slice(), message.as_bytes());
+
+ // just assess that there are no errors for now
+ queue
+ .modify_deadlines(&messages[0].metadata.ack_id, Duration::from_secs(10))
+ .await
+ .unwrap();
+ queue
+ .acknowledge(&[messages[0].metadata.ack_id.clone()])
+ .await
+ .unwrap();
+ }
+
+ #[tokio::test]
+ async fn test_acknowledge_larger_batch() {
+ let client = test_helpers::get_localstack_sqs_client().await.unwrap();
+ let queue_url = test_helpers::create_queue(&client, "test-ack-large").await;
+ let message = "hello world";
+ for _ in 0..20 {
+ test_helpers::send_message(&client, &queue_url, message).await;
+ }
+
+ let queue: Arc = Arc::new(SqsQueue::try_new(queue_url.clone()).await.unwrap());
+ let mut queue_receiver = QueueReceiver::new(queue.clone(), Duration::from_millis(200));
+ let mut messages = Vec::new();
+ for _ in 0..5 {
+ let new_messages = queue_receiver
+ .receive(20, Duration::from_secs(60))
+ .await
+ .unwrap();
+ messages.extend(new_messages.into_iter());
+ }
+ assert_eq!(messages.len(), 20);
+ let in_flight_count: usize = test_helpers::get_queue_attribute(
+ &client,
+ &queue_url,
+ QueueAttributeName::ApproximateNumberOfMessagesNotVisible,
+ )
+ .await
+ .parse()
+ .unwrap();
+ assert_eq!(in_flight_count, 20);
+
+ let ack_ids = messages
+ .iter()
+ .map(|msg| msg.metadata.ack_id.clone())
+ .collect::>();
+
+ queue.acknowledge(&ack_ids).await.unwrap();
+
+ let in_flight_count: usize = test_helpers::get_queue_attribute(
+ &client,
+ &queue_url,
+ QueueAttributeName::ApproximateNumberOfMessagesNotVisible,
+ )
+ .await
+ .parse()
+ .unwrap();
+ assert_eq!(in_flight_count, 0);
+ }
+
+ #[tokio::test]
+ async fn test_receive_wrong_queue() {
+ let client = test_helpers::get_localstack_sqs_client().await.unwrap();
+ let queue_url = test_helpers::create_queue(&client, "test-receive-existing-msg").await;
+ let bad_queue_url = format!("{}wrong", queue_url);
+ let queue = Arc::new(SqsQueue::try_new(bad_queue_url).await.unwrap());
+ tokio::time::timeout(
+ Duration::from_millis(500),
+ queue.clone().receive(5, Duration::from_secs(60)),
+ )
+ .await
+ .unwrap()
+ .unwrap_err();
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs b/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs
new file mode 100644
index 00000000000..340a6c05b95
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs
@@ -0,0 +1,341 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::sync::{Arc, Weak};
+use std::time::{Duration, Instant};
+
+use anyhow::{anyhow, Context};
+use async_trait::async_trait;
+use quickwit_actors::{
+ Actor, ActorContext, ActorExitStatus, ActorHandle, ActorState, Handler, Mailbox,
+};
+use serde_json::{json, Value as JsonValue};
+
+use super::Queue;
+use crate::source::SourceContext;
+
+#[derive(Debug, Clone)]
+pub(super) struct VisibilitySettings {
+ /// The original deadline asked from the queue when polling the messages
+ pub deadline_for_receive: Duration,
+ /// The last deadline extension when the message reading is completed
+ pub deadline_for_last_extension: Duration,
+ /// The extension applied why the VisibilityTask to maintain the message visibility
+ pub deadline_for_default_extension: Duration,
+ /// Rhe timeout for the visibility extension request
+ pub request_timeout: Duration,
+ /// an extra margin that is substracted from the expected deadline when
+ /// asserting whether we are still in time to extend the visibility
+ pub request_margin: Duration,
+}
+
+impl VisibilitySettings {
+ /// The commit timeout gives us a first estimate on how long the processing
+ /// will take for the messages. We could include other factors such as the
+ /// message size.
+ pub(super) fn from_commit_timeout(commit_timeout_secs: usize) -> Self {
+ let commit_timeout = Duration::from_secs(commit_timeout_secs as u64);
+ Self {
+ deadline_for_receive: Duration::from_secs(120) + commit_timeout,
+ deadline_for_last_extension: 2 * commit_timeout,
+ deadline_for_default_extension: Duration::from_secs(60),
+ request_timeout: Duration::from_secs(3),
+ request_margin: Duration::from_secs(1),
+ }
+ }
+}
+
+#[derive(Debug)]
+struct VisibilityTask {
+ queue: Arc,
+ ack_id: String,
+ extension_count: u64,
+ current_deadline: Instant,
+ last_extension_requested: bool,
+ visibility_settings: VisibilitySettings,
+ ref_count: Weak<()>,
+}
+
+// A handle to the visibility actor. When dropped, the actor exits and the
+// visibility isn't maintained anymore.
+pub(super) struct VisibilityTaskHandle {
+ mailbox: Mailbox,
+ actor_handle: ActorHandle,
+ ack_id: String,
+ _ref_count: Arc<()>,
+}
+
+/// Spawns actor that ensures that the visibility of a given message
+/// (represented by its ack_id) is extended when required. We prefer applying
+/// ample margins in the extension process to avoid missing deadlines while also
+/// keeping the number of extension requests (and associated cost) small.
+pub(super) fn spawn_visibility_task(
+ ctx: &SourceContext,
+ queue: Arc,
+ ack_id: String,
+ current_deadline: Instant,
+ visibility_settings: VisibilitySettings,
+) -> VisibilityTaskHandle {
+ let ref_count = Arc::new(());
+ let weak_ref = Arc::downgrade(&ref_count);
+ let task = VisibilityTask {
+ queue,
+ ack_id: ack_id.clone(),
+ extension_count: 0,
+ current_deadline,
+ last_extension_requested: false,
+ visibility_settings,
+ ref_count: weak_ref,
+ };
+ let (mailbox, actor_handle) = ctx.spawn_actor().spawn(task);
+ VisibilityTaskHandle {
+ mailbox,
+ actor_handle,
+ ack_id,
+ _ref_count: ref_count,
+ }
+}
+
+impl VisibilityTask {
+ async fn extend_visibility(
+ &mut self,
+ ctx: &ActorContext,
+ extension: Duration,
+ ) -> anyhow::Result<()> {
+ let _zone = ctx.protect_zone();
+ self.current_deadline = tokio::time::timeout(
+ self.visibility_settings.request_timeout,
+ self.queue.modify_deadlines(&self.ack_id, extension),
+ )
+ .await
+ .context("deadline extension timed out")??;
+ self.extension_count += 1;
+ Ok(())
+ }
+
+ fn next_extension(&self) -> Duration {
+ (self.current_deadline - Instant::now())
+ - self.visibility_settings.request_timeout
+ - self.visibility_settings.request_margin
+ }
+}
+
+impl VisibilityTaskHandle {
+ pub fn extension_failed(&self) -> bool {
+ self.actor_handle.state() == ActorState::Failure
+ }
+
+ pub fn ack_id(&self) -> &str {
+ &self.ack_id
+ }
+
+ pub async fn request_last_extension(self, extension: Duration) -> anyhow::Result<()> {
+ self.mailbox
+ .ask_for_res(RequestLastExtension { extension })
+ .await
+ .map_err(|e| anyhow!(e))?;
+ Ok(())
+ }
+}
+
+#[async_trait]
+impl Actor for VisibilityTask {
+ type ObservableState = JsonValue;
+
+ fn name(&self) -> String {
+ "QueueVisibilityTask".to_string()
+ }
+
+ async fn initialize(&mut self, ctx: &ActorContext) -> Result<(), ActorExitStatus> {
+ let first_extension = self.next_extension();
+ if first_extension.is_zero() {
+ return Err(anyhow!("initial visibility deadline insufficient").into());
+ }
+ ctx.schedule_self_msg(first_extension, Loop);
+ Ok(())
+ }
+
+ fn yield_after_each_message(&self) -> bool {
+ false
+ }
+
+ fn observable_state(&self) -> Self::ObservableState {
+ json!({
+ "ack_id": self.ack_id,
+ "extension_count": self.extension_count,
+ })
+ }
+}
+
+#[derive(Debug)]
+struct Loop;
+
+#[async_trait]
+impl Handler for VisibilityTask {
+ type Reply = ();
+
+ async fn handle(
+ &mut self,
+ _message: Loop,
+ ctx: &ActorContext,
+ ) -> Result<(), ActorExitStatus> {
+ if self.ref_count.strong_count() == 0 {
+ return Ok(());
+ }
+ if self.last_extension_requested {
+ return Ok(());
+ }
+ self.extend_visibility(ctx, self.visibility_settings.deadline_for_default_extension)
+ .await?;
+ ctx.schedule_self_msg(self.next_extension(), Loop);
+ Ok(())
+ }
+}
+
+/// Ensures that the visibility of the message is extended until the given
+/// deadline and then stops the extension loop.
+#[derive(Debug)]
+struct RequestLastExtension {
+ extension: Duration,
+}
+
+#[async_trait]
+impl Handler for VisibilityTask {
+ type Reply = anyhow::Result<()>;
+
+ async fn handle(
+ &mut self,
+ message: RequestLastExtension,
+ ctx: &ActorContext,
+ ) -> Result {
+ let last_deadline = Instant::now() + message.extension;
+ self.last_extension_requested = true;
+ if last_deadline > self.current_deadline {
+ Ok(self.extend_visibility(ctx, message.extension).await)
+ } else {
+ Ok(Ok(()))
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use quickwit_actors::Universe;
+ use tokio::sync::watch;
+
+ use super::*;
+ use crate::source::queue_sources::memory_queue::MemoryQueueForTests;
+
+ #[tokio::test]
+ async fn test_visibility_task_request_last_extension() {
+ // actor context
+ let universe = Universe::with_accelerated_time();
+ let (source_mailbox, _source_inbox) = universe.create_test_mailbox();
+ let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null);
+ let ctx: SourceContext =
+ ActorContext::for_test(&universe, source_mailbox, observable_state_tx);
+ // queue with test message
+ let ack_id = "ack_id".to_string();
+ let queue = Arc::new(MemoryQueueForTests::new());
+ queue.send_message("test message".to_string(), &ack_id);
+ let initial_deadline = queue
+ .clone()
+ .receive(1, Duration::from_secs(1))
+ .await
+ .unwrap()[0]
+ .metadata
+ .initial_deadline;
+ // spawn task
+ let visibility_settings = VisibilitySettings {
+ deadline_for_default_extension: Duration::from_secs(1),
+ deadline_for_last_extension: Duration::from_secs(20),
+ deadline_for_receive: Duration::from_secs(1),
+ request_timeout: Duration::from_millis(100),
+ request_margin: Duration::from_millis(100),
+ };
+ let handle = spawn_visibility_task(
+ &ctx,
+ queue.clone(),
+ ack_id.clone(),
+ initial_deadline,
+ visibility_settings.clone(),
+ );
+ // assert that the background task performs extensions
+ assert!(!handle.extension_failed());
+ tokio::time::sleep_until(initial_deadline.into()).await;
+ let next_deadline = queue.next_visibility_deadline(&ack_id).unwrap();
+ assert!(initial_deadline < next_deadline);
+ assert!(!handle.extension_failed());
+ // request last extension
+ handle
+ .request_last_extension(Duration::from_secs(5))
+ .await
+ .unwrap();
+ assert!(
+ Instant::now() + Duration::from_secs(4)
+ < queue.next_visibility_deadline(&ack_id).unwrap()
+ );
+ universe.assert_quit().await;
+ }
+
+ #[tokio::test]
+ async fn test_visibility_task_stop_on_drop() {
+ // actor context
+ let universe = Universe::with_accelerated_time();
+ let (source_mailbox, _source_inbox) = universe.create_test_mailbox();
+ let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null);
+ let ctx: SourceContext =
+ ActorContext::for_test(&universe, source_mailbox, observable_state_tx);
+ // queue with test message
+ let ack_id = "ack_id".to_string();
+ let queue = Arc::new(MemoryQueueForTests::new());
+ queue.send_message("test message".to_string(), &ack_id);
+ let initial_deadline = queue
+ .clone()
+ .receive(1, Duration::from_secs(1))
+ .await
+ .unwrap()[0]
+ .metadata
+ .initial_deadline;
+ // spawn task
+ let visibility_settings = VisibilitySettings {
+ deadline_for_default_extension: Duration::from_secs(1),
+ deadline_for_last_extension: Duration::from_secs(20),
+ deadline_for_receive: Duration::from_secs(1),
+ request_timeout: Duration::from_millis(100),
+ request_margin: Duration::from_millis(100),
+ };
+ let handle = spawn_visibility_task(
+ &ctx,
+ queue.clone(),
+ ack_id.clone(),
+ initial_deadline,
+ visibility_settings.clone(),
+ );
+ // assert that visibility is not extended after drop
+ drop(handle);
+ tokio::time::sleep_until(initial_deadline.into()).await;
+ // the message is either already expired or about to expire
+ if let Some(next_deadline) = queue.next_visibility_deadline(&ack_id) {
+ assert_eq!(next_deadline, initial_deadline);
+ }
+ // assert_eq!(q, None);
+ universe.assert_quit().await;
+ }
+}
diff --git a/quickwit/quickwit-indexing/src/source/stdin_source.rs b/quickwit/quickwit-indexing/src/source/stdin_source.rs
new file mode 100644
index 00000000000..81e9f40b3e6
--- /dev/null
+++ b/quickwit/quickwit-indexing/src/source/stdin_source.rs
@@ -0,0 +1,133 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::fmt;
+use std::time::Duration;
+
+use async_trait::async_trait;
+use quickwit_actors::{ActorExitStatus, Mailbox};
+use quickwit_common::Progress;
+use quickwit_proto::metastore::SourceType;
+use tokio::io::{AsyncBufReadExt, BufReader};
+
+use super::{BatchBuilder, BATCH_NUM_BYTES_LIMIT};
+use crate::actors::DocProcessor;
+use crate::source::{Source, SourceContext, SourceRuntime, TypedSourceFactory};
+
+pub struct StdinBatchReader {
+ reader: BufReader,
+ is_eof: bool,
+}
+
+impl StdinBatchReader {
+ pub fn new() -> Self {
+ Self {
+ reader: BufReader::new(tokio::io::stdin()),
+ is_eof: false,
+ }
+ }
+
+ async fn read_batch(&mut self, source_progress: &Progress) -> anyhow::Result {
+ let mut batch_builder = BatchBuilder::new(SourceType::Stdin);
+ while batch_builder.num_bytes < BATCH_NUM_BYTES_LIMIT {
+ let mut buf = String::new();
+ // stdin might be slow because it's depending on external
+ // input (e.g. user typing on a keyboard)
+ let bytes_read = source_progress
+ .protect_future(self.reader.read_line(&mut buf))
+ .await?;
+ if bytes_read > 0 {
+ batch_builder.add_doc(buf.into());
+ } else {
+ self.is_eof = true;
+ break;
+ }
+ }
+
+ Ok(batch_builder)
+ }
+
+ fn is_eof(&self) -> bool {
+ self.is_eof
+ }
+}
+
+pub struct StdinSource {
+ reader: StdinBatchReader,
+ num_bytes_processed: u64,
+ num_lines_processed: u64,
+}
+
+impl fmt::Debug for StdinSource {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "StdinSource")
+ }
+}
+
+#[async_trait]
+impl Source for StdinSource {
+ async fn emit_batches(
+ &mut self,
+ doc_processor_mailbox: &Mailbox,
+ ctx: &SourceContext,
+ ) -> Result {
+ let batch_builder = self.reader.read_batch(ctx.progress()).await?;
+ self.num_bytes_processed += batch_builder.num_bytes;
+ self.num_lines_processed += batch_builder.docs.len() as u64;
+ doc_processor_mailbox
+ .send_message(batch_builder.build())
+ .await?;
+ if self.reader.is_eof() {
+ ctx.send_exit_with_success(doc_processor_mailbox).await?;
+ return Err(ActorExitStatus::Success);
+ }
+
+ Ok(Duration::ZERO)
+ }
+
+ fn name(&self) -> String {
+ format!("{:?}", self)
+ }
+
+ fn observable_state(&self) -> serde_json::Value {
+ serde_json::json!({
+ "num_bytes_processed": self.num_bytes_processed,
+ "num_lines_processed": self.num_lines_processed,
+ })
+ }
+}
+
+pub struct FileSourceFactory;
+
+#[async_trait]
+impl TypedSourceFactory for FileSourceFactory {
+ type Source = StdinSource;
+ type Params = ();
+
+ async fn typed_create_source(
+ _source_runtime: SourceRuntime,
+ _params: (),
+ ) -> anyhow::Result {
+ Ok(StdinSource {
+ reader: StdinBatchReader::new(),
+ num_bytes_processed: 0,
+ num_lines_processed: 0,
+ })
+ }
+}
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/fetch.rs b/quickwit/quickwit-ingest/src/ingest_v2/fetch.rs
index 0056f9d488e..53f6b0aeeea 100644
--- a/quickwit/quickwit-ingest/src/ingest_v2/fetch.rs
+++ b/quickwit/quickwit-ingest/src/ingest_v2/fetch.rs
@@ -64,7 +64,7 @@ pub(super) struct FetchStreamTask {
}
impl fmt::Debug for FetchStreamTask {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("FetchStreamTask")
.field("client_id", &self.client_id)
.field("index_uid", &self.index_uid)
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs
index fd046d2e975..d9bc7b1be75 100644
--- a/quickwit/quickwit-ingest/src/ingest_v2/mod.rs
+++ b/quickwit/quickwit-ingest/src/ingest_v2/mod.rs
@@ -49,6 +49,7 @@ use quickwit_proto::ingest::router::{IngestRequestV2, IngestSubrequest};
use quickwit_proto::ingest::{CommitTypeV2, DocBatchV2};
use quickwit_proto::types::{DocUid, DocUidGenerator, IndexId, NodeId, SubrequestId};
use tracing::{error, info};
+use workbench::pending_subrequests;
pub use self::fetch::{FetchStreamError, MultiFetchStream};
pub use self::ingester::{wait_for_ingester_decommission, wait_for_ingester_status, Ingester};
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/router.rs b/quickwit/quickwit-ingest/src/ingest_v2/router.rs
index 470cf57bf12..4f46ed5ea92 100644
--- a/quickwit/quickwit-ingest/src/ingest_v2/router.rs
+++ b/quickwit/quickwit-ingest/src/ingest_v2/router.rs
@@ -56,7 +56,7 @@ use super::ingester::PERSIST_REQUEST_TIMEOUT;
use super::metrics::IngestResultMetrics;
use super::routing_table::RoutingTable;
use super::workbench::IngestWorkbench;
-use super::IngesterPool;
+use super::{pending_subrequests, IngesterPool};
use crate::{get_ingest_router_buffer_size, LeaderId};
/// Duration after which ingest requests time out with [`IngestV2Error::Timeout`].
@@ -171,13 +171,7 @@ impl IngestRouter {
let mut state_guard = self.state.lock().await;
- for subrequest in workbench.subworkbenches.values().filter_map(|subworbench| {
- if subworbench.is_pending() {
- Some(&subworbench.subrequest)
- } else {
- None
- }
- }) {
+ for subrequest in pending_subrequests(&workbench.subworkbenches) {
if !state_guard.routing_table.has_open_shards(
&subrequest.index_id,
&subrequest.source_id,
@@ -185,6 +179,7 @@ impl IngestRouter {
&mut debounced_request.closed_shards,
unavailable_leaders,
) {
+ // No shard available! Let's attempt to create one.
let acquire_result = state_guard
.debouncer
.acquire(&subrequest.index_id, &subrequest.source_id);
@@ -294,23 +289,36 @@ impl IngestRouter {
}
for persist_failure in persist_response.failures {
workbench.record_persist_failure(&persist_failure);
-
- if persist_failure.reason() == PersistFailureReason::ShardClosed {
- let shard_id = persist_failure.shard_id().clone();
- let index_uid: IndexUid = persist_failure.index_uid().clone();
- let source_id: SourceId = persist_failure.source_id;
- closed_shards
- .entry((index_uid, source_id))
- .or_default()
- .push(shard_id);
- } else if persist_failure.reason() == PersistFailureReason::ShardNotFound {
- let shard_id = persist_failure.shard_id().clone();
- let index_uid: IndexUid = persist_failure.index_uid().clone();
- let source_id: SourceId = persist_failure.source_id;
- deleted_shards
- .entry((index_uid, source_id))
- .or_default()
- .push(shard_id);
+ match persist_failure.reason() {
+ PersistFailureReason::ShardClosed => {
+ let shard_id = persist_failure.shard_id().clone();
+ let index_uid: IndexUid = persist_failure.index_uid().clone();
+ let source_id: SourceId = persist_failure.source_id;
+ closed_shards
+ .entry((index_uid, source_id))
+ .or_default()
+ .push(shard_id);
+ }
+ PersistFailureReason::ShardNotFound => {
+ let shard_id = persist_failure.shard_id().clone();
+ let index_uid: IndexUid = persist_failure.index_uid().clone();
+ let source_id: SourceId = persist_failure.source_id;
+ deleted_shards
+ .entry((index_uid, source_id))
+ .or_default()
+ .push(shard_id);
+ }
+ PersistFailureReason::WalFull
+ | PersistFailureReason::ShardRateLimited => {
+ // Let's record that the shard is rate limited or that the ingester
+ // that hosts has its wal full.
+ //
+ // That way we will avoid to retry the persist request on the very
+ // same node.
+ let shard_id = persist_failure.shard_id().clone();
+ workbench.rate_limited_shard.insert(shard_id);
+ }
+ _ => {}
}
}
}
@@ -349,6 +357,7 @@ impl IngestRouter {
}
async fn batch_persist(&self, workbench: &mut IngestWorkbench, commit_type: CommitTypeV2) {
+ // Let's first create the shards that might be missing.
let debounced_request = self
.make_get_or_create_open_shard_request(workbench, &self.ingester_pool)
.await;
@@ -368,11 +377,14 @@ impl IngestRouter {
// lines, validate, transform and then pack the docs into compressed batches routed
// to the right shards.
- for subrequest in workbench.pending_subrequests() {
+ let rate_limited_shards: &HashSet = &workbench.rate_limited_shard;
+ for subrequest in pending_subrequests(&workbench.subworkbenches) {
let Some(shard) = state_guard
.routing_table
.find_entry(&subrequest.index_id, &subrequest.source_id)
- .and_then(|entry| entry.next_open_shard_round_robin(&self.ingester_pool))
+ .and_then(|entry| {
+ entry.next_open_shard_round_robin(&self.ingester_pool, rate_limited_shards)
+ })
else {
no_shards_available_subrequest_ids.push(subrequest.subrequest_id);
continue;
@@ -440,7 +452,7 @@ impl IngestRouter {
&self,
ingest_request: IngestRequestV2,
max_num_attempts: usize,
- ) -> IngestV2Result {
+ ) -> IngestResponseV2 {
let commit_type = ingest_request.commit_type();
let mut workbench = IngestWorkbench::new(ingest_request.subrequests, max_num_attempts);
while !workbench.is_complete() {
@@ -471,7 +483,7 @@ impl IngestRouter {
timeout_duration.as_millis()
);
IngestV2Error::Timeout(message)
- })?
+ })
}
pub async fn debug_info(&self) -> JsonValue {
@@ -667,6 +679,7 @@ pub(super) struct PersistRequestSummary {
mod tests {
use std::collections::BTreeSet;
+ use mockall::Sequence;
use quickwit_proto::control_plane::{
GetOrCreateOpenShardsFailure, GetOrCreateOpenShardsFailureReason,
GetOrCreateOpenShardsResponse, GetOrCreateOpenShardsSuccess, MockControlPlaneService,
@@ -1685,7 +1698,7 @@ mod tests {
index_uid: Some(index_uid_clone.clone()),
source_id: "test-source".to_string(),
shard_id: Some(ShardId::from(1)),
- reason: PersistFailureReason::ShardRateLimited as i32,
+ reason: PersistFailureReason::Timeout as i32,
}],
};
Ok(response)
@@ -1877,4 +1890,164 @@ mod tests {
assert_eq!(routing_table["test-index-0"].as_array().unwrap().len(), 1);
assert_eq!(routing_table["test-index-1"].as_array().unwrap().len(), 1);
}
+
+ #[tokio::test]
+ async fn test_do_not_retry_rate_limited_shards() {
+ // We avoid retrying a shard limited shard at the scale of a workbench.
+ let self_node_id = "test-router".into();
+ let control_plane = ControlPlaneServiceClient::from_mock(MockControlPlaneService::new());
+ let ingester_pool = IngesterPool::default();
+ let replication_factor = 1;
+ let router = IngestRouter::new(
+ self_node_id,
+ control_plane,
+ ingester_pool.clone(),
+ replication_factor,
+ );
+ let mut state_guard = router.state.lock().await;
+ let index_uid: IndexUid = IndexUid::for_test("test-index-0", 0);
+
+ state_guard.routing_table.replace_shards(
+ index_uid.clone(),
+ "test-source",
+ vec![
+ Shard {
+ index_uid: Some(index_uid.clone()),
+ source_id: "test-source".to_string(),
+ shard_id: Some(ShardId::from(1)),
+ shard_state: ShardState::Open as i32,
+ leader_id: "test-ingester-0".to_string(),
+ ..Default::default()
+ },
+ Shard {
+ index_uid: Some(index_uid.clone()),
+ source_id: "test-source".to_string(),
+ shard_id: Some(ShardId::from(2)),
+ shard_state: ShardState::Open as i32,
+ leader_id: "test-ingester-0".to_string(),
+ ..Default::default()
+ },
+ ],
+ );
+ drop(state_guard);
+
+ // We have two shards.
+ // - shard 1 is rate limited
+ // - shard 2 is timeout.
+ // We expect a retry on shard 2 that is then successful.
+ let mut seq = Sequence::new();
+
+ let mut mock_ingester_0 = MockIngesterService::new();
+ mock_ingester_0
+ .expect_persist()
+ .times(1)
+ .returning(move |request| {
+ assert_eq!(request.leader_id, "test-ingester-0");
+ assert_eq!(request.commit_type(), CommitTypeV2::Auto);
+ assert_eq!(request.subrequests.len(), 1);
+ let subrequest = &request.subrequests[0];
+ assert_eq!(subrequest.subrequest_id, 0);
+ let index_uid = subrequest.index_uid().clone();
+ assert_eq!(subrequest.source_id, "test-source");
+ assert_eq!(subrequest.shard_id(), ShardId::from(1));
+ assert_eq!(
+ subrequest.doc_batch,
+ Some(DocBatchV2::for_test(["test-doc-foo"]))
+ );
+
+ let response = PersistResponse {
+ leader_id: request.leader_id,
+ successes: Vec::new(),
+ failures: vec![PersistFailure {
+ subrequest_id: 0,
+ index_uid: Some(index_uid),
+ source_id: "test-source".to_string(),
+ shard_id: Some(ShardId::from(1)),
+ reason: PersistFailureReason::ShardRateLimited as i32,
+ }],
+ };
+ Ok(response)
+ })
+ .in_sequence(&mut seq);
+
+ mock_ingester_0
+ .expect_persist()
+ .times(1)
+ .returning(move |request| {
+ assert_eq!(request.leader_id, "test-ingester-0");
+ assert_eq!(request.commit_type(), CommitTypeV2::Auto);
+ assert_eq!(request.subrequests.len(), 1);
+ let subrequest = &request.subrequests[0];
+ assert_eq!(subrequest.subrequest_id, 0);
+ let index_uid = subrequest.index_uid().clone();
+ assert_eq!(subrequest.source_id, "test-source");
+ assert_eq!(subrequest.shard_id(), ShardId::from(2));
+ assert_eq!(
+ subrequest.doc_batch,
+ Some(DocBatchV2::for_test(["test-doc-foo"]))
+ );
+
+ let response = PersistResponse {
+ leader_id: request.leader_id,
+ successes: Vec::new(),
+ failures: vec![PersistFailure {
+ subrequest_id: 0,
+ index_uid: Some(index_uid),
+ source_id: "test-source".to_string(),
+ shard_id: Some(ShardId::from(1)),
+ reason: PersistFailureReason::Timeout as i32,
+ }],
+ };
+ Ok(response)
+ })
+ .in_sequence(&mut seq);
+
+ mock_ingester_0
+ .expect_persist()
+ .times(1)
+ .returning(move |request| {
+ assert_eq!(request.leader_id, "test-ingester-0");
+ assert_eq!(request.commit_type(), CommitTypeV2::Auto);
+ assert_eq!(request.subrequests.len(), 1);
+ let subrequest = &request.subrequests[0];
+ assert_eq!(subrequest.subrequest_id, 0);
+ let index_uid = subrequest.index_uid().clone();
+ assert_eq!(subrequest.source_id, "test-source");
+ assert_eq!(subrequest.shard_id(), ShardId::from(2));
+ assert_eq!(
+ subrequest.doc_batch,
+ Some(DocBatchV2::for_test(["test-doc-foo"]))
+ );
+
+ let response = PersistResponse {
+ leader_id: request.leader_id,
+ successes: vec![PersistSuccess {
+ subrequest_id: 0,
+ index_uid: Some(index_uid),
+ source_id: "test-source".to_string(),
+ shard_id: Some(ShardId::from(1)),
+ num_persisted_docs: 1,
+ replication_position_inclusive: Some(Position::offset(0u64)),
+ parse_failures: Vec::new(),
+ }],
+ failures: Vec::new(),
+ };
+ Ok(response)
+ })
+ .in_sequence(&mut seq);
+
+ let ingester_0 = IngesterServiceClient::from_mock(mock_ingester_0);
+ ingester_pool.insert("test-ingester-0".into(), ingester_0.clone());
+
+ let ingest_request = IngestRequestV2 {
+ subrequests: vec![IngestSubrequest {
+ subrequest_id: 0,
+ index_id: "test-index-0".to_string(),
+ source_id: "test-source".to_string(),
+ doc_batch: Some(DocBatchV2::for_test(["test-doc-foo"])),
+ }],
+ commit_type: CommitTypeV2::Auto as i32,
+ };
+ router.ingest(ingest_request).await.unwrap();
+ }
}
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs
index 7c4ed6f7fed..555bb8d368a 100644
--- a/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs
+++ b/quickwit/quickwit-ingest/src/ingest_v2/routing_table.rs
@@ -147,6 +147,7 @@ impl RoutingTableEntry {
pub fn next_open_shard_round_robin(
&self,
ingester_pool: &IngesterPool,
+ rate_limited_shards: &HashSet,
) -> Option<&RoutingEntry> {
for (shards, round_robin_idx) in [
(&self.local_shards, &self.local_round_robin_idx),
@@ -157,10 +158,14 @@ impl RoutingTableEntry {
}
for _attempt in 0..shards.len() {
let shard_idx = round_robin_idx.fetch_add(1, Ordering::Relaxed);
- let shard = &shards[shard_idx % shards.len()];
-
- if shard.shard_state.is_open() && ingester_pool.contains_key(&shard.leader_id) {
- return Some(shard);
+ let shard_routing_entry: &RoutingEntry = &shards[shard_idx % shards.len()];
+ if !shard_routing_entry.shard_state.is_open()
+ || rate_limited_shards.contains(&shard_routing_entry.shard_id)
+ {
+ continue;
+ }
+ if ingester_pool.contains_key(&shard_routing_entry.leader_id) {
+ return Some(shard_routing_entry);
}
}
}
@@ -658,7 +663,10 @@ mod tests {
let table_entry = RoutingTableEntry::empty(index_uid.clone(), source_id.clone());
let ingester_pool = IngesterPool::default();
- let shard_opt = table_entry.next_open_shard_round_robin(&ingester_pool);
+ let mut rate_limited_shards = HashSet::new();
+
+ let shard_opt =
+ table_entry.next_open_shard_round_robin(&ingester_pool, &rate_limited_shards);
assert!(shard_opt.is_none());
ingester_pool.insert("test-ingester-0".into(), IngesterServiceClient::mocked());
@@ -695,17 +703,17 @@ mod tests {
remote_round_robin_idx: AtomicUsize::default(),
};
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(2));
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(3));
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(2));
@@ -753,17 +761,24 @@ mod tests {
remote_round_robin_idx: AtomicUsize::default(),
};
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(2));
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(5));
let shard = table_entry
- .next_open_shard_round_robin(&ingester_pool)
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
+ .unwrap();
+ assert_eq!(shard.shard_id, ShardId::from(2));
+
+ rate_limited_shards.insert(ShardId::from(5));
+
+ let shard = table_entry
+ .next_open_shard_round_robin(&ingester_pool, &rate_limited_shards)
.unwrap();
assert_eq!(shard.shard_id, ShardId::from(2));
}
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/workbench.rs b/quickwit/quickwit-ingest/src/ingest_v2/workbench.rs
index eeca716186a..a8bc0700270 100644
--- a/quickwit/quickwit-ingest/src/ingest_v2/workbench.rs
+++ b/quickwit/quickwit-ingest/src/ingest_v2/workbench.rs
@@ -19,6 +19,7 @@
use std::collections::{BTreeMap, HashSet};
+use quickwit_common::rate_limited_error;
use quickwit_proto::control_plane::{
GetOrCreateOpenShardsFailure, GetOrCreateOpenShardsFailureReason,
};
@@ -26,8 +27,8 @@ use quickwit_proto::ingest::ingester::{PersistFailure, PersistFailureReason, Per
use quickwit_proto::ingest::router::{
IngestFailure, IngestFailureReason, IngestResponseV2, IngestSubrequest, IngestSuccess,
};
-use quickwit_proto::ingest::{IngestV2Error, IngestV2Result, RateLimitingCause};
-use quickwit_proto::types::{NodeId, SubrequestId};
+use quickwit_proto::ingest::{IngestV2Error, RateLimitingCause};
+use quickwit_proto::types::{NodeId, ShardId, SubrequestId};
use tracing::warn;
use super::router::PersistRequestSummary;
@@ -37,6 +38,7 @@ use super::router::PersistRequestSummary;
#[derive(Default)]
pub(super) struct IngestWorkbench {
pub subworkbenches: BTreeMap,
+ pub rate_limited_shard: HashSet,
pub num_successes: usize,
/// The number of batch persist attempts. This is not sum of the number of attempts for each
/// subrequest.
@@ -51,6 +53,19 @@ pub(super) struct IngestWorkbench {
pub unavailable_leaders: HashSet,
}
+/// Returns an iterator of pending of subrequests, sorted by sub request id.
+pub(super) fn pending_subrequests(
+ subworkbenches: &BTreeMap,
+) -> impl Iterator
- {
+ subworkbenches.values().filter_map(|subworbench| {
+ if subworbench.is_pending() {
+ Some(&subworbench.subrequest)
+ } else {
+ None
+ }
+ })
+}
+
impl IngestWorkbench {
pub fn new(ingest_subrequests: Vec, max_num_attempts: usize) -> Self {
let subworkbenches: BTreeMap = ingest_subrequests
@@ -92,17 +107,6 @@ impl IngestWorkbench {
.all(|subworbench| !subworbench.is_pending())
}
- #[cfg(not(test))]
- pub fn pending_subrequests(&self) -> impl Iterator
- {
- self.subworkbenches.values().filter_map(|subworbench| {
- if subworbench.is_pending() {
- Some(&subworbench.subrequest)
- } else {
- None
- }
- })
- }
-
pub fn record_get_or_create_open_shards_failure(
&mut self,
open_shards_failure: GetOrCreateOpenShardsFailure,
@@ -155,12 +159,18 @@ impl IngestWorkbench {
}
IngestV2Error::Unavailable(_) => {
self.unavailable_leaders.insert(persist_summary.leader_id);
-
for subrequest_id in persist_summary.subrequest_ids {
self.record_ingester_unavailable(subrequest_id);
}
}
- IngestV2Error::Internal(_) | IngestV2Error::ShardNotFound { .. } => {
+ IngestV2Error::Internal(internal_err_msg) => {
+ rate_limited_error!(limit_per_min=6, err_msg=%internal_err_msg, "persist error: internal error during persist");
+ for subrequest_id in persist_summary.subrequest_ids {
+ self.record_internal_error(subrequest_id);
+ }
+ }
+ IngestV2Error::ShardNotFound { shard_id } => {
+ rate_limited_error!(limit_per_min=6, shard_id=%shard_id, "persist error: shard not found");
for subrequest_id in persist_summary.subrequest_ids {
self.record_internal_error(subrequest_id);
}
@@ -213,7 +223,7 @@ impl IngestWorkbench {
);
}
- pub fn into_ingest_result(self) -> IngestV2Result {
+ pub fn into_ingest_result(self) -> IngestResponseV2 {
let num_subworkbenches = self.subworkbenches.len();
let mut successes = Vec::with_capacity(self.num_successes);
let mut failures = Vec::with_capacity(num_subworkbenches - self.num_successes);
@@ -257,27 +267,10 @@ impl IngestWorkbench {
failures.sort_by_key(|failure| failure.subrequest_id);
}
- let response = IngestResponseV2 {
+ IngestResponseV2 {
successes,
failures,
- };
- Ok(response)
- }
-
- #[cfg(test)]
- pub fn pending_subrequests(&self) -> impl Iterator
- {
- use itertools::Itertools;
-
- self.subworkbenches
- .values()
- .filter_map(|subworbench| {
- if subworbench.is_pending() {
- Some(&subworbench.subrequest)
- } else {
- None
- }
- })
- .sorted_by_key(|subrequest| subrequest.subrequest_id)
+ }
}
}
@@ -438,7 +431,7 @@ mod tests {
},
];
let mut workbench = IngestWorkbench::new(ingest_subrequests, 1);
- assert_eq!(workbench.pending_subrequests().count(), 2);
+ assert_eq!(pending_subrequests(&workbench.subworkbenches).count(), 2);
assert!(!workbench.is_complete());
let persist_success = PersistSuccess {
@@ -448,10 +441,9 @@ mod tests {
workbench.record_persist_success(persist_success);
assert_eq!(workbench.num_successes, 1);
- assert_eq!(workbench.pending_subrequests().count(), 1);
+ assert_eq!(pending_subrequests(&workbench.subworkbenches).count(), 1);
assert_eq!(
- workbench
- .pending_subrequests()
+ pending_subrequests(&workbench.subworkbenches)
.next()
.unwrap()
.subrequest_id,
@@ -469,10 +461,9 @@ mod tests {
workbench.record_persist_failure(&persist_failure);
assert_eq!(workbench.num_successes, 1);
- assert_eq!(workbench.pending_subrequests().count(), 1);
+ assert_eq!(pending_subrequests(&workbench.subworkbenches).count(), 1);
assert_eq!(
- workbench
- .pending_subrequests()
+ pending_subrequests(&workbench.subworkbenches)
.next()
.unwrap()
.subrequest_id,
@@ -491,7 +482,7 @@ mod tests {
assert!(workbench.is_complete());
assert_eq!(workbench.num_successes, 2);
- assert_eq!(workbench.pending_subrequests().count(), 0);
+ assert_eq!(pending_subrequests(&workbench.subworkbenches).count(), 0);
}
#[test]
@@ -692,7 +683,7 @@ mod tests {
#[test]
fn test_ingest_workbench_into_ingest_result() {
let workbench = IngestWorkbench::new(Vec::new(), 0);
- let response = workbench.into_ingest_result().unwrap();
+ let response = workbench.into_ingest_result();
assert!(response.successes.is_empty());
assert!(response.failures.is_empty());
@@ -715,7 +706,7 @@ mod tests {
workbench.record_no_shards_available(1);
- let response = workbench.into_ingest_result().unwrap();
+ let response = workbench.into_ingest_result();
assert_eq!(response.successes.len(), 1);
assert_eq!(response.successes[0].subrequest_id, 0);
@@ -734,7 +725,7 @@ mod tests {
let failure = SubworkbenchFailure::Persist(PersistFailureReason::Timeout);
workbench.record_failure(0, failure);
- let ingest_response = workbench.into_ingest_result().unwrap();
+ let ingest_response = workbench.into_ingest_result();
assert_eq!(ingest_response.successes.len(), 0);
assert_eq!(
ingest_response.failures[0].reason(),
diff --git a/quickwit/quickwit-integration-tests/Cargo.toml b/quickwit/quickwit-integration-tests/Cargo.toml
index 7030d2ba004..aa6a692c236 100644
--- a/quickwit/quickwit-integration-tests/Cargo.toml
+++ b/quickwit/quickwit-integration-tests/Cargo.toml
@@ -10,10 +10,17 @@ repository.workspace = true
authors.workspace = true
license.workspace = true
+[features]
+sqs-localstack-tests = [
+ "quickwit-indexing/sqs",
+ "quickwit-indexing/sqs-localstack-tests"
+]
+
[dependencies]
[dev-dependencies]
anyhow = { workspace = true }
+aws-sdk-sqs = { workspace = true }
futures-util = { workspace = true }
hyper = { workspace = true }
itertools = { workspace = true }
@@ -23,11 +30,13 @@ tempfile = { workspace = true }
tokio = { workspace = true }
tonic = { workspace = true }
tracing = { workspace = true }
+tracing-subscriber = { workspace = true }
quickwit-actors = { workspace = true, features = ["testsuite"] }
quickwit-cli = { workspace = true }
quickwit-common = { workspace = true, features = ["testsuite"] }
quickwit-config = { workspace = true, features = ["testsuite"] }
+quickwit-indexing = { workspace = true, features = ["testsuite"] }
quickwit-metastore = { workspace = true, features = ["testsuite"] }
quickwit-proto = { workspace = true, features = ["testsuite"] }
quickwit-rest-client = { workspace = true }
diff --git a/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs b/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs
index 824e24988e1..17eed1116fb 100644
--- a/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs
+++ b/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs
@@ -24,6 +24,7 @@ use std::path::PathBuf;
use std::str::FromStr;
use std::time::{Duration, Instant};
+use anyhow::Context;
use futures_util::future;
use itertools::Itertools;
use quickwit_actors::ActorExitStatus;
@@ -419,7 +420,12 @@ impl ClusterSandbox {
input_format: quickwit_config::SourceInputFormat::Json,
overwrite: false,
vrl_script: None,
- input_path_opt: Some(tmp_data_file.path().to_path_buf()),
+ input_path_opt: Some(QuickwitUri::from_str(
+ tmp_data_file
+ .path()
+ .to_str()
+ .context("temp path could not be converted to URI")?,
+ )?),
})
.await?;
Ok(())
diff --git a/quickwit/quickwit-integration-tests/src/tests/mod.rs b/quickwit/quickwit-integration-tests/src/tests/mod.rs
index 9c50db45bbc..103fd9a9f1f 100644
--- a/quickwit/quickwit-integration-tests/src/tests/mod.rs
+++ b/quickwit/quickwit-integration-tests/src/tests/mod.rs
@@ -19,4 +19,6 @@
mod basic_tests;
mod index_tests;
+#[cfg(feature = "sqs-localstack-tests")]
+mod sqs_tests;
mod update_tests;
diff --git a/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs b/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs
new file mode 100644
index 00000000000..d2f86b00ff6
--- /dev/null
+++ b/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs
@@ -0,0 +1,167 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at hello@quickwit.io.
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+use std::io::Write;
+use std::iter;
+use std::str::FromStr;
+use std::time::Duration;
+
+use aws_sdk_sqs::types::QueueAttributeName;
+use quickwit_common::test_utils::wait_until_predicate;
+use quickwit_common::uri::Uri;
+use quickwit_config::ConfigFormat;
+use quickwit_indexing::source::sqs_queue::test_helpers as sqs_test_helpers;
+use quickwit_metastore::SplitState;
+use quickwit_serve::SearchRequestQueryString;
+use tempfile::NamedTempFile;
+use tracing::info;
+
+use crate::test_utils::ClusterSandbox;
+
+fn create_mock_data_file(num_lines: usize) -> (NamedTempFile, Uri) {
+ let mut temp_file = tempfile::NamedTempFile::new().unwrap();
+ for i in 0..num_lines {
+ writeln!(temp_file, "{{\"body\": \"hello {}\"}}", i).unwrap()
+ }
+ temp_file.flush().unwrap();
+ let path = temp_file.path().to_str().unwrap();
+ let uri = Uri::from_str(path).unwrap();
+ (temp_file, uri)
+}
+
+#[tokio::test]
+async fn test_sqs_single_node_cluster() {
+ tracing_subscriber::fmt::init();
+ let sandbox = ClusterSandbox::start_standalone_node().await.unwrap();
+ let index_id = "test-sqs-source-single-node-cluster";
+ let index_config = format!(
+ r#"
+ version: 0.8
+ index_id: {}
+ doc_mapping:
+ field_mappings:
+ - name: body
+ type: text
+ indexing_settings:
+ commit_timeout_secs: 1
+ "#,
+ index_id
+ );
+
+ info!("create SQS queue");
+ let sqs_client = sqs_test_helpers::get_localstack_sqs_client().await.unwrap();
+ let queue_url = sqs_test_helpers::create_queue(&sqs_client, "test-single-node-cluster").await;
+
+ sandbox.wait_for_cluster_num_ready_nodes(1).await.unwrap();
+
+ info!("create index");
+ sandbox
+ .indexer_rest_client
+ .indexes()
+ .create(index_config.clone(), ConfigFormat::Yaml, false)
+ .await
+ .unwrap();
+
+ let source_id: &str = "test-sqs-single-node-cluster";
+ let source_config_input = format!(
+ r#"
+ version: 0.7
+ source_id: {}
+ desired_num_pipelines: 1
+ max_num_pipelines_per_indexer: 1
+ source_type: file
+ params:
+ notifications:
+ - type: sqs
+ queue_url: {}
+ message_type: raw_uri
+ input_format: plain_text
+ "#,
+ source_id, queue_url
+ );
+
+ info!("create file source with SQS notification");
+ sandbox
+ .indexer_rest_client
+ .sources(index_id)
+ .create(source_config_input, ConfigFormat::Yaml)
+ .await
+ .unwrap();
+
+ // Send messages with duplicates
+ let tmp_mock_data_files: Vec<_> = iter::repeat_with(|| create_mock_data_file(1000))
+ .take(10)
+ .collect();
+ for (_, uri) in &tmp_mock_data_files {
+ sqs_test_helpers::send_message(&sqs_client, &queue_url, uri.as_str()).await;
+ }
+ sqs_test_helpers::send_message(&sqs_client, &queue_url, tmp_mock_data_files[0].1.as_str())
+ .await;
+ sqs_test_helpers::send_message(&sqs_client, &queue_url, tmp_mock_data_files[5].1.as_str())
+ .await;
+
+ info!("wait for split to be published");
+ sandbox
+ .wait_for_splits(index_id, Some(vec![SplitState::Published]), 1)
+ .await
+ .unwrap();
+
+ info!("count docs using search");
+ let search_result = sandbox
+ .indexer_rest_client
+ .search(
+ index_id,
+ SearchRequestQueryString {
+ query: "".to_string(),
+ max_hits: 0,
+ ..Default::default()
+ },
+ )
+ .await
+ .unwrap();
+ assert_eq!(search_result.num_hits, 10 * 1000);
+
+ wait_until_predicate(
+ || async {
+ let in_flight_count: usize = sqs_test_helpers::get_queue_attribute(
+ &sqs_client,
+ &queue_url,
+ QueueAttributeName::ApproximateNumberOfMessagesNotVisible,
+ )
+ .await
+ .parse()
+ .unwrap();
+ in_flight_count == 2
+ },
+ Duration::from_secs(5),
+ Duration::from_millis(100),
+ )
+ .await
+ .expect("Number of in-flight messages didn't reach 2 within the timeout");
+
+ info!("delete index");
+ sandbox
+ .indexer_rest_client
+ .indexes()
+ .delete(index_id, false)
+ .await
+ .unwrap();
+
+ sandbox.shutdown().await.unwrap();
+}
diff --git a/quickwit/quickwit-integration-tests/src/tests/update_tests/doc_mapping_tests.rs b/quickwit/quickwit-integration-tests/src/tests/update_tests/doc_mapping_tests.rs
index c5208741000..8d9725ab7f0 100644
--- a/quickwit/quickwit-integration-tests/src/tests/update_tests/doc_mapping_tests.rs
+++ b/quickwit/quickwit-integration-tests/src/tests/update_tests/doc_mapping_tests.rs
@@ -324,9 +324,7 @@ async fn test_update_doc_mapping_object_to_json() {
.await;
}
-// TODO expected to be fix as part of #5084
#[tokio::test]
-#[ignore]
async fn test_update_doc_mapping_tokenizer_default_to_raw() {
let index_id = "update-tokenizer-default-to-raw";
let original_doc_mappings = json!({
@@ -349,10 +347,11 @@ async fn test_update_doc_mapping_tokenizer_default_to_raw() {
ingest_after_update,
&[
("body:hello", Ok(&[json!({"body": "hello-world"})])),
- ("body:world", Ok(&[json!({"body": "bonjour-monde"})])),
+ ("body:world", Ok(&[json!({"body": "hello-world"})])),
// phrases queries won't apply to older splits that didn't support them
("body:\"hello world\"", Ok(&[])),
("body:\"hello-world\"", Ok(&[])),
+ ("body:\"hello-worl\"*", Ok(&[])),
("body:bonjour", Ok(&[])),
("body:monde", Ok(&[])),
// the raw tokenizer only returns exact matches
@@ -361,14 +360,16 @@ async fn test_update_doc_mapping_tokenizer_default_to_raw() {
"body:\"bonjour-monde\"",
Ok(&[json!({"body": "bonjour-monde"})]),
),
+ (
+ "body:\"bonjour-mond\"*",
+ Ok(&[json!({"body": "bonjour-monde"})]),
+ ),
],
)
.await;
}
-// TODO expected to be fix as part of #5084
#[tokio::test]
-#[ignore]
async fn test_update_doc_mapping_tokenizer_add_position() {
let index_id = "update-tokenizer-add-position";
let original_doc_mappings = json!({
@@ -395,6 +396,7 @@ async fn test_update_doc_mapping_tokenizer_add_position() {
// phrases queries don't apply to older splits that didn't support them
("body:\"hello-world\"", Ok(&[])),
("body:\"hello world\"", Ok(&[])),
+ ("body:\"hello-worl\"*", Ok(&[])),
("body:bonjour", Ok(&[json!({"body": "bonjour-monde"})])),
("body:monde", Ok(&[json!({"body": "bonjour-monde"})])),
(
@@ -405,6 +407,10 @@ async fn test_update_doc_mapping_tokenizer_add_position() {
"body:\"bonjour monde\"",
Ok(&[json!({"body": "bonjour-monde"})]),
),
+ (
+ "body:\"bonjour-mond\"*",
+ Ok(&[json!({"body": "bonjour-monde"})]),
+ ),
],
)
.await;
@@ -455,6 +461,39 @@ async fn test_update_doc_mapping_tokenizer_raw_to_phrase() {
.await;
}
+#[tokio::test]
+async fn test_update_doc_mapping_unindexed_to_indexed() {
+ let index_id = "update-not-indexed-to-indexed";
+ let original_doc_mappings = json!({
+ "field_mappings": [
+ {"name": "body", "type": "text", "indexed": false}
+ ]
+ });
+ let ingest_before_update = &[json!({"body": "hello"})];
+ let updated_doc_mappings = json!({
+ "field_mappings": [
+ {"name": "body", "type": "text", "tokenizer": "raw"}
+ ]
+ });
+ let ingest_after_update = &[json!({"body": "bonjour"})];
+ validate_search_across_doc_mapping_updates(
+ index_id,
+ original_doc_mappings,
+ ingest_before_update,
+ updated_doc_mappings,
+ ingest_after_update,
+ &[
+ // term query won't apply to older splits that weren't indexed
+ ("body:hello", Ok(&[])),
+ ("body:IN [hello]", Ok(&[])),
+ // works on newer data
+ ("body:bonjour", Ok(&[json!({"body": "bonjour"})])),
+ ("body:IN [bonjour]", Ok(&[json!({"body": "bonjour"})])),
+ ],
+ )
+ .await;
+}
+
#[tokio::test]
async fn test_update_doc_mapping_strict_to_dynamic() {
let index_id = "update-strict-to-dynamic";
diff --git a/quickwit/quickwit-lambda/Cargo.toml b/quickwit/quickwit-lambda/Cargo.toml
index abb15cbcca0..009020df57f 100644
--- a/quickwit/quickwit-lambda/Cargo.toml
+++ b/quickwit/quickwit-lambda/Cargo.toml
@@ -26,7 +26,7 @@ chrono = { workspace = true }
flate2 = { workspace = true }
http = { workspace = true }
lambda_http = "0.8.0"
-lambda_runtime = "0.11.1"
+lambda_runtime = "0.13.0"
mime_guess = { workspace = true }
once_cell = { workspace = true }
opentelemetry = { workspace = true }
diff --git a/quickwit/quickwit-lambda/src/indexer/handler.rs b/quickwit/quickwit-lambda/src/indexer/handler.rs
index 1282b0e54a4..c2027b25955 100644
--- a/quickwit/quickwit-lambda/src/indexer/handler.rs
+++ b/quickwit/quickwit-lambda/src/indexer/handler.rs
@@ -34,7 +34,7 @@ async fn indexer_handler(event: LambdaEvent) -> Result {
let payload = serde_json::from_value::(event.payload)?;
let ingest_res = ingest(IngestArgs {
- input_path: payload.uri(),
+ input_path: payload.uri()?,
input_format: quickwit_config::SourceInputFormat::Json,
overwrite: false,
vrl_script: None,
diff --git a/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs b/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs
index 176dc8bbc9a..282eb6622c7 100644
--- a/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs
+++ b/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs
@@ -19,7 +19,7 @@
use std::collections::HashSet;
use std::num::NonZeroUsize;
-use std::path::{Path, PathBuf};
+use std::path::Path;
use anyhow::{bail, Context};
use chitchat::transport::ChannelTransport;
@@ -138,12 +138,12 @@ pub(super) async fn send_telemetry() {
/// Convert the incoming file path to a source config
pub(super) async fn configure_source(
- input_path: PathBuf,
+ input_uri: Uri,
input_format: SourceInputFormat,
vrl_script: Option,
) -> anyhow::Result {
let transform_config = vrl_script.map(|vrl_script| TransformConfig::new(vrl_script, None));
- let source_params = SourceParams::file(input_path.clone());
+ let source_params = SourceParams::file_from_uri(input_uri);
Ok(SourceConfig {
source_id: LAMBDA_SOURCE_ID.to_owned(),
num_pipelines: NonZeroUsize::new(1).expect("1 is always non-zero."),
diff --git a/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs b/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs
index 6faf495b85f..13dd7f9d1b5 100644
--- a/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs
+++ b/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs
@@ -20,7 +20,6 @@
mod helpers;
use std::collections::HashSet;
-use std::path::PathBuf;
use anyhow::bail;
use helpers::{
@@ -31,6 +30,7 @@ use quickwit_actors::Universe;
use quickwit_cli::start_actor_runtimes;
use quickwit_cli::tool::start_statistics_reporting_loop;
use quickwit_common::runtimes::RuntimesConfig;
+use quickwit_common::uri::Uri;
use quickwit_config::service::QuickwitService;
use quickwit_config::SourceInputFormat;
use quickwit_index_management::clear_cache_directory;
@@ -43,7 +43,7 @@ use crate::utils::load_node_config;
#[derive(Debug, Eq, PartialEq)]
pub struct IngestArgs {
- pub input_path: PathBuf,
+ pub input_path: Uri,
pub input_format: SourceInputFormat,
pub overwrite: bool,
pub vrl_script: Option,
diff --git a/quickwit/quickwit-lambda/src/indexer/model.rs b/quickwit/quickwit-lambda/src/indexer/model.rs
index fe6ae14aea4..2cf785ca178 100644
--- a/quickwit/quickwit-lambda/src/indexer/model.rs
+++ b/quickwit/quickwit-lambda/src/indexer/model.rs
@@ -17,9 +17,10 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
-use std::path::PathBuf;
+use std::str::FromStr;
use aws_lambda_events::event::s3::S3Event;
+use quickwit_common::uri::Uri;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Deserialize, Serialize)]
@@ -31,17 +32,17 @@ pub enum IndexerEvent {
}
impl IndexerEvent {
- pub fn uri(&self) -> PathBuf {
- match &self {
- IndexerEvent::Custom { source_uri } => PathBuf::from(source_uri),
+ pub fn uri(&self) -> anyhow::Result {
+ let path: String = match self {
+ IndexerEvent::Custom { source_uri } => source_uri.clone(),
IndexerEvent::S3(event) => [
"s3://",
event.records[0].s3.bucket.name.as_ref().unwrap(),
event.records[0].s3.object.key.as_ref().unwrap(),
]
- .iter()
- .collect(),
- }
+ .join(""),
+ };
+ Uri::from_str(&path)
}
}
@@ -58,14 +59,14 @@ mod tests {
});
let parsed_cust_event: IndexerEvent = serde_json::from_value(cust_event).unwrap();
assert_eq!(
- parsed_cust_event.uri(),
- PathBuf::from("s3://quickwit-test/test.json"),
+ parsed_cust_event.uri().unwrap(),
+ Uri::from_str("s3://quickwit-test/test.json").unwrap(),
);
}
#[test]
fn test_s3_event_uri() {
- let cust_event = json!({
+ let s3_event = json!({
"Records": [
{
"eventVersion": "2.0",
@@ -103,10 +104,10 @@ mod tests {
}
]
});
- let parsed_cust_event: IndexerEvent = serde_json::from_value(cust_event).unwrap();
+ let s3_event: IndexerEvent = serde_json::from_value(s3_event).unwrap();
assert_eq!(
- parsed_cust_event.uri(),
- PathBuf::from("s3://quickwit-test/test.json"),
+ s3_event.uri().unwrap(),
+ Uri::from_str("s3://quickwit-test/test.json").unwrap(),
);
}
}
diff --git a/quickwit/quickwit-metastore/Cargo.toml b/quickwit/quickwit-metastore/Cargo.toml
index 1f2b3d87ce8..612ada72d52 100644
--- a/quickwit/quickwit-metastore/Cargo.toml
+++ b/quickwit/quickwit-metastore/Cargo.toml
@@ -47,7 +47,7 @@ quickwit-query = { workspace = true }
quickwit-storage = { workspace = true }
[dev-dependencies]
-dotenv = { workspace = true }
+dotenvy = { workspace = true }
futures = { workspace = true }
md5 = { workspace = true }
mockall = { workspace = true }
diff --git a/quickwit/quickwit-metastore/src/checkpoint.rs b/quickwit/quickwit-metastore/src/checkpoint.rs
index 5627af53a0e..bd8f2b7bd97 100644
--- a/quickwit/quickwit-metastore/src/checkpoint.rs
+++ b/quickwit/quickwit-metastore/src/checkpoint.rs
@@ -33,7 +33,7 @@ use thiserror::Error;
use tracing::{debug, warn};
/// A `PartitionId` uniquely identifies a partition for a given source.
-#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
pub struct PartitionId(pub Arc);
impl PartitionId {
diff --git a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs
index 120c0574831..d7c75e9a6f3 100644
--- a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs
+++ b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs
@@ -32,7 +32,6 @@ use itertools::Itertools;
use quickwit_common::pretty::PrettySample;
use quickwit_config::{
DocMapping, IndexingSettings, RetentionPolicy, SearchSettings, SourceConfig,
- INGEST_V2_SOURCE_ID,
};
use quickwit_proto::metastore::{
AcquireShardsRequest, AcquireShardsResponse, DeleteQuery, DeleteShardsRequest,
@@ -48,6 +47,7 @@ use tracing::{info, warn};
use super::MutationOccurred;
use crate::checkpoint::IndexCheckpointDelta;
+use crate::metastore::use_shard_api;
use crate::{split_tag_filter, IndexMetadata, ListSplitsQuery, Split, SplitMetadata, SplitState};
/// A `FileBackedIndex` object carries an index metadata and its split metadata.
@@ -82,6 +82,7 @@ pub(crate) struct FileBackedIndex {
#[cfg(any(test, feature = "testsuite"))]
impl quickwit_config::TestableForRegression for FileBackedIndex {
fn sample_for_regression() -> Self {
+ use quickwit_config::INGEST_V2_SOURCE_ID;
use quickwit_proto::ingest::{Shard, ShardState};
use quickwit_proto::types::{DocMappingUid, Position, ShardId};
@@ -381,8 +382,14 @@ impl FileBackedIndex {
) -> MetastoreResult<()> {
if let Some(checkpoint_delta) = checkpoint_delta_opt {
let source_id = checkpoint_delta.source_id.clone();
+ let source = self.metadata.sources.get(&source_id).ok_or_else(|| {
+ MetastoreError::NotFound(EntityKind::Source {
+ index_id: self.index_id().to_string(),
+ source_id: source_id.clone(),
+ })
+ })?;
- if source_id == INGEST_V2_SOURCE_ID {
+ if use_shard_api(&source.source_params) {
let publish_token = publish_token_opt.ok_or_else(|| {
let message = format!(
"publish token is required for publishing splits for source `{source_id}`"
diff --git a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs
index 236a4dcd9c7..c9246b1aacf 100644
--- a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs
+++ b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs
@@ -131,7 +131,7 @@ impl Shards {
follower_id: subrequest.follower_id,
doc_mapping_uid: subrequest.doc_mapping_uid,
publish_position_inclusive: Some(Position::Beginning),
- publish_token: None,
+ publish_token: subrequest.publish_token.clone(),
};
mutation_occurred = true;
entry.insert(shard.clone());
@@ -335,6 +335,7 @@ mod tests {
leader_id: "leader_id".to_string(),
follower_id: None,
doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: None,
};
let MutationOccurred::Yes(subresponse) = shards.open_shard(subrequest.clone()).unwrap()
else {
@@ -349,6 +350,7 @@ mod tests {
assert_eq!(shard.shard_state(), ShardState::Open);
assert_eq!(shard.leader_id, "leader_id");
assert_eq!(shard.follower_id, None);
+ assert_eq!(shard.publish_token, None);
assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
let MutationOccurred::No(subresponse) = shards.open_shard(subrequest).unwrap() else {
@@ -367,6 +369,7 @@ mod tests {
leader_id: "leader_id".to_string(),
follower_id: Some("follower_id".to_string()),
doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: Some("publish_token".to_string()),
};
let MutationOccurred::Yes(subresponse) = shards.open_shard(subrequest).unwrap() else {
panic!("Expected `MutationOccurred::No`");
diff --git a/quickwit/quickwit-metastore/src/metastore/mod.rs b/quickwit/quickwit-metastore/src/metastore/mod.rs
index cbddbbd4c1f..89822b549cc 100644
--- a/quickwit/quickwit-metastore/src/metastore/mod.rs
+++ b/quickwit/quickwit-metastore/src/metastore/mod.rs
@@ -33,7 +33,8 @@ pub use index_metadata::IndexMetadata;
use itertools::Itertools;
use quickwit_common::thread_pool::run_cpu_intensive;
use quickwit_config::{
- DocMapping, IndexConfig, IndexingSettings, RetentionPolicy, SearchSettings, SourceConfig,
+ DocMapping, FileSourceParams, IndexConfig, IndexingSettings, RetentionPolicy, SearchSettings,
+ SourceConfig, SourceParams,
};
use quickwit_doc_mapper::tag_pruning::TagFilterAst;
use quickwit_proto::metastore::{
@@ -903,6 +904,25 @@ impl Default for FilterRange {
}
}
+/// Maps the given source params to whether checkpoints should be stored in the index metadata
+/// (false) or the shard table (true)
+fn use_shard_api(params: &SourceParams) -> bool {
+ match params {
+ SourceParams::File(FileSourceParams::Filepath(_)) => false,
+ SourceParams::File(FileSourceParams::Notifications(_)) => true,
+ SourceParams::Ingest => true,
+ SourceParams::IngestApi => false,
+ SourceParams::IngestCli => false,
+ SourceParams::Kafka(_) => false,
+ SourceParams::Kinesis(_) => false,
+ SourceParams::PubSub(_) => false,
+ SourceParams::Pulsar(_) => false,
+ SourceParams::Stdin => panic!("stdin cannot be checkpointed"),
+ SourceParams::Vec(_) => false,
+ SourceParams::Void(_) => false,
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs b/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs
index 0289b10867e..7d916fcd29f 100644
--- a/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs
+++ b/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs
@@ -28,7 +28,6 @@ use quickwit_common::uri::Uri;
use quickwit_common::ServiceStream;
use quickwit_config::{
validate_index_id_pattern, IndexTemplate, IndexTemplateId, PostgresMetastoreConfig,
- INGEST_V2_SOURCE_ID,
};
use quickwit_proto::ingest::{Shard, ShardState};
use quickwit_proto::metastore::{
@@ -68,7 +67,7 @@ use crate::file_backed::MutationOccurred;
use crate::metastore::postgres::model::Shards;
use crate::metastore::postgres::utils::split_maturity_timestamp;
use crate::metastore::{
- IndexesMetadataResponseExt, PublishSplitsRequestExt, STREAM_SPLITS_CHUNK_SIZE,
+ use_shard_api, IndexesMetadataResponseExt, PublishSplitsRequestExt, STREAM_SPLITS_CHUNK_SIZE,
};
use crate::{
AddSourceRequestExt, CreateIndexRequestExt, IndexMetadata, IndexMetadataResponseExt,
@@ -130,16 +129,23 @@ impl PostgresqlMetastore {
}
/// Returns an Index object given an index_id or None if it does not exist.
-async fn index_opt<'a, E>(executor: E, index_id: &str) -> MetastoreResult>
-where E: sqlx::Executor<'a, Database = Postgres> {
- let index_opt: Option = sqlx::query_as::<_, PgIndex>(
+async fn index_opt<'a, E>(
+ executor: E,
+ index_id: &str,
+ lock: bool,
+) -> MetastoreResult>
+where
+ E: sqlx::Executor<'a, Database = Postgres>,
+{
+ let index_opt: Option = sqlx::query_as::<_, PgIndex>(&format!(
r#"
SELECT *
FROM indexes
WHERE index_id = $1
- FOR UPDATE
+ {}
"#,
- )
+ if lock { "FOR UPDATE" } else { "" }
+ ))
.bind(index_id)
.fetch_optional(executor)
.await?;
@@ -150,18 +156,20 @@ where E: sqlx::Executor<'a, Database = Postgres> {
async fn index_opt_for_uid<'a, E>(
executor: E,
index_uid: IndexUid,
+ lock: bool,
) -> MetastoreResult>
where
E: sqlx::Executor<'a, Database = Postgres>,
{
- let index_opt: Option = sqlx::query_as::<_, PgIndex>(
+ let index_opt: Option = sqlx::query_as::<_, PgIndex>(&format!(
r#"
SELECT *
FROM indexes
WHERE index_uid = $1
- FOR UPDATE
+ {}
"#,
- )
+ if lock { "FOR UPDATE" } else { "" }
+ ))
.bind(&index_uid)
.fetch_optional(executor)
.await?;
@@ -171,8 +179,9 @@ where
async fn index_metadata(
tx: &mut Transaction<'_, Postgres>,
index_id: &str,
+ lock: bool,
) -> MetastoreResult {
- index_opt(tx.as_mut(), index_id)
+ index_opt(tx.as_mut(), index_id, lock)
.await?
.ok_or_else(|| {
MetastoreError::NotFound(EntityKind::Index {
@@ -306,7 +315,7 @@ where
M: FnOnce(&mut IndexMetadata) -> Result, E>,
{
let index_id = &index_uid.index_id;
- let mut index_metadata = index_metadata(tx, index_id).await?;
+ let mut index_metadata = index_metadata(tx, index_id, true).await?;
if index_metadata.index_uid != index_uid {
return Err(MetastoreError::NotFound(EntityKind::Index {
index_id: index_id.to_string(),
@@ -422,9 +431,9 @@ impl MetastoreService for PostgresqlMetastore {
request: IndexMetadataRequest,
) -> MetastoreResult {
let pg_index_opt = if let Some(index_uid) = &request.index_uid {
- index_opt_for_uid(&self.connection_pool, index_uid.clone()).await?
+ index_opt_for_uid(&self.connection_pool, index_uid.clone(), false).await?
} else if let Some(index_id) = &request.index_id {
- index_opt(&self.connection_pool, index_id).await?
+ index_opt(&self.connection_pool, index_id, false).await?
} else {
let message = "invalid request: neither `index_id` nor `index_uid` is set".to_string();
return Err(MetastoreError::Internal {
@@ -676,7 +685,7 @@ impl MetastoreService for PostgresqlMetastore {
let replaced_split_ids = request.replaced_split_ids;
run_with_tx!(self.connection_pool, tx, {
- let mut index_metadata = index_metadata(tx, &index_uid.index_id).await?;
+ let mut index_metadata = index_metadata(tx, &index_uid.index_id, true).await?;
if index_metadata.index_uid != index_uid {
return Err(MetastoreError::NotFound(EntityKind::Index {
index_id: index_uid.index_id,
@@ -684,8 +693,14 @@ impl MetastoreService for PostgresqlMetastore {
}
if let Some(checkpoint_delta) = checkpoint_delta_opt {
let source_id = checkpoint_delta.source_id.clone();
+ let source = index_metadata.sources.get(&source_id).ok_or_else(|| {
+ MetastoreError::NotFound(EntityKind::Source {
+ index_id: index_uid.index_id.to_string(),
+ source_id: source_id.to_string(),
+ })
+ })?;
- if source_id == INGEST_V2_SOURCE_ID {
+ if use_shard_api(&source.source_params) {
let publish_token = request.publish_token_opt.ok_or_else(|| {
let message = format!(
"publish token is required for publishing splits for source \
@@ -933,7 +948,7 @@ impl MetastoreService for PostgresqlMetastore {
.map_err(|sqlx_error| convert_sqlx_err(&index_uid.index_id, sqlx_error))?;
if num_found_splits == 0
- && index_opt(&self.connection_pool, &index_uid.index_id)
+ && index_opt(&self.connection_pool, &index_uid.index_id, false)
.await?
.is_none()
{
@@ -1013,7 +1028,7 @@ impl MetastoreService for PostgresqlMetastore {
.map_err(|sqlx_error| convert_sqlx_err(&index_uid.index_id, sqlx_error))?;
if num_found_splits == 0
- && index_opt_for_uid(&self.connection_pool, index_uid.clone())
+ && index_opt_for_uid(&self.connection_pool, index_uid.clone(), false)
.await?
.is_none()
{
@@ -1209,7 +1224,7 @@ impl MetastoreService for PostgresqlMetastore {
// If no splits were updated, maybe the index does not exist in the first place?
if update_result.rows_affected() == 0
- && index_opt_for_uid(&self.connection_pool, index_uid.clone())
+ && index_opt_for_uid(&self.connection_pool, index_uid.clone(), false)
.await?
.is_none()
{
@@ -1622,6 +1637,7 @@ async fn open_or_fetch_shard<'e>(
.bind(&subrequest.leader_id)
.bind(&subrequest.follower_id)
.bind(subrequest.doc_mapping_uid)
+ .bind(&subrequest.publish_token)
.fetch_optional(executor.clone())
.await?;
@@ -1738,7 +1754,7 @@ impl crate::tests::DefaultForTest for PostgresqlMetastore {
// The number of connections to Postgres should not be
// too catastrophic, as it is limited by the number of concurrent
// unit tests running (= number of test-threads).
- dotenv::dotenv().ok();
+ dotenvy::dotenv().ok();
let uri: Uri = std::env::var("QW_TEST_DATABASE_URL")
.expect("environment variable `QW_TEST_DATABASE_URL` should be set")
.parse()
diff --git a/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql b/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql
index d6747fcb9f1..bd9e9240688 100644
--- a/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql
+++ b/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql
@@ -1,5 +1,5 @@
-INSERT INTO shards(index_uid, source_id, shard_id, leader_id, follower_id, doc_mapping_uid)
- VALUES ($1, $2, $3, $4, $5, $6)
+INSERT INTO shards(index_uid, source_id, shard_id, leader_id, follower_id, doc_mapping_uid, publish_token)
+ VALUES ($1, $2, $3, $4, $5, $6, $7)
ON CONFLICT
DO NOTHING
RETURNING
diff --git a/quickwit/quickwit-metastore/src/tests/shard.rs b/quickwit/quickwit-metastore/src/tests/shard.rs
index 1bbbd1e010a..ed3327774ec 100644
--- a/quickwit/quickwit-metastore/src/tests/shard.rs
+++ b/quickwit/quickwit-metastore/src/tests/shard.rs
@@ -134,6 +134,7 @@ pub async fn test_metastore_open_shards<
leader_id: "test-ingester-foo".to_string(),
follower_id: Some("test-ingester-bar".to_string()),
doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: None,
}],
};
let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap();
@@ -163,6 +164,7 @@ pub async fn test_metastore_open_shards<
leader_id: "test-ingester-foo".to_string(),
follower_id: Some("test-ingester-bar".to_string()),
doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: Some("publish-token-baz".to_string()),
}],
};
let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap();
@@ -181,6 +183,35 @@ pub async fn test_metastore_open_shards<
assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
assert!(shard.publish_token.is_none());
+ // Test open shard #2.
+ let open_shards_request = OpenShardsRequest {
+ subrequests: vec![OpenShardSubrequest {
+ subrequest_id: 0,
+ index_uid: Some(test_index.index_uid.clone()),
+ source_id: test_index.source_id.clone(),
+ shard_id: Some(ShardId::from(2)),
+ leader_id: "test-ingester-foo".to_string(),
+ follower_id: None,
+ doc_mapping_uid: Some(DocMappingUid::default()),
+ publish_token: Some("publish-token-open".to_string()),
+ }],
+ };
+ let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap();
+ assert_eq!(open_shards_response.subresponses.len(), 1);
+
+ let subresponse = &open_shards_response.subresponses[0];
+ assert_eq!(subresponse.subrequest_id, 0);
+
+ let shard = subresponse.open_shard();
+ assert_eq!(shard.index_uid(), &test_index.index_uid);
+ assert_eq!(shard.source_id, test_index.source_id);
+ assert_eq!(shard.shard_id(), ShardId::from(2));
+ assert_eq!(shard.shard_state(), ShardState::Open);
+ assert_eq!(shard.leader_id, "test-ingester-foo");
+ assert!(shard.follower_id.is_none());
+ assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
+ assert_eq!(shard.publish_token(), "publish-token-open");
+
cleanup_index(&mut metastore, test_index.index_uid).await;
}
diff --git a/quickwit/quickwit-metastore/src/tests/split.rs b/quickwit/quickwit-metastore/src/tests/split.rs
index eb8865731f6..ed9d21a36ed 100644
--- a/quickwit/quickwit-metastore/src/tests/split.rs
+++ b/quickwit/quickwit-metastore/src/tests/split.rs
@@ -21,7 +21,7 @@ use std::time::Duration;
use futures::future::try_join_all;
use quickwit_common::rand::append_random_suffix;
-use quickwit_config::IndexConfig;
+use quickwit_config::{IndexConfig, SourceConfig, SourceParams};
use quickwit_proto::metastore::{
CreateIndexRequest, DeleteSplitsRequest, EntityKind, IndexMetadataRequest, ListSplitsRequest,
ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreError, PublishSplitsRequest,
@@ -77,8 +77,10 @@ pub async fn test_metastore_publish_splits_empty_splits_array_is_allowed<
// checkpoint. This operation is allowed and used in the Indexer.
{
let index_config = IndexConfig::for_test(&index_id, &index_uri);
+ let source_configs = &[SourceConfig::for_test(&source_id, SourceParams::void())];
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -133,6 +135,7 @@ pub async fn test_metastore_publish_splits<
let index_config = IndexConfig::for_test(&index_id, &index_uri);
let source_id = format!("{index_id}--source");
+ let source_configs = &[SourceConfig::for_test(&source_id, SourceParams::void())];
let split_id_1 = format!("{index_id}--split-1");
let split_metadata_1 = SplitMetadata {
@@ -199,7 +202,8 @@ pub async fn test_metastore_publish_splits<
// Publish a non-existent split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -226,7 +230,8 @@ pub async fn test_metastore_publish_splits<
// Publish a staged split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -255,7 +260,8 @@ pub async fn test_metastore_publish_splits<
// Publish a published split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -306,7 +312,8 @@ pub async fn test_metastore_publish_splits<
// Publish a non-staged split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -369,7 +376,8 @@ pub async fn test_metastore_publish_splits<
// Publish a staged split and non-existent split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -407,7 +415,8 @@ pub async fn test_metastore_publish_splits<
// Publish a published split and non-existent split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -460,7 +469,8 @@ pub async fn test_metastore_publish_splits<
// Publish a non-staged split and non-existent split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -520,7 +530,8 @@ pub async fn test_metastore_publish_splits<
// Publish staged splits on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -559,7 +570,8 @@ pub async fn test_metastore_publish_splits<
// Publish a staged split and published split on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -617,7 +629,8 @@ pub async fn test_metastore_publish_splits<
// Publish published splits on an index
{
let create_index_request =
- CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs)
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
@@ -681,7 +694,12 @@ pub async fn test_metastore_publish_splits_concurrency<
let index_id = append_random_suffix("test-publish-concurrency");
let index_uri = format!("ram:///indexes/{index_id}");
let index_config = IndexConfig::for_test(&index_id, &index_uri);
- let create_index_request = CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ let source_id = format!("{index_id}--source");
+
+ let source_config = SourceConfig::for_test(&source_id, SourceParams::void());
+ let create_index_request =
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, &[source_config])
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
@@ -690,8 +708,6 @@ pub async fn test_metastore_publish_splits_concurrency<
.index_uid()
.clone();
- let source_id = format!("{index_id}--source");
-
let mut join_handles = Vec::with_capacity(10);
for partition_id in 0..10 {
@@ -1430,6 +1446,7 @@ pub async fn test_metastore_split_update_timestamp<
let index_config = IndexConfig::for_test(&index_id, &index_uri);
let source_id = format!("{index_id}--source");
+ let source_config = SourceConfig::for_test(&source_id, SourceParams::void());
let split_id = format!("{index_id}--split");
let split_metadata = SplitMetadata {
@@ -1440,7 +1457,9 @@ pub async fn test_metastore_split_update_timestamp<
};
// Create an index
- let create_index_request = CreateIndexRequest::try_from_index_config(&index_config).unwrap();
+ let create_index_request =
+ CreateIndexRequest::try_from_index_and_source_configs(&index_config, &[source_config])
+ .unwrap();
let index_uid: IndexUid = metastore
.create_index(create_index_request)
.await
diff --git a/quickwit/quickwit-proto/protos/quickwit/metastore.proto b/quickwit/quickwit-proto/protos/quickwit/metastore.proto
index 4daf5b6d171..b3cd3a7898d 100644
--- a/quickwit/quickwit-proto/protos/quickwit/metastore.proto
+++ b/quickwit/quickwit-proto/protos/quickwit/metastore.proto
@@ -41,6 +41,7 @@ enum SourceType {
SOURCE_TYPE_PULSAR = 9;
SOURCE_TYPE_VEC = 10;
SOURCE_TYPE_VOID = 11;
+ SOURCE_TYPE_STDIN = 13;
}
// Metastore meant to manage Quickwit's indexes, their splits and delete tasks.
@@ -406,6 +407,7 @@ message OpenShardSubrequest {
string leader_id = 5;
optional string follower_id = 6;
quickwit.common.DocMappingUid doc_mapping_uid = 7;
+ optional string publish_token = 8;
}
message OpenShardsResponse {
diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs
index 4b3d88ee7ed..1f0f36db21c 100644
--- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs
+++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs
@@ -338,6 +338,8 @@ pub struct OpenShardSubrequest {
pub follower_id: ::core::option::Option<::prost::alloc::string::String>,
#[prost(message, optional, tag = "7")]
pub doc_mapping_uid: ::core::option::Option,
+ #[prost(string, optional, tag = "8")]
+ pub publish_token: ::core::option::Option<::prost::alloc::string::String>,
}
#[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)]
#[allow(clippy::derive_partial_eq_without_eq)]
@@ -529,6 +531,7 @@ pub enum SourceType {
Pulsar = 9,
Vec = 10,
Void = 11,
+ Stdin = 13,
}
impl SourceType {
/// String value of the enum field names used in the ProtoBuf definition.
@@ -549,6 +552,7 @@ impl SourceType {
SourceType::Pulsar => "SOURCE_TYPE_PULSAR",
SourceType::Vec => "SOURCE_TYPE_VEC",
SourceType::Void => "SOURCE_TYPE_VOID",
+ SourceType::Stdin => "SOURCE_TYPE_STDIN",
}
}
/// Creates an enum from field names used in the ProtoBuf definition.
@@ -566,6 +570,7 @@ impl SourceType {
"SOURCE_TYPE_PULSAR" => Some(Self::Pulsar),
"SOURCE_TYPE_VEC" => Some(Self::Vec),
"SOURCE_TYPE_VOID" => Some(Self::Void),
+ "SOURCE_TYPE_STDIN" => Some(Self::Stdin),
_ => None,
}
}
diff --git a/quickwit/quickwit-proto/src/metastore/mod.rs b/quickwit/quickwit-proto/src/metastore/mod.rs
index 6caba6f7f1d..4782dac03c2 100644
--- a/quickwit/quickwit-proto/src/metastore/mod.rs
+++ b/quickwit/quickwit-proto/src/metastore/mod.rs
@@ -306,6 +306,7 @@ impl SourceType {
SourceType::Nats => "nats",
SourceType::PubSub => "pubsub",
SourceType::Pulsar => "pulsar",
+ SourceType::Stdin => "stdin",
SourceType::Unspecified => "unspecified",
SourceType::Vec => "vec",
SourceType::Void => "void",
@@ -325,6 +326,7 @@ impl fmt::Display for SourceType {
SourceType::Nats => "NATS",
SourceType::PubSub => "Google Cloud Pub/Sub",
SourceType::Pulsar => "Apache Pulsar",
+ SourceType::Stdin => "Stdin",
SourceType::Unspecified => "unspecified",
SourceType::Vec => "vec",
SourceType::Void => "void",
diff --git a/quickwit/quickwit-query/src/query_ast/full_text_query.rs b/quickwit/quickwit-query/src/query_ast/full_text_query.rs
index cf1142e5cf6..2e809cdd476 100644
--- a/quickwit/quickwit-query/src/query_ast/full_text_query.rs
+++ b/quickwit/quickwit-query/src/query_ast/full_text_query.rs
@@ -143,6 +143,12 @@ impl FullTextParams {
Ok(TantivyBoolQuery::build_clause(operator, leaf_queries).into())
}
FullTextMode::Phrase { slop } => {
+ if !index_record_option.has_positions() {
+ return Err(InvalidQuery::SchemaError(
+ "Applied phrase query on field which does not have positions indexed"
+ .to_string(),
+ ));
+ }
let mut phrase_query = TantivyPhraseQuery::new_with_offset(terms);
phrase_query.set_slop(slop);
Ok(phrase_query.into())
diff --git a/quickwit/quickwit-query/src/query_ast/range_query.rs b/quickwit/quickwit-query/src/query_ast/range_query.rs
index 225c26fa815..f8445d4bd8c 100644
--- a/quickwit/quickwit-query/src/query_ast/range_query.rs
+++ b/quickwit/quickwit-query/src/query_ast/range_query.rs
@@ -20,15 +20,16 @@
use std::ops::Bound;
use serde::{Deserialize, Serialize};
-use tantivy::query::{
- FastFieldRangeWeight as TantivyFastFieldRangeQuery, RangeQuery as TantivyRangeQuery,
-};
+use tantivy::fastfield::FastValue;
+use tantivy::query::FastFieldRangeQuery;
use tantivy::schema::Schema as TantivySchema;
-use tantivy::DateTime;
+use tantivy::tokenizer::TextAnalyzer;
+use tantivy::{DateTime, Term};
+use super::tantivy_query_ast::TantivyBoolQuery;
use super::QueryAst;
use crate::json_literal::InterpretUserInput;
-use crate::query_ast::tantivy_query_ast::{TantivyBoolQuery, TantivyQueryAst};
+use crate::query_ast::tantivy_query_ast::TantivyQueryAst;
use crate::query_ast::BuildTantivyAst;
use crate::tokenizers::TokenizerManager;
use crate::{InvalidQuery, JsonLiteral};
@@ -40,129 +41,6 @@ pub struct RangeQuery {
pub upper_bound: Bound,
}
-struct NumericalBoundaries {
- i64_range: (Bound, Bound),
- u64_range: (Bound, Bound),
- f64_range: (Bound, Bound),
-}
-
-fn extract_boundary_value(bound: &Bound) -> Option<&T> {
- match bound {
- Bound::Included(val) | Bound::Excluded(val) => Some(val),
- Bound::Unbounded => None,
- }
-}
-
-trait IntType {
- fn min() -> Self;
- fn max() -> Self;
- fn to_f64(self) -> f64;
- fn from_f64(val: f64) -> Self;
-}
-impl IntType for i64 {
- fn min() -> Self {
- Self::MIN
- }
- fn max() -> Self {
- Self::MAX
- }
- fn to_f64(self) -> f64 {
- self as f64
- }
- fn from_f64(val: f64) -> Self {
- val as Self
- }
-}
-impl IntType for u64 {
- fn min() -> Self {
- Self::MIN
- }
- fn max() -> Self {
- Self::MAX
- }
- fn to_f64(self) -> f64 {
- self as f64
- }
- fn from_f64(val: f64) -> Self {
- val as Self
- }
-}
-
-fn convert_lower_bound<'a, T: IntType + InterpretUserInput<'a>>(
- lower_bound_f64: f64,
- lower_bound: &'a Bound,
-) -> Bound {
- convert_bound(lower_bound).unwrap_or_else(|| {
- if lower_bound_f64 <= T::min().to_f64() {
- // All value should match
- return Bound::Unbounded;
- }
- if lower_bound_f64 > T::max().to_f64() {
- // No values should match
- return Bound::Excluded(T::max());
- }
- // The miss was due to a decimal number.
- Bound::Included(T::from_f64(lower_bound_f64.ceil()))
- })
-}
-
-fn convert_upper_bound<'a, T: IntType + InterpretUserInput<'a>>(
- upper_bound_f64: f64,
- lower_bound: &'a Bound,
-) -> Bound {
- convert_bound(lower_bound).unwrap_or_else(|| {
- if upper_bound_f64 >= T::max().to_f64() {
- // All value should match
- return Bound::Unbounded;
- }
- if upper_bound_f64 < T::min().to_f64() {
- // No values should match
- return Bound::Excluded(T::max());
- }
- // The miss was due to a decimal number.
- Bound::Included(T::from_f64(upper_bound_f64.floor()))
- })
-}
-
-/// This function interprets the lower_bound and upper_bound as numerical boundaries
-/// for JSON field.
-fn compute_numerical_boundaries(
- lower_bound: &Bound,
- upper_bound: &Bound,
-) -> Option {
- // Let's check that this range can be interpret (as in both, bounds),
- // as a numerical range.
- // Note that if interpreting as a f64 range fails, we consider the boundary non
- // numerical and do not attempt to build u64/i64 boundaries.
- let lower_bound_f64: Bound = convert_bound(lower_bound)?;
- let upper_bound_f64: Bound = convert_bound(upper_bound)?;
-
- let lower_bound_i64: Bound;
- let lower_bound_u64: Bound;
- if let Some(lower_bound_f64) = extract_boundary_value(&lower_bound_f64).copied() {
- lower_bound_i64 = convert_lower_bound(lower_bound_f64, lower_bound);
- lower_bound_u64 = convert_lower_bound(lower_bound_f64, lower_bound);
- } else {
- lower_bound_i64 = Bound::Unbounded;
- lower_bound_u64 = Bound::Unbounded;
- }
-
- let upper_bound_i64: Bound;
- let upper_bound_u64: Bound;
- if let Some(upper_bound_f64) = extract_boundary_value(&upper_bound_f64).copied() {
- upper_bound_i64 = convert_upper_bound(upper_bound_f64, upper_bound);
- upper_bound_u64 = convert_upper_bound(upper_bound_f64, upper_bound);
- } else {
- upper_bound_i64 = Bound::Unbounded;
- upper_bound_u64 = Bound::Unbounded;
- }
- Some(NumericalBoundaries {
- i64_range: (lower_bound_i64, upper_bound_i64),
- u64_range: (lower_bound_u64, upper_bound_u64),
- f64_range: (lower_bound_f64, upper_bound_f64),
- })
-}
-
/// Converts a given bound JsonLiteral bound into a bound of type T.
fn convert_bound<'a, T>(bound: &'a Bound) -> Option>
where T: InterpretUserInput<'a> {
@@ -204,19 +82,33 @@ impl From for QueryAst {
}
}
-/// Return
-/// - Some(true) if the range is guaranteed to be empty
-/// - Some(false) if the range is not empty
-/// - None if we cannot judge easily.
-fn is_empty(boundaries: &(Bound, Bound)) -> Option {
- match boundaries {
- (Bound::Included(lower), Bound::Included(upper)) => Some(lower > upper),
- (Bound::Included(lower), Bound::Excluded(upper))
- | (Bound::Excluded(lower), Bound::Included(upper)) => Some(lower >= upper),
- (Bound::Unbounded, Bound::Included(_)) | (Bound::Included(_), Bound::Unbounded) => {
- Some(false)
- }
- _ => None,
+fn term_with_fastval(term: &Term, val: T) -> Term {
+ let mut term = term.clone();
+ term.append_type_and_fast_value(val);
+ term
+}
+
+fn query_from_fast_val_range(
+ empty_term: &Term,
+ range: (Bound, Bound),
+) -> FastFieldRangeQuery {
+ let (lower_bound, upper_bound) = range;
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| term_with_fastval(empty_term, val)),
+ upper_bound.map(|val| term_with_fastval(empty_term, val)),
+ )
+}
+
+fn get_normalized_text(normalizer: &mut Option, text: &str) -> String {
+ if let Some(normalizer) = normalizer {
+ let mut token_stream = normalizer.token_stream(text);
+ let mut tokens = Vec::new();
+ token_stream.process(&mut |token| {
+ tokens.push(token.text.clone());
+ });
+ tokens[0].to_string()
+ } else {
+ text.to_string()
}
}
@@ -224,11 +116,11 @@ impl BuildTantivyAst for RangeQuery {
fn build_tantivy_ast_impl(
&self,
schema: &TantivySchema,
- _tokenizer_manager: &TokenizerManager,
+ tokenizer_manager: &TokenizerManager,
_search_fields: &[String],
_with_validation: bool,
) -> Result {
- let (_field, field_entry, _path) =
+ let (field, field_entry, json_path) =
super::utils::find_field_or_hit_dynamic(&self.field, schema)?;
if !field_entry.is_fast() {
return Err(InvalidQuery::SchemaError(format!(
@@ -237,29 +129,50 @@ impl BuildTantivyAst for RangeQuery {
)));
}
Ok(match field_entry.field_type() {
- tantivy::schema::FieldType::Str(_) => {
- return Err(InvalidQuery::RangeQueryNotSupportedForField {
- value_type: "str",
- field_name: field_entry.name().to_string(),
- });
+ tantivy::schema::FieldType::Str(options) => {
+ let mut normalizer = options
+ .get_fast_field_tokenizer_name()
+ .and_then(|tokenizer_name| tokenizer_manager.get_normalizer(tokenizer_name));
+
+ let (lower_bound, upper_bound) =
+ convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
+
+ FastFieldRangeQuery::new(
+ lower_bound.map(|text| {
+ Term::from_field_text(field, &get_normalized_text(&mut normalizer, text))
+ }),
+ upper_bound.map(|text| {
+ Term::from_field_text(field, &get_normalized_text(&mut normalizer, text))
+ }),
+ )
+ .into()
}
tantivy::schema::FieldType::U64(_) => {
let (lower_bound, upper_bound) =
convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
- TantivyFastFieldRangeQuery::new::(self.field.clone(), lower_bound, upper_bound)
- .into()
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| Term::from_field_u64(field, val)),
+ upper_bound.map(|val| Term::from_field_u64(field, val)),
+ )
+ .into()
}
tantivy::schema::FieldType::I64(_) => {
let (lower_bound, upper_bound) =
convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
- TantivyFastFieldRangeQuery::new::(self.field.clone(), lower_bound, upper_bound)
- .into()
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| Term::from_field_i64(field, val)),
+ upper_bound.map(|val| Term::from_field_i64(field, val)),
+ )
+ .into()
}
tantivy::schema::FieldType::F64(_) => {
let (lower_bound, upper_bound) =
convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
- TantivyFastFieldRangeQuery::new::(self.field.clone(), lower_bound, upper_bound)
- .into()
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| Term::from_field_f64(field, val)),
+ upper_bound.map(|val| Term::from_field_f64(field, val)),
+ )
+ .into()
}
tantivy::schema::FieldType::Bool(_) => {
return Err(InvalidQuery::RangeQueryNotSupportedForField {
@@ -272,12 +185,11 @@ impl BuildTantivyAst for RangeQuery {
convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
let truncate_datetime =
|date: &DateTime| date.truncate(date_options.get_precision());
- let truncated_lower_bound = map_bound(&lower_bound, truncate_datetime);
- let truncated_upper_bound = map_bound(&upper_bound, truncate_datetime);
- TantivyFastFieldRangeQuery::new::(
- self.field.clone(),
- truncated_lower_bound,
- truncated_upper_bound,
+ let lower_bound = map_bound(&lower_bound, truncate_datetime);
+ let upper_bound = map_bound(&upper_bound, truncate_datetime);
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| Term::from_field_date(field, val)),
+ upper_bound.map(|val| Term::from_field_date(field, val)),
)
.into()
}
@@ -288,44 +200,61 @@ impl BuildTantivyAst for RangeQuery {
});
}
tantivy::schema::FieldType::Bytes(_) => todo!(),
- tantivy::schema::FieldType::JsonObject(_) => {
- let full_path = self.field.clone();
+ tantivy::schema::FieldType::JsonObject(options) => {
let mut sub_queries: Vec = Vec::new();
- if let Some(NumericalBoundaries {
- i64_range,
- u64_range,
- f64_range,
- }) = compute_numerical_boundaries(&self.lower_bound, &self.upper_bound)
- {
- // Adding the f64 range.
- sub_queries.push(
- TantivyFastFieldRangeQuery::new(
- full_path.clone(),
- f64_range.0,
- f64_range.1,
- )
- .into(),
- );
- // Adding the i64 range.
- if !is_empty(&i64_range).unwrap_or(false) {
- sub_queries.push(
- TantivyFastFieldRangeQuery::new(
- full_path.clone(),
- i64_range.0,
- i64_range.1,
- )
- .into(),
- );
- }
- // Adding the u64 range.
- if !is_empty(&u64_range).unwrap_or(false) {
- sub_queries.push(
- TantivyFastFieldRangeQuery::new(full_path, u64_range.0, u64_range.1)
- .into(),
- );
- }
+ let empty_term =
+ Term::from_field_json_path(field, json_path, options.is_expand_dots_enabled());
+ // Try to convert the bounds into numerical values in following order i64, u64,
+ // f64. Tantivy will convert to the correct numerical type of the column if it
+ // doesn't match.
+ let bounds_range_i64: Option<(Bound, Bound)> =
+ convert_bound(&self.lower_bound).zip(convert_bound(&self.upper_bound));
+ let bounds_range_u64: Option<(Bound, Bound)> =
+ convert_bound(&self.lower_bound).zip(convert_bound(&self.upper_bound));
+ let bounds_range_f64: Option<(Bound, Bound)> =
+ convert_bound(&self.lower_bound).zip(convert_bound(&self.upper_bound));
+ if let Some(range) = bounds_range_i64 {
+ sub_queries.push(query_from_fast_val_range(&empty_term, range).into());
+ } else if let Some(range) = bounds_range_u64 {
+ sub_queries.push(query_from_fast_val_range(&empty_term, range).into());
+ } else if let Some(range) = bounds_range_f64 {
+ sub_queries.push(query_from_fast_val_range(&empty_term, range).into());
}
- // TODO add support for str range queries.
+
+ let mut normalizer = options
+ .get_fast_field_tokenizer_name()
+ .and_then(|tokenizer_name| tokenizer_manager.get_normalizer(tokenizer_name));
+
+ let bounds_range_str: Option<(Bound<&str>, Bound<&str>)> =
+ convert_bound(&self.lower_bound).zip(convert_bound(&self.upper_bound));
+ if let Some(range) = bounds_range_str {
+ let str_query = FastFieldRangeQuery::new(
+ range.0.map(|val| {
+ let val = get_normalized_text(&mut normalizer, val);
+ let mut term = empty_term.clone();
+ term.append_type_and_str(&val);
+ term
+ }),
+ range.1.map(|val| {
+ let val = get_normalized_text(&mut normalizer, val);
+ let mut term = empty_term.clone();
+ term.append_type_and_str(&val);
+ term
+ }),
+ )
+ .into();
+ sub_queries.push(str_query);
+ }
+ if sub_queries.is_empty() {
+ return Err(InvalidQuery::InvalidBoundary {
+ expected_value_type: "i64, u64, f64, str",
+ field_name: field_entry.name().to_string(),
+ });
+ }
+ if sub_queries.len() == 1 {
+ return Ok(sub_queries.pop().unwrap());
+ }
+
let bool_query = TantivyBoolQuery {
should: sub_queries,
..Default::default()
@@ -335,8 +264,11 @@ impl BuildTantivyAst for RangeQuery {
tantivy::schema::FieldType::IpAddr(_) => {
let (lower_bound, upper_bound) =
convert_bounds(&self.lower_bound, &self.upper_bound, field_entry.name())?;
- TantivyRangeQuery::new_ip_bounds(self.field.clone(), lower_bound, upper_bound)
- .into()
+ FastFieldRangeQuery::new(
+ lower_bound.map(|val| Term::from_field_ip_addr(field, val)),
+ upper_bound.map(|val| Term::from_field_ip_addr(field, val)),
+ )
+ .into()
}
})
}
@@ -357,7 +289,6 @@ mod tests {
use tantivy::schema::{DateOptions, DateTimePrecision, Schema, FAST, STORED, TEXT};
use super::RangeQuery;
- use crate::query_ast::tantivy_query_ast::TantivyBoolQuery;
use crate::query_ast::BuildTantivyAst;
use crate::{
create_default_quickwit_tokenizer_manager, InvalidQuery, JsonLiteral, MatchAllOrNone,
@@ -412,24 +343,22 @@ mod tests {
"my_i64_field",
JsonLiteral::String("1980".to_string()),
JsonLiteral::String("1989".to_string()),
- "FastFieldRangeWeight { field: \"my_i64_field\", lower_bound: \
- Included(9223372036854777788), upper_bound: Included(9223372036854777797), \
- column_type_opt: Some(I64) }",
+ "FastFieldRangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=0, \
+ type=I64, 1980)), upper_bound: Included(Term(field=0, type=I64, 1989)) } }",
);
test_range_query_typed_field_util(
"my_u64_field",
JsonLiteral::String("1980".to_string()),
JsonLiteral::String("1989".to_string()),
- "FastFieldRangeWeight { field: \"my_u64_field\", lower_bound: Included(1980), \
- upper_bound: Included(1989), column_type_opt: Some(U64) }",
+ "FastFieldRangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=1, \
+ type=U64, 1980)), upper_bound: Included(Term(field=1, type=U64, 1989)) } }",
);
test_range_query_typed_field_util(
"my_f64_field",
JsonLiteral::String("1980".to_string()),
JsonLiteral::String("1989".to_string()),
- "FastFieldRangeWeight { field: \"my_f64_field\", lower_bound: \
- Included(13879794984393113600), upper_bound: Included(13879834566811713536), \
- column_type_opt: Some(F64) }",
+ "FastFieldRangeQuery { bounds: BoundsRange { lower_bound: Included(Term(field=2, \
+ type=F64, 1980.0)), upper_bound: Included(Term(field=2, type=F64, 1989.0)) } }",
);
}
@@ -468,42 +397,6 @@ mod tests {
);
}
- #[test]
- fn test_range_query_field_unsupported_type_field() {
- let schema = make_schema(false);
- let range_query = RangeQuery {
- field: "my_str_field".to_string(),
- lower_bound: Bound::Included(JsonLiteral::String("1980".to_string())),
- upper_bound: Bound::Included(JsonLiteral::String("1989".to_string())),
- };
- // with validation
- let invalid_query: InvalidQuery = range_query
- .build_tantivy_ast_call(
- &schema,
- &create_default_quickwit_tokenizer_manager(),
- &[],
- true,
- )
- .unwrap_err();
- assert!(matches!(
- invalid_query,
- InvalidQuery::RangeQueryNotSupportedForField { .. }
- ));
- // without validation
- assert_eq!(
- range_query
- .build_tantivy_ast_call(
- &schema,
- &create_default_quickwit_tokenizer_manager(),
- &[],
- false
- )
- .unwrap()
- .const_predicate(),
- Some(MatchAllOrNone::MatchNone)
- );
- }
-
#[test]
fn test_range_dynamic() {
let range_query = RangeQuery {
@@ -520,35 +413,15 @@ mod tests {
true,
)
.unwrap();
- let TantivyBoolQuery {
- must,
- must_not,
- should,
- filter,
- } = tantivy_ast.as_bool_query().unwrap();
- assert!(must.is_empty());
- assert!(must_not.is_empty());
- assert!(filter.is_empty());
- assert_eq!(should.len(), 3);
- let range_queries: Vec<&dyn tantivy::query::Query> = should
- .iter()
- .map(|ast| ast.as_leaf().unwrap())
- .collect::>();
- assert_eq!(
- format!("{:?}", range_queries[0]),
- "FastFieldRangeWeight { field: \"hello\", lower_bound: \
- Included(13879794984393113600), upper_bound: Included(13879834566811713536), \
- column_type_opt: Some(F64) }"
- );
- assert_eq!(
- format!("{:?}", range_queries[1]),
- "FastFieldRangeWeight { field: \"hello\", lower_bound: Included(9223372036854777788), \
- upper_bound: Included(9223372036854777797), column_type_opt: Some(I64) }"
- );
assert_eq!(
- format!("{:?}", range_queries[2]),
- "FastFieldRangeWeight { field: \"hello\", lower_bound: Included(1980), upper_bound: \
- Included(1989), column_type_opt: Some(U64) }"
+ format!("{:?}", tantivy_ast),
+ "Bool(TantivyBoolQuery { must: [], must_not: [], should: [Leaf(FastFieldRangeQuery { \
+ bounds: BoundsRange { lower_bound: Included(Term(field=6, type=Json, path=hello, \
+ type=I64, 1980)), upper_bound: Included(Term(field=6, type=Json, path=hello, \
+ type=I64, 1989)) } }), Leaf(FastFieldRangeQuery { bounds: BoundsRange { lower_bound: \
+ Included(Term(field=6, type=Json, path=hello, type=Str, \"1980\")), upper_bound: \
+ Included(Term(field=6, type=Json, path=hello, type=Str, \"1989\")) } })], filter: [] \
+ })"
);
}
diff --git a/quickwit/quickwit-query/src/query_ast/utils.rs b/quickwit/quickwit-query/src/query_ast/utils.rs
index 593ac978ef8..be2bdf2f3f8 100644
--- a/quickwit/quickwit-query/src/query_ast/utils.rs
+++ b/quickwit/quickwit-query/src/query_ast/utils.rs
@@ -181,7 +181,7 @@ fn compute_tantivy_ast_query_for_json(
) -> Result {
let mut bool_query = TantivyBoolQuery::default();
let term = Term::from_field_json_path(field, json_path, json_options.is_expand_dots_enabled());
- if let Some(term) = convert_to_fast_value_and_append_to_json_term(term, text) {
+ if let Some(term) = convert_to_fast_value_and_append_to_json_term(term, text, true) {
bool_query
.should
.push(TantivyTermQuery::new(term, IndexRecordOption::Basic).into());
diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs
index b225a53c13d..e9e923c0a48 100644
--- a/quickwit/quickwit-search/src/leaf.rs
+++ b/quickwit/quickwit-search/src/leaf.rs
@@ -44,6 +44,7 @@ use tantivy::directory::FileSlice;
use tantivy::fastfield::FastFieldReaders;
use tantivy::schema::Field;
use tantivy::{DateTime, Index, ReloadPolicy, Searcher, Term};
+use tokio::task::JoinError;
use tracing::*;
use crate::collector::{make_collector_for_split, make_merge_collector, IncrementalCollector};
@@ -1217,7 +1218,8 @@ pub async fn leaf_search(
let split_filter = Arc::new(RwLock::new(split_filter));
- let mut leaf_search_single_split_futures: Vec<_> = Vec::with_capacity(split_with_req.len());
+ let mut leaf_search_single_split_join_handles: Vec<(String, tokio::task::JoinHandle<()>)> =
+ Vec::with_capacity(split_with_req.len());
let merge_collector = make_merge_collector(&request, &aggregations_limits)?;
let incremental_merge_collector = IncrementalCollector::new(merge_collector);
@@ -1236,26 +1238,45 @@ pub async fn leaf_search(
continue;
}
- leaf_search_single_split_futures.push(tokio::spawn(
- leaf_search_single_split_wrapper(
- request,
- searcher_context.clone(),
- index_storage.clone(),
- doc_mapper.clone(),
- split,
- split_filter.clone(),
- incremental_merge_collector.clone(),
- leaf_split_search_permit,
- aggregations_limits.clone(),
- )
- .in_current_span(),
+ leaf_search_single_split_join_handles.push((
+ split.split_id.clone(),
+ tokio::spawn(
+ leaf_search_single_split_wrapper(
+ request,
+ searcher_context.clone(),
+ index_storage.clone(),
+ doc_mapper.clone(),
+ split,
+ split_filter.clone(),
+ incremental_merge_collector.clone(),
+ leaf_split_search_permit,
+ aggregations_limits.clone(),
+ )
+ .in_current_span(),
+ ),
));
}
// TODO we could cancel running splits when !run_all_splits and the running split can no
// longer give better results after some other split answered.
- let split_search_results: Vec> =
- futures::future::join_all(leaf_search_single_split_futures).await;
+ let mut split_search_join_errors: Vec<(String, JoinError)> = Vec::new();
+
+ // There is no need to use `join_all`, as these are spawned tasks.
+ for (split, leaf_search_join_handle) in leaf_search_single_split_join_handles {
+ // splits that did not panic were already added to the collector
+ if let Err(join_error) = leaf_search_join_handle.await {
+ if join_error.is_cancelled() {
+ // An explicit task cancellation is not an error.
+ continue;
+ }
+ if join_error.is_panic() {
+ error!(split=%split, "leaf search task panicked");
+ } else {
+ error!(split=%split, "please report: leaf search was not cancelled, and could not extract panic. this should never happen");
+ }
+ split_search_join_errors.push((split, join_error));
+ }
+ }
// we can't use unwrap_or_clone because mutexes aren't Clone
let mut incremental_merge_collector = match Arc::try_unwrap(incremental_merge_collector) {
@@ -1263,17 +1284,12 @@ pub async fn leaf_search(
Err(filter_merger) => filter_merger.lock().unwrap().clone(),
};
- for result in split_search_results {
- // splits that did not panic were already added to the collector
- if let Err(e) = result {
- incremental_merge_collector.add_failed_split(SplitSearchError {
- // we could reasonably add a wrapper to the JoinHandle to give us the
- // split_id anyway
- split_id: "unknown".to_string(),
- error: format!("{}", SearchError::from(e)),
- retryable_error: true,
- })
- }
+ for (split_id, split_search_join_error) in split_search_join_errors {
+ incremental_merge_collector.add_failed_split(SplitSearchError {
+ split_id,
+ error: SearchError::from(split_search_join_error).to_string(),
+ retryable_error: true,
+ });
}
let result = crate::search_thread_pool()
diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs
index 778911e8d10..827c2b027ca 100644
--- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs
+++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs
@@ -23,7 +23,8 @@ use bytes::Bytes;
use quickwit_common::uri::Uri;
use quickwit_config::{
load_index_config_update, load_source_config_from_user_config, validate_index_id_pattern,
- ConfigFormat, NodeConfig, SourceConfig, SourceParams, CLI_SOURCE_ID, INGEST_API_SOURCE_ID,
+ ConfigFormat, FileSourceParams, NodeConfig, SourceConfig, SourceParams, CLI_SOURCE_ID,
+ INGEST_API_SOURCE_ID,
};
use quickwit_doc_mapper::{analyze_text, TokenizerConfig};
use quickwit_index_management::{IndexService, IndexServiceError};
@@ -708,12 +709,15 @@ async fn create_source(
let source_config: SourceConfig =
load_source_config_from_user_config(config_format, &source_config_bytes)
.map_err(IndexServiceError::InvalidConfig)?;
- if let SourceParams::File(_) = &source_config.source_params {
- return Err(IndexServiceError::OperationNotAllowed(
- "file sources are limited to a local usage. please use the CLI command `quickwit tool \
- local-ingest` to ingest data from a file"
- .to_string(),
- ));
+ // Note: This check is performed here instead of the source config serde
+ // because many tests use the file source, and can't store that config in
+ // the metastore without going through the validation.
+ if let SourceParams::File(FileSourceParams::Filepath(_)) = &source_config.source_params {
+ return Err(IndexServiceError::InvalidConfig(anyhow::anyhow!(
+ "path based file sources are limited to a local usage, please use the CLI command \
+ `quickwit tool local-ingest` to ingest data from a specific file or setup a \
+ notification based file source"
+ )));
}
let index_metadata_request = IndexMetadataRequest::for_index_id(index_id.to_string());
let index_uid: IndexUid = index_service
@@ -1641,28 +1645,6 @@ mod tests {
assert!(indexes.is_empty());
}
- #[tokio::test]
- async fn test_create_file_source_returns_403() {
- let metastore = metastore_for_test();
- let index_service = IndexService::new(metastore.clone(), StorageResolver::unconfigured());
- let mut node_config = NodeConfig::for_test();
- node_config.default_index_root_uri = Uri::for_test("file:///default-index-root-uri");
- let index_management_handler =
- super::index_management_handlers(index_service, Arc::new(node_config))
- .recover(recover_fn);
- let source_config_body = r#"{"version": "0.7", "source_id": "file-source", "source_type":
- "file", "params": {"filepath": "FILEPATH"}}"#;
- let resp = warp::test::request()
- .path("/indexes/hdfs-logs/sources")
- .method("POST")
- .body(source_config_body)
- .reply(&index_management_handler)
- .await;
- assert_eq!(resp.status(), 403);
- let response_body = std::str::from_utf8(resp.body()).unwrap();
- assert!(response_body.contains("limited to a local usage"))
- }
-
#[tokio::test]
async fn test_create_index_with_yaml() {
let metastore = metastore_for_test();
@@ -1886,6 +1868,34 @@ mod tests {
sources"
));
}
+ {
+ let resp = warp::test::request()
+ .path("/indexes/hdfs-logs/sources")
+ .method("POST")
+ .body(
+ r#"{"version": "0.8", "source_id": "my-stdin-source", "source_type": "stdin"}"#,
+ )
+ .reply(&index_management_handler)
+ .await;
+ assert_eq!(resp.status(), 400);
+ let response_body = std::str::from_utf8(resp.body()).unwrap();
+ assert!(
+ response_body.contains("stdin can only be used as source through the CLI command")
+ )
+ }
+ {
+ let resp = warp::test::request()
+ .path("/indexes/hdfs-logs/sources")
+ .method("POST")
+ .body(
+ r#"{"version": "0.8", "source_id": "my-local-file-source", "source_type": "file", "params": {"filepath": "localfile"}}"#,
+ )
+ .reply(&index_management_handler)
+ .await;
+ assert_eq!(resp.status(), 400);
+ let response_body = std::str::from_utf8(resp.body()).unwrap();
+ assert!(response_body.contains("limited to a local usage"))
+ }
}
#[tokio::test]
diff --git a/quickwit/quickwit-storage/src/storage.rs b/quickwit/quickwit-storage/src/storage.rs
index 9ca1dd2041f..9e8f1c54000 100644
--- a/quickwit/quickwit-storage/src/storage.rs
+++ b/quickwit/quickwit-storage/src/storage.rs
@@ -103,7 +103,9 @@ pub trait Storage: fmt::Debug + Send + Sync + 'static {
/// Downloads a slice of a file from the storage, and returns an in memory buffer
async fn get_slice(&self, path: &Path, range: Range) -> StorageResult;
- /// Open a stream handle on the file from the storage
+ /// Opens a stream handle on the file from the storage.
+ ///
+ /// Might panic, return an error or an empty stream if the range is empty.
async fn get_slice_stream(
&self,
path: &Path,
diff --git a/quickwit/rest-api-tests/scenarii/aggregations/0001-aggregations.yaml b/quickwit/rest-api-tests/scenarii/aggregations/0001-aggregations.yaml
index 1df59efd2c8..f81c2215f40 100644
--- a/quickwit/rest-api-tests/scenarii/aggregations/0001-aggregations.yaml
+++ b/quickwit/rest-api-tests/scenarii/aggregations/0001-aggregations.yaml
@@ -375,4 +375,21 @@ expected:
aggregations:
response_stats:
sum_of_squares: 55300.0
-
+# Test term aggs number precision
+method: [GET]
+engines:
+ - quickwit
+endpoint: _elastic/aggregations/_search
+json:
+ query: { match_all: {} }
+ size: 0
+ aggs:
+ names:
+ terms:
+ field: "high_prec_test"
+expected:
+ aggregations:
+ names:
+ buckets:
+ - doc_count: 1
+ key: 1769070189829214200
diff --git a/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml b/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml
index e98e59a156c..08b437015b9 100644
--- a/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml
+++ b/quickwit/rest-api-tests/scenarii/aggregations/0002-doc-len.yaml
@@ -12,7 +12,7 @@ json:
expected:
aggregations:
doc_len:
- value: 913.0
+ value: 952.0
---
# Test doc len isn't shown when querying documents
method: [GET]
diff --git a/quickwit/rest-api-tests/scenarii/aggregations/_setup.quickwit.yaml b/quickwit/rest-api-tests/scenarii/aggregations/_setup.quickwit.yaml
index 11ee82ec67f..60580ca576a 100644
--- a/quickwit/rest-api-tests/scenarii/aggregations/_setup.quickwit.yaml
+++ b/quickwit/rest-api-tests/scenarii/aggregations/_setup.quickwit.yaml
@@ -25,6 +25,9 @@ json:
- rfc3339
fast_precision: seconds
fast: true
+ - name: high_prec_test
+ type: u64
+ fast: true
store_document_size: true
---
# Create empty index
@@ -64,7 +67,7 @@ endpoint: aggregations/ingest
params:
commit: force
ndjson:
- - {"name": "Fritz", "response": 30, "id": 0}
+ - {"name": "Fritz", "high_prec_test": 1769070189829214200, "response": 30, "id": 0}
- {"name": "Fritz", "response": 30, "id": 0}
- {"name": "Holger", "response": 30, "id": 4, "date": "2015-02-06T00:00:00Z", "host": "192.168.0.10"}
- {"name": "Werner", "response": 20, "id": 5, "date": "2015-01-02T00:00:00Z", "host": "192.168.0.10"}
diff --git a/quickwit/rest-api-tests/scenarii/es_compatibility/0007-range_queries.yaml b/quickwit/rest-api-tests/scenarii/es_compatibility/0007-range_queries.yaml
index c3c625395c4..5337325c229 100644
--- a/quickwit/rest-api-tests/scenarii/es_compatibility/0007-range_queries.yaml
+++ b/quickwit/rest-api-tests/scenarii/es_compatibility/0007-range_queries.yaml
@@ -187,3 +187,61 @@ expected:
total:
value: 1
relation: "eq"
+---
+# This field is not a JSON field and doesn not have fast field normalization.
+# That means it is case sensitive
+json:
+ query:
+ range:
+ repo.name:
+ gte: "h"
+ lte: "z"
+expected:
+ hits:
+ total:
+ value: 62
+ relation: "eq"
+---
+# This field is a JSON field and has fast field normalization.
+# That means it is case insensitive
+json:
+ query:
+ range:
+ actor.login:
+ gte: "H" # should automatically be normalized
+ lte: "Z"
+expected:
+ hits:
+ total:
+ value: 68
+ relation: "eq"
+---
+# This field is a JSON field and has fast field normalization.
+# That means it is case insensitive
+json:
+ query:
+ range:
+ actor.login:
+ gte: "h" # should automatically be normalized
+ lte: "z"
+expected:
+ hits:
+ total:
+ value: 68
+ relation: "eq"
+---
+# This field is a JSON field and has fast field normalization.
+# That means it is case insensitive
+json:
+ query:
+ range:
+ actor.login:
+ gte: "H" # should automatically be normalized
+ lte: "Z"
+expected:
+ hits:
+ total:
+ value: 68
+ relation: "eq"
+
+
diff --git a/quickwit/rest-api-tests/scenarii/es_compatibility/_setup.quickwit.yaml b/quickwit/rest-api-tests/scenarii/es_compatibility/_setup.quickwit.yaml
index bad0012b553..ec6e9f81a3d 100644
--- a/quickwit/rest-api-tests/scenarii/es_compatibility/_setup.quickwit.yaml
+++ b/quickwit/rest-api-tests/scenarii/es_compatibility/_setup.quickwit.yaml
@@ -36,6 +36,10 @@ json:
timestamp_field: created_at
mode: dynamic
field_mappings:
+ - name: repo.name
+ type: text
+ fast: true
+ indexed: true
- name: actor.id
type: u64
fast: true