diff --git a/doctest/test_docs.py b/doctest/test_docs.py index dce37112492..6d724ec7b40 100644 --- a/doctest/test_docs.py +++ b/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" diff --git a/sql-cli/setup.py b/sql-cli/setup.py index 2da8f2b2971..c393dab28ff 100644 --- a/sql-cli/setup.py +++ b/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", ) diff --git a/sql-cli/src/opensearch_sql_cli/opensearch_connection.py b/sql-cli/src/opensearch_sql_cli/opensearch_connection.py index 505e8aadcef..1af79e593b1 100644 --- a/sql-cli/src/opensearch_sql_cli/opensearch_connection.py +++ b/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): diff --git a/sql-cli/src/opensearch_sql_cli/opensearchsql_cli.py b/sql-cli/src/opensearch_sql_cli/opensearchsql_cli.py index 04ca4608bbf..e490938cd8d 100644 --- a/sql-cli/src/opensearch_sql_cli/opensearchsql_cli.py +++ b/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): diff --git a/sql-cli/tests/test_formatter.py b/sql-cli/tests/test_formatter.py index 91da8651e78..7c322b99c84 100644 --- a/sql-cli/tests/test_formatter.py +++ b/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") diff --git a/sql-cli/tests/test_opensearch_connection.py b/sql-cli/tests/test_opensearch_connection.py index ba871e7700d..1a1aa76eccf 100644 --- a/sql-cli/tests/test_opensearch_connection.py +++ b/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): diff --git a/sql-cli/tests/test_opensearchsql_cli.py b/sql-cli/tests/test_opensearchsql_cli.py index a3d2723d7d7..8b76a2e9316 100644 --- a/sql-cli/tests/test_opensearchsql_cli.py +++ b/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) diff --git a/sql-cli/tests/utils.py b/sql-cli/tests/utils.py index 536f00fc89b..e9ecee64001 100644 --- a/sql-cli/tests/utils.py +++ b/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