From ec523acc63ea26c94b462ccffa8a76e431fb6cf6 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 11 Sep 2025 23:43:00 +0000 Subject: [PATCH 1/4] Add pretty calcite output Signed-off-by: Simeon Widdis --- src/opensearch_sql_cli/calcite.py | 71 +++++++++++++++++++++ src/opensearch_sql_cli/opensearchsql_cli.py | 11 +++- 2 files changed, 81 insertions(+), 1 deletion(-) create mode 100644 src/opensearch_sql_cli/calcite.py diff --git a/src/opensearch_sql_cli/calcite.py b/src/opensearch_sql_cli/calcite.py new file mode 100644 index 0000000..eb61e2d --- /dev/null +++ b/src/opensearch_sql_cli/calcite.py @@ -0,0 +1,71 @@ +import json +from pygments import highlight +from pygments.lexer import RegexLexer +from pygments.token import * +from pygments.formatters import TerminalFormatter + +class CalcitePlanLexer(RegexLexer): + name = 'CalcitePlan' + tokens = { + 'root': [ + (r'^=.*', Comment), + (r'[A-Z][a-z]\w+', Name.Function), + # (r'\[[^\]]*\]', Number), + (r'\$[\d\w]+', Name.Variable), + (r'"[^"]*"', String), + (r'\d+', Number), + (r'[a-zA-Z_][\w#\.]*=', Name.Attribute), + (r'\s+', Text), + (r'null', Keyword), + (r'[\w\-]+', Text), + (r'[^\w]', Punctuation), + ] + } + +def process_physical_calcite_line(line): + if "sourceBuilder=" not in line: + return line + + start_idx = line.index("sourceBuilder=") + len("sourceBuilder=") + + stack = 0 + for i in range(start_idx, len(line)): + if line[i] == '{': + stack += 1 + elif line[i] == '}': + stack -= 1 + if stack == 0: + end_idx = i + 1 + break + + json_str = line[start_idx:end_idx] + leading_space = len(line) - len(line.lstrip()) + indent_str = ' ' * leading_space + + try: + parsed_json = json.loads(json_str) + pretty_json = json.dumps(parsed_json, indent=2) + pretty_json = '\n'.join(indent_str + line for line in pretty_json.splitlines()).lstrip('| ') + return line[:start_idx] + pretty_json + line[end_idx:] + except json.JSONDecodeError: + return line + + +def format_calcite_explain(err): + report = """= Calcite Plan = +== Logical == +$LOGICAL + +== Physical == +$PHYSICAL""" + logical_plan = err['calcite']['logical'].splitlines() + logical_plan = '\n'.join(logical_plan) + + physical_plan = err['calcite']['physical'].splitlines() + physical_plan = map(process_physical_calcite_line, physical_plan) + physical_plan = '\n'.join(physical_plan) + + report = report.replace("$LOGICAL", logical_plan) + report = report.replace("$PHYSICAL", physical_plan) + highlighted = highlight(report, CalcitePlanLexer(), TerminalFormatter()) + return highlighted diff --git a/src/opensearch_sql_cli/opensearchsql_cli.py b/src/opensearch_sql_cli/opensearchsql_cli.py index 8b14144..0ac08e1 100644 --- a/src/opensearch_sql_cli/opensearchsql_cli.py +++ b/src/opensearch_sql_cli/opensearchsql_cli.py @@ -2,6 +2,7 @@ from os.path import expanduser, expandvars +from opensearch_sql_cli.calcite import format_calcite_explain from prompt_toolkit.history import FileHistory """ @@ -10,11 +11,14 @@ """ +import ast import click import re import pyfiglet import os import json +import pprint +from opensearchpy.exceptions import OpenSearchException from prompt_toolkit.completion import WordCompleter from prompt_toolkit.enums import DEFAULT_BUFFER @@ -141,7 +145,12 @@ def run_cli(self): formatter = Formatter(settings) formatted_output = formatter.format_output(output) self.echo_via_pager("\n".join(formatted_output)) - + except OpenSearchException as e: + msg = ast.literal_eval(str(e)) + if 'explain' in text.lower() and msg.get('calcite', False): + print(format_calcite_explain(msg)) + else: + pprint.pprint(msg) except Exception as e: print(repr(e)) From aa0c531c75c24a536e39ff4273c75ed3271c07b6 Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 11 Sep 2025 23:47:07 +0000 Subject: [PATCH 2/4] Tweaks Signed-off-by: Simeon Widdis --- src/opensearch_sql_cli/calcite.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/opensearch_sql_cli/calcite.py b/src/opensearch_sql_cli/calcite.py index eb61e2d..60d656d 100644 --- a/src/opensearch_sql_cli/calcite.py +++ b/src/opensearch_sql_cli/calcite.py @@ -14,9 +14,9 @@ class CalcitePlanLexer(RegexLexer): (r'\$[\d\w]+', Name.Variable), (r'"[^"]*"', String), (r'\d+', Number), - (r'[a-zA-Z_][\w#\.]*=', Name.Attribute), + (r'[a-zA-Z_][\w\.#\d]*(?:=)', Name.Attribute), (r'\s+', Text), - (r'null', Keyword), + (r'null|true|false', Keyword), (r'[\w\-]+', Text), (r'[^\w]', Punctuation), ] From 0fe2a45f6918e35110742aa1f0f2b48d5b50dfbc Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 11 Sep 2025 23:53:07 +0000 Subject: [PATCH 3/4] Fix upload-artifact Signed-off-by: Simeon Widdis --- .github/workflows/sql-cli-test-and-build-workflow.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/sql-cli-test-and-build-workflow.yml b/.github/workflows/sql-cli-test-and-build-workflow.yml index e095d96..2ed4dbb 100644 --- a/.github/workflows/sql-cli-test-and-build-workflow.yml +++ b/.github/workflows/sql-cli-test-and-build-workflow.yml @@ -60,7 +60,7 @@ jobs: cp -r ./dist/*.tar.gz ./dist/*.whl opensearchsql-builds/ - name: Upload Artifact - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v4 with: name: opensearchsql path: opensearchsql-builds From d0f195d30865301475db661f39ff2302d44917af Mon Sep 17 00:00:00 2001 From: Simeon Widdis Date: Thu, 11 Sep 2025 23:58:34 +0000 Subject: [PATCH 4/4] Update test explain output Signed-off-by: Simeon Widdis --- tests/test_main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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": [], }