use opensearch-py client

doctest/test_docs.py

@@ -16,7 +16,7 @@
 from opensearch_sql_cli.opensearch_connection import OpenSearchConnection
 from opensearch_sql_cli.utils import OutputSettings
 from opensearch_sql_cli.formatter import Formatter
-from elasticsearch import Elasticsearch as OpenSearch, helpers
+from opensearchpy import OpenSearch, helpers
 
 ENDPOINT = "http://localhost:9200"
 ACCOUNTS = "accounts"

sql-cli/setup.py

@@ -13,7 +13,7 @@
     "prompt_toolkit == 2.0.6",
     "Pygments >= 2.7.4",
     "cli_helpers[styles] == 1.2.1",
-    "elasticsearch == 7.10.1",
+    "opensearch-py == 1.0.0",
     "pyfiglet == 0.8.post1",
     "boto3 == 1.16.29",
     "requests-aws4auth == 0.9",
@@ -22,9 +22,7 @@
 _version_re = re.compile(r"__version__\s+=\s+(.*)")
 
 with open("src/opensearch_sql_cli/__init__.py", "rb") as f:
-    version = str(
-        ast.literal_eval(_version_re.search(f.read().decode("utf-8")).group(1))
-    )
+    version = str(ast.literal_eval(_version_re.search(f.read().decode("utf-8")).group(1)))
 
 description = "OpenSearch SQL CLI with auto-completion and syntax highlighting"
 
@@ -38,8 +36,8 @@
     version=version,
     license="Apache 2.0",
     url="https://docs-beta.opensearch.org/search-plugins/sql/cli/",
-    packages=find_packages('src'),
-    package_dir={'': 'src'},
+    packages=find_packages("src"),
+    package_dir={"": "src"},
     package_data={"opensearch_sql_cli": ["conf/clirc", "opensearch_literals/opensearch_literals.json"]},
     description=description,
     long_description=long_description,
@@ -64,5 +62,5 @@
         "Topic :: Software Development",
         "Topic :: Software Development :: Libraries :: Python Modules",
     ],
-    python_requires='>=3.0'
+    python_requires=">=3.0",
 )

sql-cli/src/opensearch_sql_cli/opensearch_connection.py

@@ -10,9 +10,9 @@
 import sys
 import urllib3
 
-from elasticsearch import Elasticsearch as OpenSearch, RequestsHttpConnection
-from elasticsearch.exceptions import ConnectionError, RequestError
-from elasticsearch.connection import create_ssl_context
+from opensearchpy import OpenSearch, RequestsHttpConnection
+from opensearchpy.exceptions import ConnectionError, RequestError
+from opensearchpy.connection import create_ssl_context
 from requests_aws4auth import AWS4Auth
 
 
@@ -21,7 +21,13 @@ class OpenSearchConnection:
     as well as send user's SQL query to OpenSearch.
     """
 
-    def __init__(self, endpoint=None, http_auth=None, use_aws_authentication=False, query_language="sql"):
+    def __init__(
+        self,
+        endpoint=None,
+        http_auth=None,
+        use_aws_authentication=False,
+        query_language="sql",
+    ):
         """Initialize an OpenSearchConnection instance.
 
         Set up client and get indices list.
@@ -54,7 +60,10 @@ def get_aes_client(self):
         if credentials is not None:
             self.aws_auth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service)
         else:
-            click.secho(message="Can not retrieve your AWS credentials, check your AWS config", fg="red")
+            click.secho(
+                message="Can not retrieve your AWS credentials, check your AWS config",
+                fg="red",
+            )
 
         aes_client = OpenSearch(
             hosts=[self.endpoint],
@@ -131,7 +140,10 @@ def handle_server_close_connection(self):
             self.set_connection(is_reconnect=True)
             click.secho(message="Reconnected! Please run query again", fg="green")
         except ConnectionError as reconnection_err:
-            click.secho(message="Connection Failed. Check your OpenSearch is running and then come back", fg="red")
+            click.secho(
+                message="Connection Failed. Check your OpenSearch is running and then come back",
+                fg="red",
+            )
             click.secho(repr(reconnection_err), err=True, fg="red")
 
     def execute_query(self, query, output_format="jdbc", explain=False, use_console=True):

sql-cli/src/opensearch_sql_cli/opensearchsql_cli.py

@@ -159,7 +159,9 @@ def echo_via_pager(self, text, color=None):
             click.echo(text, color=color)
 
     def connect(self, endpoint, http_auth=None):
-        self.opensearch_executor = OpenSearchConnection(endpoint, http_auth, self.use_aws_authentication, self.query_language)
+        self.opensearch_executor = OpenSearchConnection(
+            endpoint, http_auth, self.use_aws_authentication, self.query_language
+        )
         self.opensearch_executor.set_connection()
 
     def _get_literals(self):

sql-cli/tests/test_formatter.py

@@ -135,7 +135,9 @@ def test_format_output_vertical(self):
             "age  | 24",
         ]
 
