Creation and updating 2.x branch

.github/workflows/sql-cli-test-and-build-workflow.yml

@@ -1,6 +1,7 @@
 name: SQL CLI Test and Build
 
 on:
+  workflow_dispatch:
   pull_request:
   push:
     branches-ignore:

src/opensearch_sql_cli/__init__.py

@@ -3,4 +3,4 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
-__version__ = "1.0.0"
+__version__ = "1.1.0"

src/opensearch_sql_cli/conf/clirc

@@ -31,13 +31,13 @@ multi_line = True
 multi_line_mode = opensearchsql_cli
 
 # log_file location.
-# In Unix/Linux: ~/.conf/opensearchsql-cli/log
+# In Unix/Linux: ~/.config/opensearchsql-cli/log
 # In Windows: %USERPROFILE%\AppData\Local\dbcli\opensearchsql-cli\log
 # %USERPROFILE% is typically C:\Users\{username}
 log_file = default
 
 # history_file location.
-# In Unix/Linux: ~/.conf/opensearchsql-cli/history
+# In Unix/Linux: ~/.config/opensearchsql-cli/history
 # In Windows: %USERPROFILE%\AppData\Local\dbcli\opensearchsql-cli\history
 # %USERPROFILE% is typically C:\Users\{username}
 history_file = default

src/opensearch_sql_cli/main.py

@@ -61,7 +61,7 @@
     "use_aws_authentication",
     is_flag=True,
     default=False,
-    help="Use AWS sigV4 to connect to AWS ELasticsearch domain",
+    help="Use AWS sigV4 to connect to AWS OpenSearch domain",
 )
 @click.option(
     "-l",
@@ -71,6 +71,14 @@
     default="sql",
     help="SQL OR PPL",
 )
