diff --git a/README.md b/README.md index 1f9f734..6e7eafd 100644 --- a/README.md +++ b/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 diff --git a/src/opensearch_sql_cli/main.py b/src/opensearch_sql_cli/main.py index ec2092d..84f5a15 100644 --- a/src/opensearch_sql_cli/main.py +++ b/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) diff --git a/src/opensearch_sql_cli/opensearch_connection.py b/src/opensearch_sql_cli/opensearch_connection.py index 4e20686..6e8700f 100644 --- a/src/opensearch_sql_cli/opensearch_connection.py +++ b/src/opensearch_sql_cli/opensearch_connection.py @@ -27,7 +27,8 @@ 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. @@ -35,6 +36,9 @@ def __init__( :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: diff --git a/tests/test_main.py b/tests/test_main.py index ea4fe8f..bf8f08a 100644 --- a/tests/test_main.py +++ b/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": [], }