[Backport 3.1] Fix PIT context leak in Legacy SQL for non-paginated queries

doctest/bootstrap.sh

@@ -2,13 +2,15 @@
 
 DIR=$(dirname "$0")
 
-if hash python3.8 2> /dev/null; then
-    PYTHON=python3.8
+# Try to find Python 3.12 or newer first, then fall back to any python3
+if hash python3.13 2> /dev/null; then
+    PYTHON=python3.13
+elif hash python3.12 2> /dev/null; then
+    PYTHON=python3.12
 elif hash python3 2> /dev/null; then
-    # fallback to python3 in case there is no python3.8 alias; should be 3.8
     PYTHON=python3
 else
-    echo 'python3.8 required'
+    echo 'python3 required'
     exit 1
 fi
 
@@ -21,4 +23,16 @@ fi
 
 $DIR/.venv/bin/pip install -U pip setuptools wheel
 $DIR/.venv/bin/pip install -r $DIR/requirements.txt
-$DIR/.venv/bin/pip install -e ./sql-cli
+
+# Only install sql-cli if Python version is 3.12+
+PYTHON_VERSION=$($DIR/.venv/bin/python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
+PYTHON_MAJOR=$(echo $PYTHON_VERSION | cut -d. -f1)
+PYTHON_MINOR=$(echo $PYTHON_VERSION | cut -d. -f2)
+
+if [ "$PYTHON_MAJOR" -gt 3 ] || ([ "$PYTHON_MAJOR" -eq 3 ] && [ "$PYTHON_MINOR" -ge 12 ]); then
+    echo "Installing sql-cli with Python $PYTHON_VERSION..."
+    $DIR/.venv/bin/pip install -e $DIR/sql-cli
+else
+    echo "Warning: Python $PYTHON_VERSION is too old for sql-cli (requires >=3.12). Skipping sql-cli installation."
+    echo "Doctest will continue without sql-cli support."
+fi

doctest/build.gradle

@@ -24,14 +24,15 @@ def plugin_path = project(':doctest').projectDir
 task cloneSqlCli(type: Exec) {
     def repoDir = new File("${project.projectDir}/sql-cli")
 
-    if (repoDir.exists()) {
-        // Repository already exists, fetch and checkout latest
-        commandLine 'git', '-C', repoDir.absolutePath, 'fetch', 'origin', 'main'
-        commandLine 'git', '-C', repoDir.absolutePath, 'checkout', 'origin/main'
-    } else {
-        // Repository doesn't exist, clone it
-        commandLine 'git', 'clone', 'https://github.com/opensearch-project/sql-cli.git', repoDir.absolutePath
-    }
+    commandLine 'sh', '-c', """
+        if [ -d "${repoDir.absolutePath}/.git" ]; then
+            echo "Updating existing sql-cli repository..."
+            cd "${repoDir.absolutePath}" && git fetch origin main && git checkout origin/main
+        else
+            echo "Cloning sql-cli repository..."
+            git clone https://github.com/opensearch-project/sql-cli.git "${repoDir.absolutePath}"
+        fi
+    """
 }
 
 task bootstrap(type: Exec, dependsOn: ['cloneSqlCli', 'spotlessJava']) {
@@ -76,8 +77,21 @@ task startOpenSearch(type: SpawnProcessTask) {
 }
 
 task doctest(type: Exec, dependsOn: ['bootstrap']) {
-
-    commandLine "$projectDir/bin/test-docs"
+    // Check if sql-cli was installed during bootstrap by looking for opensearch-sql-cli in venv
+    def venvLibDirs = file("$projectDir/.venv/lib").listFiles()?.findAll { it.isDirectory() && it.name.startsWith("python") }
+    def sqlCliInstalled = venvLibDirs?.any { pythonDir ->
+        def sitePackages = new File(pythonDir, "site-packages")
+        sitePackages.exists() && sitePackages.listFiles()?.any { 
+            it.name.toLowerCase().contains("opensearch") && it.name.toLowerCase().contains("sql") && it.name.toLowerCase().contains("cli")
+        }
+    } ?: false
+
+    if (sqlCliInstalled) {
+        commandLine "$projectDir/bin/test-docs"
+    } else {
+        // Skip doctest if sql-cli not available (Python < 3.12)
+        commandLine 'echo', 'Skipping doctest: opensearch-sql-cli not available (requires Python >=3.12)'
+    }
 
     doLast {
         // remove the cloned sql-cli folder

doctest/test_docs.py

@@ -13,11 +13,21 @@
 import click
 
 from functools import partial
-from opensearch_sql_cli.opensearch_connection import OpenSearchConnection
-from opensearch_sql_cli.utils import OutputSettings
-from opensearch_sql_cli.formatter import Formatter
 from opensearchpy import OpenSearch, helpers
 
+# Try to import sql-cli components, skip tests if not available
+try:
+    from opensearch_sql_cli.opensearch_connection import OpenSearchConnection
+    from opensearch_sql_cli.utils import OutputSettings
+    from opensearch_sql_cli.formatter import Formatter
+    SQL_CLI_AVAILABLE = True
+except ImportError:
+    SQL_CLI_AVAILABLE = False
+    # Create dummy classes to prevent NameError during module loading
+    OpenSearchConnection = object
+    OutputSettings = object
+    Formatter = object
+
 ENDPOINT = "http://localhost:9200"
 ACCOUNTS = "accounts"
 EMPLOYEES = "employees"
@@ -86,8 +96,14 @@ def pretty_print(s):
         print(s)
 
 
-sql_cmd = DocTestConnection(query_language="sql")
-ppl_cmd = DocTestConnection(query_language="ppl")
+# Only instantiate DocTestConnection if sql-cli is available
+if SQL_CLI_AVAILABLE:
+    sql_cmd = DocTestConnection(query_language="sql")
+    ppl_cmd = DocTestConnection(query_language="ppl")
+else:
+    sql_cmd = None
+    ppl_cmd = None
+
 test_data_client = OpenSearch([ENDPOINT], verify_certs=True)
 
 
@@ -204,6 +220,14 @@ def doc_suite(fn):
 
 
 def load_tests(loader, suite, ignore):
+    # Skip all tests if sql-cli is not available (requires Python >=3.12)
+    if not SQL_CLI_AVAILABLE:
+        class SkippedDocTests(unittest.TestCase):
+            @unittest.skip("opensearch-sql-cli not available (requires Python >=3.12)")
+            def test_skip_all_doctests(self):
+                pass
+        return unittest.TestSuite([SkippedDocTests('test_skip_all_doctests')])
+
     tests = []
     # Load doctest docs by category
     with open('../docs/category.json') as json_file:

integ-test/build.gradle

@@ -286,8 +286,15 @@ testClusters {
 
 task startPrometheus(type: SpawnProcessTask) {
     mustRunAfter ':doctest:doctest'
+    pidLockFileName ".prom.pid.lock"
 
     doFirst {
+        // Kill any existing Prometheus processes to prevent "Server already running" error
+        exec {
+            commandLine 'sh', '-c', 'pkill -f "prometheus.*--storage.tsdb.path" || true'
+            ignoreExitValue true
+        }
+
         download.run {
             src getPrometheusBinaryLocation()
             dest new File("$projectDir/bin", 'prometheus.tar.gz')
@@ -308,9 +315,16 @@ task startPrometheus(type: SpawnProcessTask) {
 }
 
 task stopPrometheus(type: KillProcessTask) {
+    pidLockFileName ".prom.pid.lock"
     doLast {
         file("$projectDir/bin/prometheus").deleteDir()
         file("$projectDir/bin/prometheus.tar.gz").delete()
+
+        // Forcefully kill any remaining Prometheus processes
+        exec {
+            commandLine 'sh', '-c', 'pkill -f "prometheus.*--storage.tsdb.path" || true'
+            ignoreExitValue true
+        }
     }
 }
 

integ-test/src/test/java/org/opensearch/sql/legacy/PointInTimeLeakIT.java

@@ -0,0 +1,173 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+
+import java.io.IOException;
+import org.json.JSONArray;
+import org.json.JSONObject;
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
+import org.opensearch.client.ResponseException;
+import org.opensearch.sql.legacy.utils.StringUtils;
+
+/**
+ * Integration test verifying PIT contexts are created only when needed and properly cleaned up.
+ *
+ * @see <a href="https://github.com/opensearch-project/sql/issues/5002">Issue #5002</a>
+ */
+public class PointInTimeLeakIT extends SQLIntegTestCase {
+
+  private static final String TEST_INDEX = "test-logs-2025.01.01";
+  private static final String PIT_STATS_ENDPOINT =
+      "/_nodes/stats/indices/search?filter_path=nodes.*.indices.search.point_in_time_current";
+
+  @Before
+  public void setUpTestIndex() throws IOException {
+    try {
+      executeRequest(new Request("DELETE", "/" + TEST_INDEX));
+    } catch (ResponseException e) {
+      if (e.getResponse().getStatusLine().getStatusCode() != 404) {
+        throw e;
+      }
+    }
+
+    Request createIndex = new Request("PUT", "/" + TEST_INDEX);
+    createIndex.setJsonEntity(
+        "{  \"mappings\": {    \"properties\": {      \"action\": {\"type\": \"text\", \"fields\":"
+            + " {\"keyword\": {\"type\": \"keyword\"}}},      \"timestamp\": {\"type\": \"date\"}  "
+            + "  }  }}");
+    executeRequest(createIndex);
+
+    Request bulkRequest = new Request("POST", "/" + TEST_INDEX + "/_bulk");
+    bulkRequest.addParameter("refresh", "true");
+    bulkRequest.setJsonEntity(
+        "{\"index\":{}}\n"
+            + "{\"action\":\"login_success\",\"timestamp\":\"2025-01-01T10:00:00Z\"}\n"
+            + "{\"index\":{}}\n"
+            + "{\"action\":\"login_success\",\"timestamp\":\"2025-01-01T10:01:00Z\"}\n"
+            + "{\"index\":{}}\n"
+            + "{\"action\":\"login_failed\",\"timestamp\":\"2025-01-01T10:02:00Z\"}\n");
+    executeRequest(bulkRequest);
+  }
+
+  @Test
+  public void testNoPitLeakWithoutFetchSize() throws IOException, InterruptedException {
+    int baselinePitCount = getCurrentPitCount();
+
+    int numQueries = 10;
+
+    for (int i = 0; i < numQueries; i++) {
+      String query =
+          StringUtils.format(
+              "SELECT * FROM %s WHERE action LIKE 'login%%' ORDER BY timestamp ASC", TEST_INDEX);
+
+      JSONObject response = executeQueryWithoutFetchSize(query);
+
+      assertTrue("Query should succeed", response.has("datarows"));
+      JSONArray dataRows = response.getJSONArray("datarows");
+      assertThat("Should return results", dataRows.length(), greaterThan(0));
+      assertFalse("Should not have cursor for non-paginated query", response.has("cursor"));
+    }
+
+    int currentPitCount = getCurrentPitCount();
+    int leakedPits = currentPitCount - baselinePitCount;
+
+    assertThat("No PITs should leak after fix", leakedPits, equalTo(0));
+  }
+
+  @Test
+  public void testPitManagedProperlyWithFetchSize() throws IOException {
+    int baselinePitCount = getCurrentPitCount();
+
+    String query =
+        StringUtils.format(
+            "SELECT * FROM %s WHERE action LIKE 'login%%' ORDER BY timestamp ASC", TEST_INDEX);
+
+    JSONObject response = executeQueryWithFetchSize(query, 2);
+
+    assertTrue("Should have cursor with fetch_size", response.has("cursor"));
+    String cursor = response.getString("cursor");
+
+    JSONObject closeResponse = executeCursorCloseQuery(cursor);
+    assertTrue("Cursor close should succeed", closeResponse.getBoolean("succeeded"));
+
+    int finalPitCount = getCurrentPitCount();
+
+    assertThat(
+        "PIT should be cleaned up after cursor close", finalPitCount, equalTo(baselinePitCount));
+  }
+
+  @Test
+  public void testCompareV1AndV2EnginePitBehavior() throws IOException {
+    int baselinePitCount = getCurrentPitCount();
+
+    String v1Query =
+        StringUtils.format(
+            "SELECT * FROM %s WHERE action LIKE 'login%%' ORDER BY timestamp ASC", TEST_INDEX);
+
+    JSONObject v1Response = executeQueryWithoutFetchSize(v1Query);
+    int afterV1PitCount = getCurrentPitCount();
+    int v1Leaked = afterV1PitCount - baselinePitCount;
+
+    String v2Query =
+        StringUtils.format(
+            "SELECT * FROM `%s` WHERE action LIKE 'login%%' ORDER BY timestamp ASC", TEST_INDEX);
+
+    JSONObject v2Response = executeQueryWithoutFetchSize(v2Query);
+    int afterV2PitCount = getCurrentPitCount();
+    int v2Leaked = afterV2PitCount - afterV1PitCount;
+
+    assertTrue("V1 should return results", v1Response.has("datarows"));
+    assertTrue("V2 should return results", v2Response.has("datarows"));
+
+    assertThat("V1 Legacy SQL should not leak PITs", v1Leaked, equalTo(0));
+    assertThat("V2 SQL should not leak PITs", v2Leaked, equalTo(0));
+  }
+
+  private JSONObject executeQueryWithoutFetchSize(String query) throws IOException {
+    Request sqlRequest = new Request("POST", "/_plugins/_sql?format=jdbc");
+    sqlRequest.setJsonEntity(String.format("{\"query\": \"%s\"}", query));
+
+    Response response = client().performRequest(sqlRequest);
+    return new JSONObject(TestUtils.getResponseBody(response));
+  }
+
+  private JSONObject executeQueryWithFetchSize(String query, int fetchSize) throws IOException {
+    Request sqlRequest = new Request("POST", "/_plugins/_sql?format=jdbc");
+    sqlRequest.setJsonEntity(
+        String.format("{\"query\": \"%s\", \"fetch_size\": %d}", query, fetchSize));
+
+    Response response = client().performRequest(sqlRequest);
+    return new JSONObject(TestUtils.getResponseBody(response));
+  }
+
+  private int getCurrentPitCount() throws IOException {
+    Request statsRequest = new Request("GET", PIT_STATS_ENDPOINT);
+    Response response = client().performRequest(statsRequest);
+    JSONObject stats = new JSONObject(TestUtils.getResponseBody(response));
+
+    if (!stats.has("nodes")) {
+      return 0;
+    }
+
+    int totalPits = 0;
+    JSONObject nodes = stats.getJSONObject("nodes");
+    for (String nodeId : nodes.keySet()) {
+      Object pitValue =
+          stats.optQuery("/nodes/" + nodeId + "/indices/search/point_in_time_current");
+      if (pitValue instanceof Number) {
+        totalPits += ((Number) pitValue).intValue();
+      }
+    }
+
+    return totalPits;
+  }
+}