Incompatibility with data streams

[timo-obrecht]

Current class definition of an index is incompatible with data streams. It does not allow to initialize the data stream from the class, nor to access the usual methods and utils (init(), save(), bulk(), etc)

Current way of handling data streams :

from elasticsearch import Elasticsearch
from elasticsearch.dsl import Date, Document, Float
from testcontainers.elasticsearch import ElasticSearchContainer

PASSWORD = "test-password-123"
ES_IMAGE = "docker.elastic.co/elasticsearch/elasticsearch:9.3.3"


class Foo(Document):
    data = Float()
    timestamp = Date(name="@timestamp")  # Python attr -> ES field

    class Index:
        name = "foo"


with (
    ElasticSearchContainer(ES_IMAGE, mem_limit="2G")
    .with_env("discovery.type", "single-node")
    .with_env("ELASTIC_PASSWORD", PASSWORD) as es
):
    url = f"http://{es.get_container_host_ip()}:{es.get_exposed_port(9200)}"

    with Elasticsearch(
        url,
        basic_auth=("elastic", PASSWORD),
        verify_certs=False,
        request_timeout=120,
    ) as client:
        client.cluster.health(wait_for_status="yellow", timeout="30s")

        print("1) DSL init() does not support data streams")
        try:
            Foo.init(using=client)
            print("UNEXPECTED: Foo.init() succeeded")
        except Exception as exc:
            print(f"Expected failure: {exc!r}")

        print("\n2) Create the data stream manually")
        client.indices.put_index_template(
            name="foo-template",
            index_patterns=[Foo.Index.name],
            data_stream={},
            template={
                "mappings": {
                    "properties": {
                        "@timestamp": {"type": "date"},
                        "data": {"type": "float"},
                    }
                }
            },
        )
        client.indices.create_data_stream(name=Foo.Index.name)

        doc = Foo(timestamp="2024-01-01T00:00:00+00:00", data=0.01)
        payload = doc.to_dict()

        print("\n3) BUG: DSL serializes `timestamp`, not `@timestamp`")
        print(payload)

        print("\n4) save() fails because the payload is not valid for a data stream")
        try:
            doc.save(using=client)
            print("UNEXPECTED: save() succeeded")
        except Exception as exc:
            print(f"Expected failure: {exc!r}")

        payload["@timestamp"] = payload.pop("timestamp")
        payload["_op_type"] = "create"

        print("\n5) Manual fix works: rename field and use op_type=create")
        Foo.bulk(using=client, actions=[payload])
        print("Bulk indexing succeeded")

I guess simply allowing DataStream in the class definition and making sure .to_dict() correctly set '@timestamp' as the key would allow a cleaner syntax like so :

class Foo(Document):
    data = Float()
    timestamp = Date(name="@timestamp")  # Python attr -> ES field

    class DataStream:
        name = "foo"

with (
    ElasticSearchContainer(ES_IMAGE)
    .with_env("discovery.type", "single-node")
    .with_env("ELASTIC_PASSWORD", PASSWORD) as es
):
    url = f"http://{es.get_container_host_ip()}:{es.get_exposed_port(9200)}"

    with Elasticsearch(
        url,
        basic_auth=("elastic", PASSWORD),
        verify_certs=False,
    ) as client:
        client.cluster.health(wait_for_status="yellow", timeout="30s")

        Foo.init(using=client)
        
        doc = Foo(timestamp="2024-01-01T00:00:00+00:00", data=0.01)
        doc.save(using=client)

[miguelgrinberg]

Thanks for reaching out! This is indeed something that is important enough for the Document class to support.

[miguelgrinberg]

I have spent some time with this issue and have a preliminary solution in place.

Here is how you define a document class for a data stream:

class Foo(dsl.Document):
    timestamp: datetime = dsl.mapped_field(es_name="@timestamp")
    data: float

    class Index:
        name = "foo"
        data_stream = True

If you prefer to use the older type mappings, then it is done like this:

class Foo(dsl.Document):
    timestamp = dsl.Date(_es_name="@timestamp")
    data = dsl.Float()

    class Index:
        name = "foo"
        data_stream = True

There are two changes here. First, I've added a es_name/_es_name arguments in the type definition, which allows you to define an Elasticsearch name for the attribute that is different than the one you use in the Python class. With this you can create a timestamp attribute that will be serialized as @timestamp.

The second change is the addition of the data_stream = True attribute in the Index definition. This tells us that the index template for this index needs to be marked as a data stream template.

Unfortunately I have not got around to improve the Foo.init() function to work with data streams, so when you intend to work on a data stream, instead of calling init() you can create your data stream via direct API calls. But you can take advantage of the document's template generation:

template = Foo._index.as_composable_template("foo-template", "foo")  # arguments are the template name and the pattern
template.save()
dsl.connections.get_connection().indices.create_data_stream(name="foo")

And now you are good to go. You can either use the document instance to save entries to the data stream:

doc = Foo(timestamp=datetime.now(), data=random.random() * 10000)
doc.save()

Or you can use the bulk API:

def get_next_doc():
    for i in range(100000):
        yield {"_source": Foo(timestamp=datetime.now(), data=random.random() * 10000), "_op_type": "create"}

Foo.bulk(get_next_doc())

The changes that enable all of this are in this PR. We are actually very close to a release, and I need to decide if I slip this one in or if I keep working on this for a more complete integration to be released later, but aside from that you are welcome to play with the code in the PR and provide feedback.