Add Amazon OpenSearch Serverless Support (#29)

* Fix failing tests expected formatting Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Add serverless support for OS connections Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Add serverless tests Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Remove unused import Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Capitalize Serverless Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Add vcrpy dep Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Revert adding vcrpy dep Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Add vcrpy again Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Revert "Fix failing tests expected formatting" This reverts commit d8b24df. Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Add vcr to tox config Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Fix usage of credentials fixture Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Correct limit in `TestMain.test_explain` Signed-off-by: Simeon Widdis <sawiddis@amazon.com> * Fix lychee Signed-off-by: Simeon Widdis <sawiddis@amazon.com> --------- Signed-off-by: Simeon Widdis <sawiddis@amazon.com>
opensearch-project · Sep 25, 2024 · 9a1c0f3 · 9a1c0f3
1 parent d4e7789
commit 9a1c0f3
Showing 11 changed files with 167 additions and 17 deletions.
diff --git a/.lycheeignore b/.lycheeignore
@@ -0,0 +1 @@
+http://localhost:9200/
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -58,11 +58,11 @@ If you've thought of a way that OpenSearch could be better, we want to hear abou
 
 As with other types of contributions, the first step is to [**open an issue on GitHub**](https://github.com/opensearch-project/OpenSearch/issues/new/choose).  Opening an issue before you make changes makes sure that someone else isn't already working on that particular problem.  It also lets us all work together to find the right approach before you spend a bunch of time on a PR.  So again, when in doubt, open an issue.
 
-Once you've opened an issue, check out our [Developer Guide](../DEVELOPER_GUIDE.rst) for instructions on how to get started.
+Once you've opened an issue, check out our [Developer Guide](./development_guide.md) for instructions on how to get started.
 
 ## Developer Certificate of Origin
 
-OpenSearch is an open source product released under the Apache 2.0 license (see either [the Apache site](https://www.apache.org/licenses/LICENSE-2.0) or the [LICENSE.txt file](../LICENSE.txt)).  The Apache 2.0 license allows you to freely use, modify, distribute, and sell your own products that include Apache 2.0 licensed software.
+OpenSearch is an open source product released under the Apache 2.0 license (see either [the Apache site](https://www.apache.org/licenses/LICENSE-2.0) or the [LICENSE.txt file](./LICENSE)).  The Apache 2.0 license allows you to freely use, modify, distribute, and sell your own products that include Apache 2.0 licensed software.
 
 We respect intellectual property rights of others and we want to make sure all incoming contributions are correctly attributed and licensed. A Developer Certificate of Origin (DCO) is a lightweight mechanism to do that.
 

diff --git a/README.md b/README.md
@@ -93,11 +93,11 @@ You can also configure the following connection properties:
 
 * `endpoint`: You do not need to specify an option, anything that follows the launch command `opensearchsql` is considered as the endpoint. If you do not provide an endpoint, by default, the SQL CLI connects to [http://localhost:9200](http://localhost:9200/).
 * `-u/-w`: Supports username and password for HTTP basic authentication, such as:
-    * OpenSearch with [OpenSearch Security Plugin](https://docs-beta.opensearch.org/security-plugin/index/) installed
+    * OpenSearch with [OpenSearch Security Plugin](https://opensearch.org/docs/latest/security/) installed
     * Amazon OpenSearch Service domain with [Fine Grained Access Control](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/fgac.html) enabled
 * `--aws-auth`: Turns on AWS sigV4 authentication to connect to an Amazon Elasticsearch Service endpoint. Use with the AWS CLI (`aws configure`) to retrieve the local AWS configuration to authenticate and connect.
 
-For a list of all available configurations, see [clirc](https://github.com/opensearch-project/sql/blob/main/sql-cli/src/opensearch_sql_cli/conf/clirc).
+For a list of all available configurations, see [clirc](src/opensearch_sql_cli/conf/clirc).
 
 
 

diff --git a/development_guide.md b/development_guide.md
@@ -16,7 +16,7 @@ https uses 443 by default.
 - Prerequisites
     - Build the application
     - Start a local OpenSearch instance with 
-    [OpenSearch SQL plugin](https://docs-beta.opensearch.org/search-plugins/sql/index/) installed
+    [OpenSearch SQL plugin](https://opensearch.org/docs/latest/search-plugins/sql/sql/index/) installed
     and listening at http://localhost:9200.
 - Pytest
     - `pip install -r requirements-dev.txt` Install test frameworks including Pytest and mock.

diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -1,5 +1,6 @@
-pytest==7.4.2
+pytest==8.3.3
 mock==3.0.5
 pexpect==3.3
 twine==1.13.0
-tox>=1.9.2
+tox>=1.9.2
+pytest-vcr>=1.0.2
diff --git a/src/opensearch_sql_cli/opensearch_connection.py b/src/opensearch_sql_cli/opensearch_connection.py
@@ -47,14 +47,15 @@ def __init__(
         self.use_aws_authentication = use_aws_authentication
         self.query_language = query_language
         self.response_timeout = response_timeout
+        self.is_aws_serverless = self.use_aws_authentication and ".aoss.amazonaws.com" in self.endpoint
 
     def get_indices(self):
         if self.client:
             res = self.client.indices.get_alias().keys()
             self.indices_list = list(res)
 
     def get_aes_client(self):
-        service = "es"
+        service = "es" if not self.is_aws_serverless else "aoss"
         session = boto3.Session()
         credentials = session.get_credentials()
         region = session.region_name
@@ -92,7 +93,11 @@ def get_opensearch_client(self):
 
         return opensearch_client
 
-    def is_sql_plugin_installed(self, opensearch_client):
+    def is_sql_plugin_installed(self, opensearch_client: OpenSearch) -> bool:
+        if self.is_aws_serverless:
+            # If using serverless there's no _cat/plugins endpoint, SQL is always installed. See:
+            # https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-genref.html#serverless-plugins
+            return True
         self.plugins = opensearch_client.cat.plugins(params={"s": "component", "v": "true"})
         sql_plugin_name_list = ["opensearch-sql"]
         return any(x in self.plugins for x in sql_plugin_name_list)
@@ -103,7 +108,6 @@ def set_connection(self, is_reconnect=False):
 
         if self.http_auth:
             opensearch_client = self.get_opensearch_client()
-
         elif self.use_aws_authentication:
             opensearch_client = self.get_aes_client()
         else:
@@ -120,9 +124,14 @@ def set_connection(self, is_reconnect=False):
                 click.echo(self.plugins)
                 sys.exit()
 
-            # info() may throw ConnectionError, if connection fails to establish
-            info = opensearch_client.info()
-            self.opensearch_version = info["version"]["number"]
+            if self.is_aws_serverless:
+                # Serverless is versionless
+                self.opensearch_version = "Serverless"
+            else:
+                # info() may throw ConnectionError, if connection fails to establish
+                info = opensearch_client.info()
+                self.opensearch_version = info["version"]["number"]
+
             self.client = opensearch_client
             self.get_indices()
 
@@ -167,14 +176,14 @@ def execute_query(self, query, output_format="jdbc", explain=False, use_console=
         try:
             if self.query_language == "sql":
                 data = self.client.transport.perform_request(
-                    url="/_plugins/_sql/_explain" if explain else "/_plugins/_sql/",
+                    url="/_plugins/_sql/_explain" if explain else "/_plugins/_sql",
                     method="POST",
                     params=None if explain else {"format": output_format, "request_timeout": self.response_timeout},
                     body={"query": final_query},
                 )
             else:
                 data = self.client.transport.perform_request(
-                    url="/_plugins/_ppl/_explain" if explain else "/_plugins/_ppl/",
+                    url="/_plugins/_ppl/_explain" if explain else "/_plugins/_ppl",
                     method="POST",
                     params=None if explain else {"format": output_format, "request_timeout": self.response_timeout},
                     body={"query": final_query},

diff --git a/tests/cassettes/serverless_show_tables.yaml b/tests/cassettes/serverless_show_tables.yaml
@@ -0,0 +1,102 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Authorization:
+      - AWS4-HMAC-SHA256 Credential=testing/20240924/us-east-1/aoss/aws4_request,
+        SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token,
+        Signature=6835b7bf753ce295f7bcfe9eceb26c6150e3ac95b3f8b829d8ab280b851b5c6e
+      content-type:
+      - application/json
+      user-agent:
+      - opensearch-py/1.0.0 (Python 3.12.5)
+      x-amz-content-sha256:
+      - e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
+      x-amz-date:
+      - 20240924T221825Z
+      x-amz-security-token:
+      - redacted
+    method: GET
+    uri: https://example_endpoint.beta-us-east-1.aoss.amazonaws.com:443/_alias
+  response:
+    body:
+      string: '{"target_index":{"aliases":{}},".opensearch_dashboards_1":{"aliases":{".opensearch_dashboards":{}}},"sample-index1":{"aliases":{}}}'
+    headers:
+      content-length:
+      - '131'
+      content-type:
+      - application/json; charset=UTF-8
+      date:
+      - Tue, 24 Sep 2024 22:18:26 GMT
+      server:
+      - aoss-amazon-m
+      x-envoy-upstream-service-time:
+      - '49'
+      x-request-id:
+      - d77aa3ab-e474-98fc-b33d-1841f27e8b29
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '{"query":"SHOW TABLES LIKE %"}'
+    headers:
+      Accept-Charset:
+      - UTF-8
+      Authorization:
+      - AWS4-HMAC-SHA256 Credential=testing/20240924/us-east-1/aoss/aws4_request,
+        SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date;x-amz-security-token,
+        Signature=272ec9c51f0be839579f97311adce6f2c757e77498f3552ac9c794e8e951cf68
+      Content-Length:
+      - '30'
+      Content-Type:
+      - application/json
+      user-agent:
+      - opensearch-py/1.0.0 (Python 3.12.5)
+      x-amz-content-sha256:
+      - e53aaf06a94a0ef15a30e8ccb84866025b8ec4fec42338a4683a735b0fae6e9d
+      x-amz-date:
+      - 20240924T221826Z
+      x-amz-security-token:
+      - redacted
+    method: POST
+    uri: https://example_endpoint.beta-us-east-1.aoss.amazonaws.com:443/_plugins/_sql
+  response:
+    body:
+      string: "{\n  \"schema\": [\n    {\n      \"name\": \"TABLE_CAT\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TABLE_SCHEM\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TABLE_NAME\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TABLE_TYPE\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"REMARKS\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TYPE_CAT\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TYPE_SCHEM\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"TYPE_NAME\",\n      \"type\":
+        \"keyword\"\n    },\n    {\n      \"name\": \"SELF_REFERENCING_COL_NAME\",\n
+        \     \"type\": \"keyword\"\n    },\n    {\n      \"name\": \"REF_GENERATION\",\n
+        \     \"type\": \"keyword\"\n    }\n  ],\n  \"datarows\": [\n    [\n      \"opensearch\",\n
+        \     null,\n      \".opensearch_dashboards_1\",\n      \"BASE TABLE\",\n
+        \     null,\n      null,\n      null,\n      null,\n      null,\n      null\n
+        \   ],\n    [\n      \"opensearch\",\n      null,\n      \"sample-index1\",\n
+        \     \"BASE TABLE\",\n      null,\n      null,\n      null,\n      null,\n
+        \     null,\n      null\n    ],\n    [\n      \"opensearch\",\n      null,\n
+        \     \"target_index\",\n      \"BASE TABLE\",\n      null,\n      null,\n
+        \     null,\n      null,\n      null,\n      null\n    ],\n    [\n      \"opensearch\",\n
+        \     null,\n      \".opensearch_dashboards\",\n      \"BASE TABLE\",\n      null,\n
+        \     null,\n      null,\n      null,\n      null,\n      null\n    ]\n  ],\n
+        \ \"total\": 4,\n  \"size\": 4,\n  \"status\": 200\n}"
+    headers:
+      content-length:
+      - '1402'
+      content-type:
+      - application/json; charset=UTF-8
+      date:
+      - Tue, 24 Sep 2024 22:18:26 GMT
+      server:
+      - aoss-amazon-s
+      x-envoy-upstream-service-time:
+      - '27'
+      x-request-id:
+      - 644a3560-53d4-49d0-a46c-959353a62724
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -32,7 +32,7 @@ def test_explain(self, connection):
                     {
                         "name": "OpenSearchIndexScan",
                         "description": {
-                            "request": 'OpenSearchQueryRequest(indexName=opensearchsql_cli_test, sourceBuilder={"from":0,"size":200,"timeout":"1m","_source":{"includes":["a"],"excludes":[]}}, searchDone=false)'
+                            "request": 'OpenSearchQueryRequest(indexName=opensearchsql_cli_test, sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["a"],"excludes":[]}}, searchDone=false)'
                         },
                         "children": [],
                     }

diff --git a/tests/test_opensearch_serverless.py b/tests/test_opensearch_serverless.py
@@ -0,0 +1,36 @@
+import os
+import pytest
+from opensearch_sql_cli.opensearch_connection import OpenSearchConnection
+import vcr
+
+
+sql_cli_vcr = vcr.VCR(
+    cassette_library_dir="tests/cassettes",
+    record_mode="once",
+    match_on=["method", "path", "query", "body"],
+)
+
+class TestServerless:
+    @pytest.fixture(scope="function")
+    def aws_serverless_credentials(self):
+        os.environ["TEST_ENDPOINT_URL"] = "https://example_endpoint.beta-us-east-1.aoss.amazonaws.com:443"
+        os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+        os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+        os.environ["AWS_SESSION_TOKEN"] = "testing"
+
+
+    def test_show_tables(self, aws_serverless_credentials):
+        with sql_cli_vcr.use_cassette("serverless_show_tables.yaml"):
+            aes_test_executor = OpenSearchConnection(endpoint=os.environ["TEST_ENDPOINT_URL"], use_aws_authentication=True)
+            aes_test_executor.set_connection()
+
+            response = aes_test_executor.client.transport.perform_request(
+                method="POST",
+                url="/_plugins/_sql",
+                body={"query": "SHOW TABLES LIKE %"},
+                headers={"Content-Type": "application/json", "Accept-Charset": "UTF-8"},
+            )
+
+            assert response["status"] == 200
+            assert response["total"] == 4
+            assert response["datarows"][0][2] == ".opensearch_dashboards_1"
diff --git a/tests/test_plan.md b/tests/test_plan.md
@@ -22,7 +22,7 @@
 
 * [ ] Test connection to a local OpenSearch instance
     * [ ] OpenSearch, no authentication
-    * [ ] OpenSearch, install [OpenSearch Security plugin](https://docs-beta.opensearch.org/security-plugin/index/) to enable authentication and SSL
+    * [ ] OpenSearch, install [OpenSearch Security plugin](https://opensearch.org/docs/latest/security/) to enable authentication and SSL
     * Run command like `opensearchsql <endpoint> -u <username> -w <password>` to connect to instance with authentication.
 * [ ] Test connection to [Amazon Elasticsearch domain](https://docs.aws.amazon.com/elasticsearch-service/latest/developerguide/es-gsg.html) with
 [Fine Grained Access Control](https://docs.aws.amazon.com/elasticsearch-service/latest/developerguide/fgac.html) enabled. 

diff --git a/tox.ini b/tox.ini
@@ -4,4 +4,5 @@ envlist = py38
 deps = pytest==4.6.3
     mock==3.0.5
     pexpect==3.3
+    pytest-vcr==1.0.2
 commands = pytest