+@click.option(
+    "-t",
+    "--timeout",
+    "response_timeout",
+    type=click.INT,
+    default=10,
+    help="Timeout in seconds to await a response from the server"
+)
 def cli(
     endpoint,
     query,
@@ -83,6 +91,7 @@ def cli(
     always_use_pager,
     use_aws_authentication,
     query_language,
+    response_timeout
 ):
     """
     Provide endpoint for OpenSearch client.
@@ -114,12 +123,9 @@ def cli(
         sys.exit(0)
 
     # use console to interact with user
-    opensearchsql_cli = OpenSearchSqlCli(
-        clirc_file=clirc,
-        always_use_pager=always_use_pager,
-        use_aws_authentication=use_aws_authentication,
-        query_language=query_language,
-    )
+    opensearchsql_cli = OpenSearchSqlCli(clirc_file=clirc, always_use_pager=always_use_pager,
+                                         use_aws_authentication=use_aws_authentication, query_language=query_language,
+                                         response_timeout=response_timeout)
     opensearchsql_cli.connect(endpoint, http_auth)
     opensearchsql_cli.run_cli()
 

src/opensearch_sql_cli/opensearch_connection.py

@@ -27,6 +27,7 @@ def __init__(
         http_auth=None,
         use_aws_authentication=False,
         query_language="sql",
+        response_timeout=10
     ):
         """Initialize an OpenSearchConnection instance.
 
@@ -45,6 +46,7 @@ def __init__(
         self.http_auth = http_auth
         self.use_aws_authentication = use_aws_authentication
         self.query_language = query_language
+        self.response_timeout = response_timeout
 
     def get_indices(self):
         if self.client:
@@ -58,7 +60,7 @@ def get_aes_client(self):
         region = session.region_name
 
         if credentials is not None:
-            self.aws_auth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service)
+            self.aws_auth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token)
         else:
             click.secho(
                 message="Can not retrieve your AWS credentials, check your AWS config",
@@ -167,14 +169,14 @@ def execute_query(self, query, output_format="jdbc", explain=False, use_console=
                 data = self.client.transport.perform_request(
                     url="/_plugins/_sql/_explain" if explain else "/_plugins/_sql/",
                     method="POST",
-                    params=None if explain else {"format": output_format},
+                    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/",
                     method="POST",
-                    params=None if explain else {"format": output_format},
+                    params=None if explain else {"format": output_format, "request_timeout": self.response_timeout},
                     body={"query": final_query},
                 )
             return data

src/opensearch_sql_cli/opensearchsql_cli.py

@@ -1,5 +1,9 @@
 from __future__ import unicode_literals
 
+from os.path import expanduser, expandvars
+
+from prompt_toolkit.history import FileHistory
+
 """
 Copyright OpenSearch Contributors
 SPDX-License-Identifier: Apache-2.0
@@ -21,7 +25,7 @@
 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from pygments.lexers.sql import SqlLexer
 
-from .config import get_config
+from .config import get_config, config_location
 from .opensearch_connection import OpenSearchConnection
 from .opensearch_buffer import opensearch_is_multiline
 from .opensearch_style import style_factory, style_factory_output
@@ -39,7 +43,8 @@
 class OpenSearchSqlCli:
     """OpenSearchSqlCli instance is used to build and run the OpenSearch SQL CLI."""
 
-    def __init__(self, clirc_file=None, always_use_pager=False, use_aws_authentication=False, query_language="sql"):
+    def __init__(self, clirc_file=None, always_use_pager=False, use_aws_authentication=False, query_language="sql",
+                 response_timeout=10):
         # Load conf file
         config = self.config = get_config(clirc_file)
         literal = self.literal = self._get_literals()
@@ -49,6 +54,7 @@ def __init__(self, clirc_file=None, always_use_pager=False, use_aws_authenticati
         self.query_language = query_language
         self.always_use_pager = always_use_pager
         self.use_aws_authentication = use_aws_authentication
+        self.response_timeout = response_timeout
         self.keywords_list = literal["keywords"]
         self.functions_list = literal["functions"]
         self.syntax_style = config["main"]["syntax_style"]
@@ -57,9 +63,15 @@ def __init__(self, clirc_file=None, always_use_pager=False, use_aws_authenticati
         self.multiline_continuation_char = config["main"]["multiline_continuation_char"]
         self.multi_line = config["main"].as_bool("multi_line")
         self.multiline_mode = config["main"].get("multi_line_mode", "src")
+        self.history_file = config["main"]["history_file"]
         self.null_string = config["main"].get("null_string", "null")
         self.style_output = style_factory_output(self.syntax_style, self.cli_style)
 
+        if self.history_file == "default":
+            self.history_file = os.path.join(config_location(), "history")
+        else:
+            self.history_file = expandvars(expanduser(self.history_file))
+
     def build_cli(self):
         # TODO: Optimize index suggestion to serve indices options only at the needed position, such as 'from'
         indices_list = self.opensearch_executor.indices_list
@@ -74,8 +86,7 @@ def get_continuation(width, *_):
             lexer=PygmentsLexer(SqlLexer),
             completer=sql_completer,
             complete_while_typing=True,
-            # TODO: add history, refer to pgcli approach
-            # history=history,
+            history=FileHistory(self.history_file),
             style=style_factory(self.syntax_style, self.cli_style),
             prompt_continuation=get_continuation,
             multiline=opensearch_is_multiline(self),
@@ -160,7 +171,7 @@ def echo_via_pager(self, text, color=None):
 
     def connect(self, endpoint, http_auth=None):
         self.opensearch_executor = OpenSearchConnection(
-            endpoint, http_auth, self.use_aws_authentication, self.query_language
+            endpoint, http_auth, self.use_aws_authentication, self.query_language, self.response_timeout
         )
         self.opensearch_executor.set_connection()
 

tests/test_opensearchsql_cli.py

@@ -18,6 +18,7 @@
 QUERY_WITH_CTRL_D = "select * from %s;\r\x04\r" % TEST_INDEX_NAME
 USE_AWS_CREDENTIALS = False
 QUERY_LANGUAGE = "sql"
+RESPONSE_TIMEOUT = 10
 
 
 @pytest.fixture()
@@ -34,7 +35,8 @@ def test_connect(self, cli):
         ) as mock_set_connectiuon:
             cli.connect(endpoint=ENDPOINT)
 
-            mock_OpenSearchConnection.assert_called_with(ENDPOINT, AUTH, USE_AWS_CREDENTIALS, QUERY_LANGUAGE)
+            mock_OpenSearchConnection.assert_called_with(ENDPOINT, AUTH, USE_AWS_CREDENTIALS, QUERY_LANGUAGE,
+                                                         RESPONSE_TIMEOUT)
             mock_set_connectiuon.assert_called()
 
     @estest