diff --git a/.ci/certs/ca.crt b/.ci/certs/ca.crt deleted file mode 100755 index 71f9bfc8..00000000 --- a/.ci/certs/ca.crt +++ /dev/null @@ -1,20 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIDSjCCAjKgAwIBAgIVAJQLm8V2LcaCTHUcoIfO+KL63nG3MA0GCSqGSIb3DQEB -CwUAMDQxMjAwBgNVBAMTKUVsYXN0aWMgQ2VydGlmaWNhdGUgVG9vbCBBdXRvZ2Vu -ZXJhdGVkIENBMB4XDTIwMDIyNjA1NTA1N1oXDTIzMDIyNTA1NTA1N1owNDEyMDAG -A1UEAxMpRWxhc3RpYyBDZXJ0aWZpY2F0ZSBUb29sIEF1dG9nZW5lcmF0ZWQgQ0Ew -ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDYyajkPvGtUOE5M1OowQfB -kWVrWjo1+LIxzgCeRHp0YztLtdVJ0sk2xoSrt2uZpxcPepdyOseLTjFJex1D2yCR -AEniIqcFif4G72nDih2LlbhpUe/+/MTryj8ZTkFTzI+eMmbQi5FFMaH+kwufmdt/ -5/w8YazO18SxxJUlzMqzfNUrhM8vvvVdxgboU7PWhk28wZHCMHQovomHmzclhRpF -N0FMktA98vHHeRjH19P7rNhifSd7hZzoH3H148HVAKoPgqnZ6vW2O2YfAWOP6ulq -cyszr57p8fS9B2wSdlWW7nVHU1JuKcYD67CxbBS23BeGFgCj4tiNrmxO8S5Yf85v -AgMBAAGjUzBRMB0GA1UdDgQWBBSWAlip9eoPmnG4p4OFZeOUBlAbNDAfBgNVHSME -GDAWgBSWAlip9eoPmnG4p4OFZeOUBlAbNDAPBgNVHRMBAf8EBTADAQH/MA0GCSqG -SIb3DQEBCwUAA4IBAQA19qqrMTWl7YyId+LR/QIHDrP4jfxmrEELrAL58q5Epc1k -XxZLzOBSXoBfBrPdv+3XklWqXrZjKWfdkux0Xmjnl4qul+srrZDLJVZG3I7IrITh -AmQUmL9MuPiMnAcxoGZp1xpijtW8Qmd2qnambbljWfkuVaa4hcVRfrAX6TciIQ21 -bS5aeLGrPqR14h30YzDp0RMmTujEa1o6ExN0+RSTkE9m89Q6WdM69az8JW7YkWqm -I+UCG3TcLd3TXmN1zNQkq4y2ObDK4Sxy/2p6yFPI1Fds5w/zLfBOvvPQY61vEqs8 -SCCcQIe7f6NDpIRIBlty1C9IaEHj7edyHjF6rtYb ------END CERTIFICATE----- diff --git a/.ci/certs/ca.pem b/.ci/certs/ca.pem deleted file mode 100644 index 71f9bfc8..00000000 --- a/.ci/certs/ca.pem +++ /dev/null @@ -1,20 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIDSjCCAjKgAwIBAgIVAJQLm8V2LcaCTHUcoIfO+KL63nG3MA0GCSqGSIb3DQEB -CwUAMDQxMjAwBgNVBAMTKUVsYXN0aWMgQ2VydGlmaWNhdGUgVG9vbCBBdXRvZ2Vu -ZXJhdGVkIENBMB4XDTIwMDIyNjA1NTA1N1oXDTIzMDIyNTA1NTA1N1owNDEyMDAG -A1UEAxMpRWxhc3RpYyBDZXJ0aWZpY2F0ZSBUb29sIEF1dG9nZW5lcmF0ZWQgQ0Ew -ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDYyajkPvGtUOE5M1OowQfB -kWVrWjo1+LIxzgCeRHp0YztLtdVJ0sk2xoSrt2uZpxcPepdyOseLTjFJex1D2yCR -AEniIqcFif4G72nDih2LlbhpUe/+/MTryj8ZTkFTzI+eMmbQi5FFMaH+kwufmdt/ -5/w8YazO18SxxJUlzMqzfNUrhM8vvvVdxgboU7PWhk28wZHCMHQovomHmzclhRpF -N0FMktA98vHHeRjH19P7rNhifSd7hZzoH3H148HVAKoPgqnZ6vW2O2YfAWOP6ulq -cyszr57p8fS9B2wSdlWW7nVHU1JuKcYD67CxbBS23BeGFgCj4tiNrmxO8S5Yf85v -AgMBAAGjUzBRMB0GA1UdDgQWBBSWAlip9eoPmnG4p4OFZeOUBlAbNDAfBgNVHSME -GDAWgBSWAlip9eoPmnG4p4OFZeOUBlAbNDAPBgNVHRMBAf8EBTADAQH/MA0GCSqG -SIb3DQEBCwUAA4IBAQA19qqrMTWl7YyId+LR/QIHDrP4jfxmrEELrAL58q5Epc1k -XxZLzOBSXoBfBrPdv+3XklWqXrZjKWfdkux0Xmjnl4qul+srrZDLJVZG3I7IrITh -AmQUmL9MuPiMnAcxoGZp1xpijtW8Qmd2qnambbljWfkuVaa4hcVRfrAX6TciIQ21 -bS5aeLGrPqR14h30YzDp0RMmTujEa1o6ExN0+RSTkE9m89Q6WdM69az8JW7YkWqm -I+UCG3TcLd3TXmN1zNQkq4y2ObDK4Sxy/2p6yFPI1Fds5w/zLfBOvvPQY61vEqs8 -SCCcQIe7f6NDpIRIBlty1C9IaEHj7edyHjF6rtYb ------END CERTIFICATE----- diff --git a/.ci/certs/testnode_san.crt b/.ci/certs/testnode_san.crt deleted file mode 100644 index 8abba55b..00000000 --- a/.ci/certs/testnode_san.crt +++ /dev/null @@ -1,20 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIDVjCCAj6gAwIBAgIULh42yRefYlRRl1hvt055LrUH0HwwDQYJKoZIhvcNAQEL -BQAwNDEyMDAGA1UEAxMpRWxhc3RpYyBDZXJ0aWZpY2F0ZSBUb29sIEF1dG9nZW5l -cmF0ZWQgQ0EwHhcNMjAwMjI4MDMzNzIwWhcNMjMwMjI3MDMzNzIwWjATMREwDwYD -VQQDEwhpbnN0YW5jZTCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAIUP -t267NN21z+3ukajej8eojSXwP6zHxy7CUAp+sQ7bTq2XCKxkYX3CW9ThcS4cV9mL -ayYdWEYnbEDGYPQDo7Wk3Ih5OEXTMZb/yNEx5D4S2lGMOS5bCDdYx6GvwCMG4jNx -aMktosaxpprAJiHh2oLgQk0hQc/a9JfMo6kJKtuhjxsxjxLwcOHhuaUD7NS0Pjop -CJkSYcrL+nnQPQjKe4uLhAbSyiX914h4QX0CJ0e4z1ccdDX2PFWTrwaIf//vQhCR -wP2YKdfjR0JB4oDAlu85GsIs2cFLPysM5ufuNZO4fCr8uOwloKI8zZ2HhlIfBEcY -Gcy4g9N/9epmxMXZlGcCAwEAAaOBgDB+MB0GA1UdDgQWBBRefYm8DHHDdkTPHhS1 -HEUwTb2uiDAfBgNVHSMEGDAWgBSWAlip9eoPmnG4p4OFZeOUBlAbNDAxBgNVHREE -KjAogglsb2NhbGhvc3SHBH8AAAGHEAAAAAAAAAAAAAAAAAAAAAGCA2VzMTAJBgNV -HRMEAjAAMA0GCSqGSIb3DQEBCwUAA4IBAQC+pauqM2wJjQaHyHu+kIm59P4b/5Oj -IH1cYCQfMB7Y2UMLxp0ew+f7o7zzE2DA52YYFDWy6J5DVWtSBPyeFGgX+RH+aA+9 -Iv4cc9QpAs6aFjncorHrzNOrWLgCHIeRAxTR0CAkeP2dUZfDBuMpRyP6rAsYzyLH -Rb3/BfYJSI5vxgt5Ke49Y/ljDKFJTyDmAVrHQ4JWrseYE1UZ2eDkBXeiRlYE/QtB -YsrUSqdL6zvFZyUcilxDUUabNcA+GgeGZ2lAEA90F8vwi62QwRXo3Iv1Hz+6xc43 -nFofDK9D8/qkrUD9iuhpx1974QwPhwWyjn9RZRpbZA4ngRL+szdRXR4N ------END CERTIFICATE----- diff --git a/.ci/certs/testnode_san.key b/.ci/certs/testnode_san.key deleted file mode 100644 index 75d19539..00000000 --- a/.ci/certs/testnode_san.key +++ /dev/null @@ -1,27 +0,0 @@ ------BEGIN RSA PRIVATE KEY----- -MIIEogIBAAKCAQEAhQ+3brs03bXP7e6RqN6Px6iNJfA/rMfHLsJQCn6xDttOrZcI -rGRhfcJb1OFxLhxX2YtrJh1YRidsQMZg9AOjtaTciHk4RdMxlv/I0THkPhLaUYw5 -LlsIN1jHoa/AIwbiM3FoyS2ixrGmmsAmIeHaguBCTSFBz9r0l8yjqQkq26GPGzGP -EvBw4eG5pQPs1LQ+OikImRJhysv6edA9CMp7i4uEBtLKJf3XiHhBfQInR7jPVxx0 -NfY8VZOvBoh//+9CEJHA/Zgp1+NHQkHigMCW7zkawizZwUs/Kwzm5+41k7h8Kvy4 -7CWgojzNnYeGUh8ERxgZzLiD03/16mbExdmUZwIDAQABAoIBAEwhjulLMVc9JEfV -PP/qv0cUOBYh3LzF3T/yq4slq7Z9YgnOJYdFM8aZgqNNjc09KEJvE5JOLeiNu9Ff -768Nugg+2HM5MCo7SN9FYCfZLOcbMFCCM2FDcnMAV9A512vzD08xryuT8dNPZ6yZ -DfhK2hQRrb2lrpr3gwSrcGRRu3THqvq7X1RIjpLV3teDMeP8rQPAlpj8fmP+kdVV -5y1ihiDIo87McihG9FMavJtBDXQkUEuVw6eIeir8L/zHHD/ZwhYjNHZGWbrB88sz -CkJkfWh/FlA63tCVdJzkmnERALLTVy9mR0Sq6sUlnFhFNO2BRdWgYLrcp9McfTJC -e8+WsSECgYEAuwQ3nAaFL0jqYu1AREyKT/f3WUenf2UsX7dwwV2/yFtQvkzW7ji4 -uZLnfUnZBojtHf35dRo+hDgtvhZhgZNAuPPsbOl/EIMTcbChEqV/3CSTFlhLFM1d -hfM9PoM+Bt/pyUNabjD1sWM0X7WeUhzcddshY3S4daBsNsLuOzweRRcCgYEAtiSS -4qiiGafYsY7gOHuAlOhs/00+1uWIFEHKgoHM9vzCxDN3LCmBdynHk8ZE2TAdhw+l -7xpu6LUxKQDfGmVZa9Epg0kQmVq9c54oQP57pJ3tR+68++insEkfnaZH8jblfq2s -sSkFrY3pdS19edq60nuft64kswKRUUkamCXTXTECgYBdoSfiMpV9bekC7DsPtq5M -iR3KEgi2zEViCmomNTRuL+GF1NyKWdWJ+xVwcYd5MRZdvKimyyPfeGzWTUg14i42 -KtEEWgZmkukqMz8BIeCYq6sENeIpIQQgqv3PjU+Bi5r1S4Y7wsFPNRakkD4aaB6r -1rCppWcwZMeoxwEUoO2aswKBgBdDIIdWJi3EpAY5SyWrkEZ0UMdiZC4p7nE33ddB -IJ5CtdU9BXFcc652ZYjX/58FaCABvZ2F8LhDu92SwOusGfmNIxIjWL1dO2jywA1c -8wmZKd7P/M7nbdMz45fMzs9+d1zwbWfK53C8+R4AC1BuwQF0zHc3BHTgVRLelUjt -O8thAoGAdO2gHIqEsZzTgbvLbsh52eVbumjfNGnrnEv1fjb+o+/wAol8dymcmzbL -bZCRzoyA0qwU9kdPFgX46H6so6o1tUM2GQtVFoT6kDnPv7EkLQK0C4cDh6OOHxDU -NPvr/9fHhQd9EDWDvS1JnVMAdKDO6ELp3SoKGGmCXR2QplnqWAk= ------END RSA PRIVATE KEY----- diff --git a/.ci/run-elasticsearch.sh b/.ci/run-elasticsearch.sh index 4ce7ccc0..64ba2489 100755 --- a/.ci/run-elasticsearch.sh +++ b/.ci/run-elasticsearch.sh @@ -27,10 +27,6 @@ CLUSTER_NAME=${CLUSTER_NAME-${moniker}${suffix}} HTTP_PORT=${HTTP_PORT-9200} ELASTIC_PASSWORD=${ELASTIC_PASSWORD-changeme} -SSL_CERT=${SSL_CERT-"${SCRIPT_PATH}/certs/testnode_san.crt"} -SSL_KEY=${SSL_KEY-"${SCRIPT_PATH}/certs/testnode_san.key"} -SSL_CA=${SSL_CA-"${SCRIPT_PATH}/certs/ca.crt"} -SSL_CA_PEM=${SSL_CA-"${SCRIPT_PATH}/certs/ca.pem"} DETACH=${DETACH-false} CLEANUP=${CLEANUP-false} @@ -122,35 +118,16 @@ if [[ "$ELASTICSEARCH_VERSION" != *oss* ]]; then environment+=($(cat <<-END --env ELASTIC_PASSWORD=$ELASTIC_PASSWORD --env xpack.license.self_generated.type=trial - --env xpack.security.enabled=true - --env xpack.security.http.ssl.enabled=true - --env xpack.security.http.ssl.verification_mode=certificate - --env xpack.security.http.ssl.key=certs/testnode_san.key - --env xpack.security.http.ssl.certificate=certs/testnode_san.crt - --env xpack.security.http.ssl.certificate_authorities=certs/ca.crt - --env xpack.security.transport.ssl.enabled=true - --env xpack.security.transport.ssl.key=certs/testnode_san.key - --env xpack.security.transport.ssl.certificate=certs/testnode_san.crt - --env xpack.security.transport.ssl.certificate_authorities=certs/ca.crt -END -)) - volumes+=($(cat <<-END - --volume $SSL_CERT:/usr/share/elasticsearch/config/certs/testnode_san.crt - --volume $SSL_KEY:/usr/share/elasticsearch/config/certs/testnode_san.key - --volume $SSL_CA:/usr/share/elasticsearch/config/certs/ca.crt - --volume $SSL_CA_PEM:/usr/share/elasticsearch/config/certs/ca.pem + --env xpack.security.enabled=false + --env xpack.security.http.ssl.enabled=false + --env xpack.security.transport.ssl.enabled=false END )) fi url="http://$NODE_NAME" if [[ "$ELASTICSEARCH_VERSION" != *oss* ]]; then - url="https://elastic:$ELASTIC_PASSWORD@$NODE_NAME" -fi - -cert_validation_flags="--insecure" -if [[ "$NODE_NAME" == "es1" ]]; then - cert_validation_flags="--cacert /usr/share/elasticsearch/config/certs/ca.pem --resolve ${NODE_NAME}:443:127.0.0.1" + url="http://elastic:$ELASTIC_PASSWORD@$NODE_NAME" fi echo -e "\033[34;1mINFO:\033[0m Starting container $NODE_NAME \033[0m" @@ -165,7 +142,7 @@ docker run \ --ulimit nofile=65536:65536 \ --ulimit memlock=-1:-1 \ --detach="$DETACH" \ - --health-cmd="curl $cert_validation_flags --fail $url:9200/_cluster/health || exit 1" \ + --health-cmd="curl --insecure --fail $url:9200/_cluster/health || exit 1" \ --health-interval=2s \ --health-retries=20 \ --health-timeout=2s \ diff --git a/.ci/run-tests b/.ci/run-tests index eb6dfcf0..0f73a094 100755 --- a/.ci/run-tests +++ b/.ci/run-tests @@ -11,11 +11,11 @@ set -euxo pipefail TEST_SUITE=${TEST_SUITE-xpack} -NODE_NAME=es1 +NODE_NAME=localhost elasticsearch_image=elasticsearch -elasticsearch_url=https://elastic:changeme@${NODE_NAME}:9200 +elasticsearch_url=http://elastic:changeme@${NODE_NAME}:9200 if [[ $TEST_SUITE != "xpack" ]]; then elasticsearch_image=elasticsearch-${TEST_SUITE} elasticsearch_url=http://${NODE_NAME}:9200 @@ -44,14 +44,14 @@ echo -e "\033[1m>>>>> Start [$ELASTICSEARCH_VERSION container] >>>>>>>>>>>>>>>>> ELASTICSEARCH_VERSION=${elasticsearch_image}:${ELASTICSEARCH_VERSION} \ NODE_NAME=${NODE_NAME} \ - NETWORK_NAME=elasticsearch \ + NETWORK_NAME=host \ DETACH=true \ bash .ci/run-elasticsearch.sh echo -e "\033[1m>>>>> Repository specific tests >>>>>>>>>>>>>>>>>>>>>>>>>>>>>\033[0m" ELASTICSEARCH_CONTAINER=${elasticsearch_image}:${ELASTICSEARCH_VERSION} \ - NETWORK_NAME=elasticsearch \ + NETWORK_NAME=host \ NODE_NAME=${NODE_NAME} \ ELASTICSEARCH_URL=${elasticsearch_url} \ TEST_SUITE=${TEST_SUITE} \ diff --git a/docs/source/reference/api/eland.DataFrame.es_query.rst b/docs/source/reference/api/eland.DataFrame.es_query.rst new file mode 100644 index 00000000..e0e78791 --- /dev/null +++ b/docs/source/reference/api/eland.DataFrame.es_query.rst @@ -0,0 +1,6 @@ +eland.DataFrame.es_query +======================== + +.. currentmodule:: eland + +.. automethod:: DataFrame.es_query diff --git a/docs/source/reference/dataframe.rst b/docs/source/reference/dataframe.rst index 4757c0c1..9da5a297 100644 --- a/docs/source/reference/dataframe.rst +++ b/docs/source/reference/dataframe.rst @@ -76,6 +76,14 @@ Plotting DataFrame.hist +Elasticsearch Functions +~~~~~~~~~~~~~~~~~~~~~~~ +.. autosummary:: + :toctree: api/ + + DataFrame.info_es + DataFrame.es_query + Serialization / IO / conversion ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autosummary:: @@ -86,10 +94,3 @@ Serialization / IO / conversion DataFrame.to_csv DataFrame.to_html DataFrame.to_string - -Elasticsearch utilities -~~~~~~~~~~~~~~~~~~~~~~~ -.. autosummary:: - :toctree: api/ - - DataFrame.info_es diff --git a/eland/dataframe.py b/eland/dataframe.py index 4a95cf0e..c4345675 100644 --- a/eland/dataframe.py +++ b/eland/dataframe.py @@ -569,6 +569,62 @@ def info_es(self): return buf.getvalue() + def es_query(self, query): + """Applies an Elasticsearch DSL query to the current DataFrame. + + Parameters + ---------- + query: + Dictionary of the Elasticsearch DSL query to apply + + Returns + ------- + eland.DataFrame: + eland DataFrame with the query applied + + Examples + -------- + + Apply a `geo-distance query`_ to a dataset with a geo-point column ``geoip.location``. + + .. _geo-distance query: https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-geo-distance-query.html + + >>> df = ed.DataFrame('localhost', 'ecommerce', columns=['customer_first_name', 'geoip.city_name']) + >>> df.es_query({"bool": {"filter": {"geo_distance": {"distance": "1km", "geoip.location": [55.3, 25.3]}}}}).head() + customer_first_name geoip.city_name + 1 Mary Dubai + 9 Rabbia Al Dubai + 10 Rabbia Al Dubai + 22 Mary Dubai + 30 Robbie Dubai + + [5 rows x 2 columns] + + If using an occurrence like ``must`` or ``filter`` you must + nest it within ``bool``: + + .. code-block:: python + + # Correct: + df.es_query({ + "bool": { + "filter": {...} + } + }) + + # Incorrect, needs to be nested under 'bool': + df.es_query({ + "filter": {...} + }) + """ + # Unpack the {'query': ...} which some + # users may use due to documentation. + if not isinstance(query, dict): + raise TypeError("'query' must be of type 'dict'") + if tuple(query) == ("query",): + query = query["query"] + return DataFrame(query_compiler=self._query_compiler.es_query(query)) + def _index_summary(self): # Print index summary e.g. # Index: 103 entries, 0 to 102 diff --git a/eland/filter.py b/eland/filter.py index 592fe9da..fb2bbf27 100644 --- a/eland/filter.py +++ b/eland/filter.py @@ -178,3 +178,9 @@ def __init__(self, inline, lang=None, params=None): if params is not None: script["params"] = params self._filter = {"script": {"script": script}} + + +class QueryFilter(BooleanFilter): + def __init__(self, query): + super().__init__() + self._filter = query diff --git a/eland/query.py b/eland/query.py index 3ab814c3..701e6ad8 100644 --- a/eland/query.py +++ b/eland/query.py @@ -135,27 +135,20 @@ def hist_aggs(self, name, field, min_aggs, max_aggs, num_bins): self._aggs[name] = agg def to_search_body(self): - if self._query.empty(): - if self._aggs: - body = {"aggs": self._aggs} - else: - body = {} - else: - if self._aggs: - body = {"query": self._query.build(), "aggs": self._aggs} - else: - body = {"query": self._query.build()} + body = {} + if self._aggs: + body["aggs"] = self._aggs + if not self._query.empty(): + body["query"] = self._query.build() return body def to_count_body(self): if len(self._aggs) > 0: warnings.warn("Requesting count for agg query {}", self) if self._query.empty(): - body = None + return None else: - body = {"query": self._query.build()} - - return body + return {"query": self._query.build()} def update_boolean_filter(self, boolean_filter): if self._query.empty(): diff --git a/eland/query_compiler.py b/eland/query_compiler.py index 7b0bd728..6a7aae08 100644 --- a/eland/query_compiler.py +++ b/eland/query_compiler.py @@ -25,6 +25,7 @@ from eland import FieldMappings from eland import Index from eland import Operations +from eland.filter import QueryFilter class QueryCompiler: @@ -397,6 +398,9 @@ def tail(self, n): return result + def es_query(self, query): + return self._update_query(QueryFilter(query)) + # To/From Pandas def to_pandas(self, show_progress=False): """Converts Eland DataFrame to Pandas DataFrame. diff --git a/eland/tests/__init__.py b/eland/tests/__init__.py index 9205338f..d98542e0 100644 --- a/eland/tests/__init__.py +++ b/eland/tests/__init__.py @@ -25,17 +25,8 @@ # Define client to use in tests TEST_SUITE = os.environ.get("TEST_SUITE", "xpack") if TEST_SUITE == "xpack": - print("Running xpack tests requires SSL. Setting up SSL enabled client") - certpath = os.path.abspath( - os.path.join(os.path.dirname(__file__), "../../.ci/certs/ca.crt") - ) - print(certpath) ES_TEST_CLIENT = Elasticsearch( - ELASTICSEARCH_HOST, - http_auth=("elastic", "changeme"), - use_ssl=True, - verify_certs=True, - ca_certs=certpath, + ELASTICSEARCH_HOST, http_auth=("elastic", "changeme"), ) else: ES_TEST_CLIENT = Elasticsearch(ELASTICSEARCH_HOST) diff --git a/eland/tests/dataframe/test_es_query_pytest.py b/eland/tests/dataframe/test_es_query_pytest.py new file mode 100644 index 00000000..9e9d3ce9 --- /dev/null +++ b/eland/tests/dataframe/test_es_query_pytest.py @@ -0,0 +1,68 @@ +# Copyright 2019 Elasticsearch BV +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# File called _pytest for PyCharm compatability + +import pytest +from eland.tests.common import TestData +from eland.tests.common import assert_eland_frame_equal + + +class TestDataEsQuery(TestData): + def test_flights_match_query(self): + ed_flights = self.ed_flights() + + left = ed_flights.es_query({"match": {"OriginCityName": "Rome"}})[ + ed_flights["Carrier"] == "Kibana Airlines" + ] + + right = ed_flights[ed_flights["Carrier"] == "Kibana Airlines"].es_query( + {"match": {"OriginCityName": "Rome"}} + ) + + assert len(left) > 0 + assert_eland_frame_equal(left, right) + + def test_es_query_allows_query_in_dict(self): + ed_flights = self.ed_flights() + + left = ed_flights.es_query({"match": {"OriginCityName": "Rome"}}) + right = ed_flights.es_query({"query": {"match": {"OriginCityName": "Rome"}}}) + + assert len(left) > 0 + assert_eland_frame_equal(left, right) + + def test_es_query_geo_location(self): + df = self.ed_ecommerce() + cur_nearby = df.es_query( + { + "bool": { + "filter": { + "geo_distance": { + "distance": "10km", + "geoip.location": {"lon": 55.3, "lat": 25.3}, + } + } + } + } + )["currency"].value_counts() + + assert cur_nearby["EUR"] == 476 + + @pytest.mark.parametrize("query", [(), [], 1, True]) + def test_es_query_wrong_type(self, query): + ed_flights = self.ed_flights_small() + + with pytest.raises(TypeError): + ed_flights.es_query(query) diff --git a/noxfile.py b/noxfile.py index aec7c992..41dc3af8 100644 --- a/noxfile.py +++ b/noxfile.py @@ -27,4 +27,4 @@ def lint(session): def test(session): session.install("-r", "requirements-dev.txt") session.run("python", "-m", "eland.tests.setup_tests") - session.run("pytest", *(session.posargs or ("eland/tests/",))) + session.run("pytest", "--doctest-modules", *(session.posargs or ("eland/",)))