Add serverless override

README.md

@@ -138,6 +138,8 @@ Run single query from command line with options
 * `-l`: Query language option. Available options are [sql, ppl]. By default it's using sql.
 * `-p`: always use pager to display output
 * `--clirc`: provide path of config file to load
+* `--serverless` / `--no-serverless`:  Force whether to connect to OpenSearch Serverless instead of
+  a managed cluster. Default: guess from URL
 
 ## Code of Conduct
 

src/opensearch_sql_cli/main.py

@@ -79,6 +79,12 @@
     default=10,
     help="Timeout in seconds to await a response from the server"
 )
+@click.option(
+    "--serverless/--no-serverless",
+    default=None,
+    is_flag=True,
+    help="Force whether to connect to OpenSearch Serverless instead of a managed cluster. Default: guess from URL",
+)
 def cli(
     endpoint,
     query,
@@ -91,7 +97,8 @@ def cli(
     always_use_pager,
     use_aws_authentication,
     query_language,
-    response_timeout
+    response_timeout,
+    serverless,
 ):
     """
     Provide endpoint for OpenSearch client.
@@ -107,7 +114,9 @@ def cli(
 
     # handle single query without more interaction with user
     if query:
-        opensearch_executor = OpenSearchConnection(endpoint, http_auth, use_aws_authentication)
+        opensearch_executor = OpenSearchConnection(
+            endpoint, http_auth, use_aws_authentication, serverless_override=serverless
+        )
         opensearch_executor.set_connection()
         if explain:
             output = opensearch_executor.execute_query(query, explain=True, use_console=False)

src/opensearch_sql_cli/opensearch_connection.py

@@ -27,14 +27,18 @@ def __init__(
         http_auth=None,
         use_aws_authentication=False,
         query_language="sql",
-        response_timeout=10
+        response_timeout=10,
+        serverless_override=None,
     ):
         """Initialize an OpenSearchConnection instance.
 
         Set up client and get indices list.
 
         :param endpoint: an url in the format of "http://localhost:9200"
         :param http_auth: a tuple in the format of (username, password)
+        :param serverless_override: by default we analyze the provided endpoint string to determine
+            whether the endpoint is serverless or not. This option allows overriding that logic and
+            forcing the connection to be serverless (true) or not (false).
         """
         self.client = None
         self.ssl_context = None
@@ -47,7 +51,10 @@ 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
+        if serverless_override is not None:
+            self.is_aws_serverless = serverless_override
+        else:
+            self.is_aws_serverless = self.use_aws_authentication and ".aoss.amazonaws.com" in self.endpoint
 
     def get_indices(self):
         if self.client:

tests/test_main.py

@@ -33,7 +33,7 @@ def test_explain(self, connection):
                     {
                         "name": "OpenSearchIndexScan",
                         "description": {
-                            "request": 'OpenSearchQueryRequest(indexName=opensearchsql_cli_test, sourceBuilder={"from":0,"size":150,"timeout":"1m","_source":{"includes":["a"],"excludes":[]}}, searchDone=false)'
+                            "request": 'OpenSearchQueryRequest(indexName=opensearchsql_cli_test, sourceBuilder={"from":0,"size":150,"timeout":"1m","_source":{"includes":["a"],"excludes":[]}}, needClean=true, searchDone=false, pitId=null, cursorKeepAlive=null, searchAfter=null, searchResponse=null)'
                         },
                         "children": [],
                     }