-        with mock.patch("src.opensearch_sql_cli.main.click.secho") as mock_secho, mock.patch("src.opensearch_sql_cli.main.click.confirm") as mock_confirm:
+        with mock.patch("src.opensearch_sql_cli.main.click.secho") as mock_secho, mock.patch(
+            "src.opensearch_sql_cli.main.click.confirm"
+        ) as mock_confirm:
             expanded_results = formatter.format_output(data)
 
         mock_secho.assert_called_with(message="Output longer than terminal width", fg="red")

sql-cli/tests/test_opensearch_connection.py

@@ -7,8 +7,8 @@
 import mock
 from textwrap import dedent
 
-from elasticsearch.exceptions import ConnectionError
-from elasticsearch import Elasticsearch as OpenSearch, RequestsHttpConnection
+from opensearchpy.exceptions import ConnectionError
+from opensearchpy import OpenSearch, RequestsHttpConnection
 
 from .utils import estest, load_data, run, TEST_INDEX_NAME
 from src.opensearch_sql_cli.opensearch_connection import OpenSearchConnection
@@ -59,7 +59,9 @@ def test_connection_fail(self):
         test_executor = OpenSearchConnection(endpoint=INVALID_ENDPOINT)
         err_message = "Can not connect to endpoint %s" % INVALID_ENDPOINT
 
-        with mock.patch("sys.exit") as mock_sys_exit, mock.patch("src.opensearch_sql_cli.opensearch_connection.click.secho") as mock_secho:
+        with mock.patch("sys.exit") as mock_sys_exit, mock.patch(
+            "src.opensearch_sql_cli.opensearch_connection.click.secho"
+        ) as mock_secho:
             test_executor.set_connection()
 
         mock_sys_exit.assert_called()
@@ -121,8 +123,11 @@ def test_get_od_client(self):
             od_test_executor.get_opensearch_client()
 
             mock_es.assert_called_with(
-                [OPENSEARCH_ENDPOINT], http_auth=AUTH, verify_certs=False, ssl_context=od_test_executor.ssl_context,
-                connection_class=RequestsHttpConnection
+                [OPENSEARCH_ENDPOINT],
+                http_auth=AUTH,
+                verify_certs=False,
+                ssl_context=od_test_executor.ssl_context,
+                connection_class=RequestsHttpConnection,
             )
 
     def test_get_aes_client(self):

sql-cli/tests/test_opensearchsql_cli.py

@@ -27,7 +27,9 @@ def cli(default_config_location):
 
 class TestOpenSearchSqlCli:
     def test_connect(self, cli):
-        with mock.patch.object(OpenSearchConnection, "__init__", return_value=None) as mock_OpenSearchConnection, mock.patch.object(
+        with mock.patch.object(
+            OpenSearchConnection, "__init__", return_value=None
+        ) as mock_OpenSearchConnection, mock.patch.object(
             OpenSearchConnection, "set_connection"
         ) as mock_set_connectiuon:
             cli.connect(endpoint=ENDPOINT)

sql-cli/tests/utils.py

@@ -5,7 +5,7 @@
 
 import json
 import pytest
-from elasticsearch import ConnectionError, helpers, ConnectionPool
+from opensearchpy import ConnectionError, helpers, ConnectionPool
 
 from src.opensearch_sql_cli.opensearch_connection import OpenSearchConnection
 from src.opensearch_sql_cli.utils import OutputSettings