def test_search_body(instrument, elasticapm_client, elasticsearch_async):
    await elasticsearch_async.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    search_query = {"query": {"term": {"user": "kimchy"}}}
    result = await elasticsearch_async.search(body=search_query, params=None)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]
    assert result["hits"]["hits"][0]["_source"] == {"user": "kimchy", "text": "hola"}
    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 1
    span = spans[0]
    # Depending on ES_VERSION, could be /_all/_search or /_search, and GET or POST
    assert span["name"] in ("ES GET /_search", "ES GET /_all/_search", "ES POST /_search")
    assert span["type"] == "db"
    assert span["subtype"] == "elasticsearch"
    assert span["action"] == "query"
    assert span["context"]["db"]["type"] == "elasticsearch"
    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}'
    assert span["sync"] is False 

def test_search_body(instrument, elasticapm_client, async_elasticsearch):
    await async_elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    search_query = {"query": {"term": {"user": "kimchy"}}}
    result = await async_elasticsearch.search(body=search_query, params=None)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]
    assert result["hits"]["hits"][0]["_source"] == {"user": "kimchy", "text": "hola"}
    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 1
    span = spans[0]
    # Depending on ES_VERSION, could be /_all/_search or /_search, and GET or POST
    assert span["name"] in ("ES GET /_search", "ES GET /_all/_search", "ES POST /_search")
    assert span["type"] == "db"
    assert span["subtype"] == "elasticsearch"
    assert span["action"] == "query"
    assert span["context"]["db"]["type"] == "elasticsearch"
    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}'
    assert span["sync"] is False 

def test_exists_source(instrument, elasticapm_client, elasticsearch):
    elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    if ES_VERSION[0] < 7:
        assert elasticsearch.exists_source("tweets", document_type, 1) is True
    else:
        assert elasticsearch.exists_source(index="tweets", id=1, doc_type=document_type) is True
    assert elasticsearch.exists_source(index="tweets", doc_type=document_type, id=1) is True
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]

    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for span in spans:
        assert span["name"] == "ES HEAD /tweets/%s/1/_source" % document_type
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def test_search_body(instrument, elasticapm_client, elasticsearch):
    elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    search_query = {"query": {"term": {"user": "kimchy"}}}
    result = elasticsearch.search(body=search_query, params=None)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]
    assert result["hits"]["hits"][0]["_source"] == {"user": "kimchy", "text": "hola"}
    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 1
    span = spans[0]
    # Depending on ES_VERSION, could be /_all/_search or /_search, and GET or POST
    assert span["name"] in ("ES GET /_search", "ES GET /_all/_search", "ES POST /_search")
    assert span["type"] == "db"
    assert span["subtype"] == "elasticsearch"
    assert span["action"] == "query"
    assert span["context"]["db"]["type"] == "elasticsearch"
    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}' 

def test_count_body(instrument, elasticapm_client, elasticsearch):
    elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    search_query = {"query": {"term": {"user": "kimchy"}}}
    result = elasticsearch.count(body=search_query)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]
    assert result["count"] == 1
    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 1
    span = spans[0]
    # Depending on ES_VERSION, could be /_all/_count or /_count, and either GET
    # or POST. None of these details actually matter much for this test.
    # Technically no version does `POST /_all/_count` but I added it anyway
    assert span["name"] in ("ES GET /_count", "ES GET /_all/_count", "ES POST /_count", "ES POST /_all/_count")
    assert span["type"] == "db"
    assert span["subtype"] == "elasticsearch"
    assert span["action"] == "query"
    assert span["context"]["db"]["type"] == "elasticsearch"
    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}' 

def test_create(instrument, elasticapm_client, elasticsearch_async):
    elasticapm_client.begin_transaction("test")
    if ES_VERSION[0] < 5:
        r1 = await elasticsearch_async.create("tweets", document_type, {"user": "kimchy", "text": "hola"}, 1)
    elif ES_VERSION[0] < 7:
        r1 = await elasticsearch_async.create("tweets", document_type, 1, body={"user": "kimchy", "text": "hola"})
    else:
        r1 = await elasticsearch_async.create("tweets", 1, body={"user": "kimchy", "text": "hola"})
    r2 = await elasticsearch_async.create(
        index="tweets", doc_type=document_type, id=2, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]

    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for i, span in enumerate(spans):
        if ES_VERSION[0] >= 5:
            assert span["name"] in (
                "ES PUT /tweets/%s/%d/_create" % (document_type, i + 1),
                "ES PUT /tweets/_create/%d" % (i + 1),
            )
        else:
            assert span["name"] == "ES PUT /tweets/%s/%d" % (document_type, i + 1)
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def test_create(instrument, elasticapm_client, async_elasticsearch):
    elasticapm_client.begin_transaction("test")
    if ES_VERSION[0] < 5:
        r1 = await async_elasticsearch.create("tweets", document_type, {"user": "kimchy", "text": "hola"}, 1)
    elif ES_VERSION[0] < 7:
        r1 = await async_elasticsearch.create("tweets", document_type, 1, body={"user": "kimchy", "text": "hola"})
    else:
        r1 = await async_elasticsearch.create("tweets", 1, body={"user": "kimchy", "text": "hola"})
    r2 = await async_elasticsearch.create(
        index="tweets", doc_type=document_type, id=2, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]

    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for i, span in enumerate(spans):
        if ES_VERSION[0] >= 5:
            assert span["name"] in (
                "ES PUT /tweets/%s/%d/_create" % (document_type, i + 1),
                "ES PUT /tweets/_create/%d" % (i + 1),
            )
        else:
            assert span["name"] == "ES PUT /tweets/%s/%d" % (document_type, i + 1)
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def test_create(instrument, elasticapm_client, elasticsearch):
    elasticapm_client.begin_transaction("test")
    if ES_VERSION[0] < 7:
        r1 = elasticsearch.create(index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"})
    else:
        r1 = elasticsearch.create(index="tweets", id=1, body={"user": "kimchy", "text": "hola"})
    r2 = elasticsearch.create(
        index="tweets", doc_type=document_type, id=2, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]

    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for i, span in enumerate(spans):
        if ES_VERSION[0] >= 5:
            assert span["name"] in (
                "ES PUT /tweets/%s/%d/_create" % (document_type, i + 1),
                "ES PUT /tweets/_create/%d" % (i + 1),
            )
        else:
            assert span["name"] == "ES PUT /tweets/%s/%d" % (document_type, i + 1)
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def test_get(instrument, elasticapm_client, elasticsearch):
    elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    # this is a fun one. Order pre-6x was (index, id, doc_type), changed to (index, doc_type, id) in 6.x, and reverted
    # to (index, id, doc_type) in 7.x. OK then.
    if ES_VERSION[0] == 6:
        r1 = elasticsearch.get("tweets", document_type, 1)
    else:
        r1 = elasticsearch.get(index="tweets", id=1, doc_type=document_type)
    r2 = elasticsearch.get(index="tweets", doc_type=document_type, id=1)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]
    for r in (r1, r2):
        assert r["found"]
        assert r["_source"] == {"user": "kimchy", "text": "hola"}
    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for span in spans:
        assert span["name"] == "ES GET /tweets/%s/1" % document_type
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def test_get_source(instrument, elasticapm_client, elasticsearch):
    elasticsearch.create(
        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
    )
    elasticapm_client.begin_transaction("test")
    if ES_VERSION[0] < 7:
        r1 = elasticsearch.get_source("tweets", document_type, 1)
    else:
        r1 = elasticsearch.get_source(index="tweets", id=1, doc_type=document_type)
    r2 = elasticsearch.get_source(index="tweets", doc_type=document_type, id=1)
    elasticapm_client.end_transaction("test", "OK")

    transaction = elasticapm_client.events[TRANSACTION][0]

    for r in (r1, r2):
        assert r == {"user": "kimchy", "text": "hola"}

    spans = elasticapm_client.spans_for_transaction(transaction)
    assert len(spans) == 2

    for span in spans:
        assert span["name"] == "ES GET /tweets/%s/1/_source" % document_type
        assert span["type"] == "db"
        assert span["subtype"] == "elasticsearch"
        assert span["action"] == "query"
        assert span["context"]["db"]["type"] == "elasticsearch"
        assert "statement" not in span["context"]["db"] 

def _restframework_version():
    import rest_framework
    return tuple(map(int, rest_framework.VERSION.split("."))) 

def __init__(self, options, columns):
        super(ElasticsearchFDW, self).__init__(options, columns)

        self.index = options.pop("index", "")
        self.doc_type = options.pop("type", "")
        self.query_column = options.pop("query_column", None)
        self.score_column = options.pop("score_column", None)
        self.scroll_size = int(options.pop("scroll_size", "1000"))
        self.scroll_duration = options.pop("scroll_duration", "10m")
        self._rowid_column = options.pop("rowid_column", "id")
        username = options.pop("username", None)
        password = options.pop("password", None)

        if ELASTICSEARCH_VERSION[0] >= 7:
            self.path = "/{index}".format(index=self.index)
            self.arguments = {"index": self.index}
        else:
            self.path = "/{index}/{doc_type}".format(
                index=self.index, doc_type=self.doc_type
            )
            self.arguments = {"index": self.index, "doc_type": self.doc_type}

        if (username is None) != (password is None):
            raise ValueError("Must provide both username and password")
        if username is not None:
            auth = (username, password)
        else:
            auth = None

        host = options.pop("host", "localhost")
        port = int(options.pop("port", "9200"))
        timeout = int(options.pop("timeout", "10"))
        self.client = Elasticsearch(
            [{"host": host, "port": port}], http_auth=auth, timeout=timeout, **options
        )

        self.columns = columns
        self.json_columns = {
            column.column_name
            for column in columns.values()
            if column.base_type_name.upper() in {"JSON", "JSONB"}
        }