From c798324d17b4b1db47a22e4aa214ed6d8ddf9780 Mon Sep 17 00:00:00 2001 From: peacewong Date: Fri, 13 Sep 2024 20:42:51 +0800 Subject: [PATCH] JDBC Driver support pull Multiple result sets --- .../linkis/ujes/jdbc/UJESSQLDriver.java | 1 + .../ujes/jdbc/LinkisSQLConnection.scala | 2 - .../linkis/ujes/jdbc/LinkisSQLStatement.scala | 48 ++++- .../linkis/ujes/jdbc/UJESSQLDriverMain.scala | 4 + .../linkis/ujes/jdbc/UJESSQLResultSet.scala | 57 +++-- .../ujes/jdbc/LinkisSQLStatementTest.java | 194 +++++++++++++++++- .../ujes/jdbc/UJESSQLResultSetTest.java | 108 ++++++++++ 7 files changed, 374 insertions(+), 40 deletions(-) diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/main/java/org/apache/linkis/ujes/jdbc/UJESSQLDriver.java b/linkis-computation-governance/linkis-jdbc-driver/src/main/java/org/apache/linkis/ujes/jdbc/UJESSQLDriver.java index dd12d1414c..0bc0b08c52 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/main/java/org/apache/linkis/ujes/jdbc/UJESSQLDriver.java +++ b/linkis-computation-governance/linkis-jdbc-driver/src/main/java/org/apache/linkis/ujes/jdbc/UJESSQLDriver.java @@ -50,6 +50,7 @@ public class UJESSQLDriver extends UJESSQLDriverMain implements Driver { static String PASSWORD = "password"; static boolean TABLEAU_SERVER = false; static String FIXED_SESSION = "fixedSession"; + static String ENABLE_MULTI_RESULT = "enableMultiResult"; static String USE_SSL = "useSSL"; static String VERSION = "version"; diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLConnection.scala b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLConnection.scala index f75db82cdf..0be96b2c15 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLConnection.scala +++ b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLConnection.scala @@ -448,8 +448,6 @@ class LinkisSQLConnection(private[jdbc] val ujesClient: UJESClient, props: Prope val runType = EngineType.mapStringToEngineType(engine) match { case EngineType.SPARK => RunType.SQL case EngineType.HIVE => RunType.HIVE - case EngineType.REPL => RunType.REPL - case EngineType.DORIS => RunType.DORIS case EngineType.TRINO => RunType.TRINO_SQL case EngineType.PRESTO => RunType.PRESTO_SQL case EngineType.NEBULA => RunType.NEBULA_SQL diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLStatement.scala b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLStatement.scala index f00d870978..e3a1475d2b 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLStatement.scala +++ b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/LinkisSQLStatement.scala @@ -37,6 +37,10 @@ class LinkisSQLStatement(private[jdbc] val ujesSQLConnection: LinkisSQLConnectio with Logging { private var jobExecuteResult: JobExecuteResult = _ + + private val openedResultSets: util.ArrayList[UJESSQLResultSet] = + new util.ArrayList[UJESSQLResultSet]() + private var resultSet: UJESSQLResultSet = _ private var closed = false private var maxRows: Int = 0 @@ -190,7 +194,7 @@ class LinkisSQLStatement(private[jdbc] val ujesSQLConnection: LinkisSQLConnectio override def getUpdateCount: Int = throwWhenClosed(-1) - override def getMoreResults: Boolean = false + override def getMoreResults: Boolean = getMoreResults(Statement.CLOSE_CURRENT_RESULT) override def setFetchDirection(direction: Int): Unit = throwWhenClosed(if (direction != ResultSet.FETCH_FORWARD) { @@ -230,7 +234,45 @@ class LinkisSQLStatement(private[jdbc] val ujesSQLConnection: LinkisSQLConnectio override def getConnection: Connection = throwWhenClosed(ujesSQLConnection) - override def getMoreResults(current: Int): Boolean = false + override def getMoreResults(current: Int): Boolean = { + if (this.resultSet == null) { + false + } else { + this.resultSet.getMetaData + val nextResultSet = this.resultSet.getNextResultSet + current match { + case Statement.CLOSE_CURRENT_RESULT => + // 1 - CLOSE CURRENT RESULT SET + this.resultSet.close() + this.resultSet.clearNextResultSet + case Statement.KEEP_CURRENT_RESULT => + // 2 - KEEP CURRENT RESULT SET + this.openedResultSets.add(this.resultSet) + this.resultSet.clearNextResultSet + case Statement.CLOSE_ALL_RESULTS => + // 3 - CLOSE ALL RESULT SET + this.openedResultSets.add(this.resultSet) + closeAllOpenedResultSet() + case _ => + throw new LinkisSQLException( + LinkisSQLErrorCode.NOSUPPORT_STATEMENT, + "getMoreResults with current not in 1,2,3 is not supported, see Statement.getMoreResults" + ) + } + this.resultSet = nextResultSet + this.resultSet != null + } + } + + private def closeAllOpenedResultSet(): Any = { + val iterator = this.openedResultSets.iterator() + while (iterator.hasNext) { + val set = iterator.next() + if (!set.isClosed) { + set.close() + } + } + } override def getGeneratedKeys: ResultSet = throw new LinkisSQLException( LinkisSQLErrorCode.NOSUPPORT_STATEMENT, @@ -302,6 +344,7 @@ class LinkisSQLStatement(private[jdbc] val ujesSQLConnection: LinkisSQLConnectio /** * log[0] error log[1] warn log[2] info log[3] all (info + warn + error) + * * @return */ def getAllLog(): Array[String] = { @@ -316,6 +359,7 @@ class LinkisSQLStatement(private[jdbc] val ujesSQLConnection: LinkisSQLConnectio /** * log[0] error log[1] warn log[2] info log[3] all (info + warn + error) + * * @return */ def getIncrementalLog(): util.List[String] = { diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLDriverMain.scala b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLDriverMain.scala index c162b8c924..44686981e8 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLDriverMain.scala +++ b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLDriverMain.scala @@ -78,6 +78,9 @@ class UJESSQLDriverMain extends Driver with Logging { case Array(USE_SSL, value) => props.setProperty(USE_SSL, value) false + case Array(ENABLE_MULTI_RESULT, value) => + props.setProperty(ENABLE_MULTI_RESULT, value) + false case Array(key, _) => if (StringUtils.isBlank(key)) { throw new LinkisSQLException( @@ -141,6 +144,7 @@ object UJESSQLDriverMain { val PASSWORD = UJESSQLDriver.PASSWORD val TABLEAU_SERVER = UJESSQLDriver.TABLEAU_SERVER val FIXED_SESSION = UJESSQLDriver.FIXED_SESSION + val ENABLE_MULTI_RESULT = UJESSQLDriver.ENABLE_MULTI_RESULT val USE_SSL = UJESSQLDriver.USE_SSL diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLResultSet.scala b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLResultSet.scala index 39418a42d1..02e8551722 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLResultSet.scala +++ b/linkis-computation-governance/linkis-jdbc-driver/src/main/scala/org/apache/linkis/ujes/jdbc/UJESSQLResultSet.scala @@ -20,6 +20,7 @@ package org.apache.linkis.ujes.jdbc import org.apache.linkis.common.utils.Logging import org.apache.linkis.ujes.client.request.ResultSetAction import org.apache.linkis.ujes.client.response.ResultSetResult +import org.apache.linkis.ujes.client.utils.UJESClientUtils import org.apache.commons.lang3.StringUtils @@ -76,6 +77,7 @@ class UJESSQLResultSet( private var path: String = _ private var metaData: util.List[util.Map[String, String]] = _ private val statement: LinkisSQLStatement = ujesStatement + private var nextResultSet: UJESSQLResultSet = _ private val connection: LinkisSQLConnection = ujesStatement.getConnection.asInstanceOf[LinkisSQLConnection] @@ -102,7 +104,15 @@ class UJESSQLResultSet( private def getResultSetPath(resultSetList: Array[String]): String = { if (resultSetList.length > 0) { - resultSetList(resultSetList.length - 1) + val enableMultiResult = connection.getProps.getProperty(UJESSQLDriverMain.ENABLE_MULTI_RESULT) + enableMultiResult match { + case "Y" => + // 配置开启时，返回首个结果集 + resultSetList(0) + case _ => + // 配置关闭时，返回以最后一个结果集为准 + resultSetList(resultSetList.length - 1) + } } else { "" } @@ -110,6 +120,12 @@ class UJESSQLResultSet( private def resultSetResultInit(): Unit = { if (path == null) path = getResultSetPath(resultSetList) + // 设置下一个结果集 + val enableMultiResult = connection.getProps.getProperty(UJESSQLDriverMain.ENABLE_MULTI_RESULT) + if (resultSetList.length > 1 && "Y".equals(enableMultiResult)) { + this.nextResultSet = + new UJESSQLResultSet(resultSetList.drop(1), this.statement, maxRows, fetchSize) + } val user = connection.getProps.getProperty("user") if (StringUtils.isNotBlank(path)) { val resultAction = @@ -235,38 +251,7 @@ class UJESSQLResultSet( } private def evaluate(dataType: String, value: String): Any = { - - if (value == null || value.equals("null") || value.equals("NULL") || value.equals("Null")) { - dataType.toLowerCase(Locale.getDefault) match { - case "string" | "char" | "varchar" | "nvarchar" => value - case _ => null - } - } else { - dataType.toLowerCase(Locale.getDefault) match { - case null => throw new LinkisSQLException(LinkisSQLErrorCode.METADATA_EMPTY) - case "char" | "varchar" | "nvarchar" | "string" => value - case "short" => value.toShort - case "smallint" => value.toShort - case "tinyint" => value.toShort - case "int" => value.toInt - case "long" => value.toLong - case "float" => value.toFloat - case "double" => value.toDouble - case "boolean" => value.toBoolean - case "byte" => value.toByte - case "timestamp" => value - case "date" => value - case "bigint" => value.toLong - case "decimal" => value.toDouble - case "array" => value.toArray - case "map" => value - case _ => - throw new LinkisSQLException( - LinkisSQLErrorCode.PREPARESTATEMENT_TYPEERROR, - s"Can't infer the SQL type to use for an instance of ${dataType}. Use getObject() with an explicit Types value to specify the type to use" - ) - } - } + UJESClientUtils.evaluate(dataType, value) } private def getColumnValue(columnIndex: Int): Any = { @@ -303,6 +288,10 @@ class UJESSQLResultSet( } } + def clearNextResultSet: Any = { + this.nextResultSet = null + } + override def getBoolean(columnIndex: Int): Boolean = { val any = getColumnValue(columnIndex) if (wasNull()) { @@ -683,6 +672,8 @@ class UJESSQLResultSet( true } + def getNextResultSet: UJESSQLResultSet = this.nextResultSet + override def setFetchDirection(direction: Int): Unit = { throw new LinkisSQLException(LinkisSQLErrorCode.NOSUPPORT_RESULTSET) } diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/LinkisSQLStatementTest.java b/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/LinkisSQLStatementTest.java index 3ebd21ae70..e319cd0254 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/LinkisSQLStatementTest.java +++ b/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/LinkisSQLStatementTest.java @@ -17,16 +17,26 @@ package org.apache.linkis.ujes.jdbc; +import org.apache.linkis.governance.common.entity.ExecutionNodeStatus; +import org.apache.linkis.governance.common.entity.task.RequestPersistTask; +import org.apache.linkis.ujes.client.UJESClient; +import org.apache.linkis.ujes.client.response.JobExecuteResult; +import org.apache.linkis.ujes.client.response.JobInfoResult; +import org.apache.linkis.ujes.client.response.ResultSetResult; + import java.sql.SQLException; +import java.sql.Statement; +import java.util.Properties; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; /* * Notice: @@ -143,6 +153,184 @@ public void getConnWhenIsClosed() { } } + /** + * single query without next result set check point 1: getMoreResults returns false check point 2: + * default getMoreResults, use Statement.CLOSE_CURRENT_RESULT. The current result set is closed. + */ + @Test + public void singleQueryWithNoMoreResult() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + LinkisSQLConnection linkisSQLConnection = Mockito.spy(new LinkisSQLConnection(ujesClient, t)); + LinkisSQLStatement linkisSQLStatement = new LinkisSQLStatement(linkisSQLConnection); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + + JobExecuteResult jobExecuteResult = new JobExecuteResult(); + Mockito.doReturn(jobExecuteResult).when(linkisSQLConnection).toSubmit(anyString()); + JobInfoResult jobInfoResult = Mockito.spy(new JobInfoResult()); + Mockito.when(ujesClient.getJobInfo(jobExecuteResult)).thenReturn(jobInfoResult); + Mockito.doReturn(ExecutionNodeStatus.Succeed.name()).when(jobInfoResult).getJobStatus(); + Mockito.doReturn(new RequestPersistTask()).when(jobInfoResult).getRequestPersistTask(); + + Mockito.doReturn(new String[] {"path 1"}).when(jobInfoResult).getResultSetList(ujesClient); + + linkisSQLStatement.execute("select 1"); + UJESSQLResultSet resultSet = linkisSQLStatement.getResultSet(); + assertNotNull(resultSet); + assertFalse(resultSet.isClosed()); + // it will close current result set with default value 1 + boolean moreResults = linkisSQLStatement.getMoreResults(); + assertFalse(moreResults); + assertTrue(resultSet.isClosed()); + } + + /** + * multiple query without multiple result param, return one result check point 1: 2 sql executed. + * 1 result set + */ + @Test + public void multiQueryWithNoMoreResult() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + LinkisSQLConnection linkisSQLConnection = Mockito.spy(new LinkisSQLConnection(ujesClient, t)); + LinkisSQLStatement linkisSQLStatement = new LinkisSQLStatement(linkisSQLConnection); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + JobExecuteResult jobExecuteResult = new JobExecuteResult(); + Mockito.doReturn(jobExecuteResult).when(linkisSQLConnection).toSubmit(anyString()); + JobInfoResult jobInfoResult = Mockito.spy(new JobInfoResult()); + Mockito.when(ujesClient.getJobInfo(jobExecuteResult)).thenReturn(jobInfoResult); + Mockito.doReturn(ExecutionNodeStatus.Succeed.name()).when(jobInfoResult).getJobStatus(); + Mockito.doReturn(new RequestPersistTask()).when(jobInfoResult).getRequestPersistTask(); + + Mockito.doReturn(new String[] {"path 1", "path 2"}) + .when(jobInfoResult) + .getResultSetList(ujesClient); + + linkisSQLStatement.execute("select 1;select 2;"); + UJESSQLResultSet resultSet = linkisSQLStatement.getResultSet(); + assertNotNull(resultSet); + assertFalse(resultSet.isClosed()); + // it will close current result set with default value 1 + boolean moreResults = linkisSQLStatement.getMoreResults(); + assertFalse(moreResults); + assertTrue(resultSet.isClosed()); + } + + /** + * multiple query executed with multiple result param is Y check point 1: getMoreResults returns + * true check point 2: current result is closed check point 3: second getMoreResults returns false + */ + @Test + public void multiQueryWithMoreResult() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + t.put(UJESSQLDriverMain.ENABLE_MULTI_RESULT(), "Y"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + LinkisSQLConnection linkisSQLConnection = Mockito.spy(new LinkisSQLConnection(ujesClient, t)); + LinkisSQLStatement linkisSQLStatement = new LinkisSQLStatement(linkisSQLConnection); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + + JobExecuteResult jobExecuteResult = new JobExecuteResult(); + Mockito.doReturn(jobExecuteResult).when(linkisSQLConnection).toSubmit(anyString()); + JobInfoResult jobInfoResult = Mockito.spy(new JobInfoResult()); + Mockito.when(ujesClient.getJobInfo(jobExecuteResult)).thenReturn(jobInfoResult); + Mockito.doReturn(ExecutionNodeStatus.Succeed.name()).when(jobInfoResult).getJobStatus(); + Mockito.doReturn(new RequestPersistTask()).when(jobInfoResult).getRequestPersistTask(); + + Mockito.doReturn(new String[] {"path 1", "path 2"}) + .when(jobInfoResult) + .getResultSetList(ujesClient); + + linkisSQLStatement.execute("select 1;select 2;"); + UJESSQLResultSet resultSet = linkisSQLStatement.getResultSet(); + assertNotNull(resultSet); + assertFalse(resultSet.isClosed()); + // it will close current result set with default value 1 + boolean moreResults = linkisSQLStatement.getMoreResults(); + assertTrue(moreResults); + assertTrue(resultSet.isClosed()); + moreResults = linkisSQLStatement.getMoreResults(); + assertFalse(moreResults); + } + + /** + * multiple query executed with multiple result param is Y, and use + * LinkisSQLStatement.KEEP_CURRENT_RESULT check point 1: getMoreResults returns true check point + * 2: current result is not close check point 3: second getMoreResults returns false + */ + @Test + public void multiQueryWithMoreResultNotCloseCurrent() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + t.put(UJESSQLDriverMain.ENABLE_MULTI_RESULT(), "Y"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + LinkisSQLConnection linkisSQLConnection = Mockito.spy(new LinkisSQLConnection(ujesClient, t)); + LinkisSQLStatement linkisSQLStatement = new LinkisSQLStatement(linkisSQLConnection); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + + JobExecuteResult jobExecuteResult = new JobExecuteResult(); + Mockito.doReturn(jobExecuteResult).when(linkisSQLConnection).toSubmit(anyString()); + JobInfoResult jobInfoResult = Mockito.spy(new JobInfoResult()); + Mockito.when(ujesClient.getJobInfo(jobExecuteResult)).thenReturn(jobInfoResult); + Mockito.doReturn(ExecutionNodeStatus.Succeed.name()).when(jobInfoResult).getJobStatus(); + Mockito.doReturn(new RequestPersistTask()).when(jobInfoResult).getRequestPersistTask(); + + Mockito.doReturn(new String[] {"path 1", "path 2"}) + .when(jobInfoResult) + .getResultSetList(ujesClient); + + linkisSQLStatement.execute("select 1;select 2;"); + UJESSQLResultSet resultSet = linkisSQLStatement.getResultSet(); + assertNotNull(resultSet); + assertFalse(resultSet.isClosed()); + // it will close current result set with default value 1 + boolean moreResults = linkisSQLStatement.getMoreResults(LinkisSQLStatement.KEEP_CURRENT_RESULT); + assertTrue(moreResults); + assertFalse(resultSet.isClosed()); + } + + /** + * multiple query executed with multiple result param is Y, and use + * LinkisSQLStatement.CLOSE_ALL_RESULTS check point 1: getMoreResults returns true check point 2: + * current result is not close check point 3: second getMoreResults returns false check point 4: + * first result set is closed after second invoke getMoreResults + */ + @Test + public void multiQueryWithMoreResultCloseAllOpenedCurrent() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + t.put(UJESSQLDriverMain.ENABLE_MULTI_RESULT(), "Y"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + LinkisSQLConnection linkisSQLConnection = Mockito.spy(new LinkisSQLConnection(ujesClient, t)); + LinkisSQLStatement linkisSQLStatement = new LinkisSQLStatement(linkisSQLConnection); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + + JobExecuteResult jobExecuteResult = new JobExecuteResult(); + Mockito.doReturn(jobExecuteResult).when(linkisSQLConnection).toSubmit(anyString()); + JobInfoResult jobInfoResult = Mockito.spy(new JobInfoResult()); + Mockito.when(ujesClient.getJobInfo(jobExecuteResult)).thenReturn(jobInfoResult); + Mockito.doReturn(ExecutionNodeStatus.Succeed.name()).when(jobInfoResult).getJobStatus(); + Mockito.doReturn(new RequestPersistTask()).when(jobInfoResult).getRequestPersistTask(); + + Mockito.doReturn(new String[] {"path 1", "path 2"}) + .when(jobInfoResult) + .getResultSetList(ujesClient); + + linkisSQLStatement.execute("select 1;select 2;"); + UJESSQLResultSet resultSet = linkisSQLStatement.getResultSet(); + assertNotNull(resultSet); + assertFalse(resultSet.isClosed()); + // it will close current result set with default value 1 + boolean moreResults = linkisSQLStatement.getMoreResults(Statement.KEEP_CURRENT_RESULT); + assertTrue(moreResults); + assertFalse(resultSet.isClosed()); + moreResults = linkisSQLStatement.getMoreResults(Statement.CLOSE_ALL_RESULTS); + assertFalse(moreResults); + assertTrue(resultSet.isClosed()); + } + @AfterAll public static void closeStateAndConn() { if (statement != null) { diff --git a/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/UJESSQLResultSetTest.java b/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/UJESSQLResultSetTest.java index a8f0a179d0..c0631427ea 100644 --- a/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/UJESSQLResultSetTest.java +++ b/linkis-computation-governance/linkis-jdbc-driver/src/test/java/org/apache/linkis/ujes/jdbc/UJESSQLResultSetTest.java @@ -17,7 +17,14 @@ package org.apache.linkis.ujes.jdbc; +import org.apache.linkis.ujes.client.UJESClient; +import org.apache.linkis.ujes.client.request.ResultSetAction; +import org.apache.linkis.ujes.client.response.ResultSetResult; + import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; @@ -25,6 +32,10 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; /* * Notice: @@ -137,4 +148,101 @@ public void next() { Assertions.assertTrue(resultSet.isAfterLast()); } } + + /** single query result with no multiple result set check point 1: nextResultSet is null */ + @Test + public void singleQueryWithNoMoreResultSet() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + Mockito.when(ujesClient.resultSet(any())).thenReturn(new ResultSetResult()); + + LinkisSQLConnection linkisSQLConnection = new LinkisSQLConnection(ujesClient, t); + + UJESSQLResultSet ujessqlResultSet = + new UJESSQLResultSet( + new String[] {"path1"}, new LinkisSQLStatement(linkisSQLConnection), 0, 0); + + ujessqlResultSet.next(); + + assertNull(ujessqlResultSet.getNextResultSet()); + } + + /** + * multiple result set with multi result switch is Y check point 1: queryResult has two path, + * return first path. check point 2: the second result set returned check point 3: the third + * result set is null + */ + @Test + public void nultiQueryWithMoreResultSet() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + t.put(UJESSQLDriverMain.ENABLE_MULTI_RESULT(), "Y"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + List pathList = new ArrayList<>(); + Mockito.when(ujesClient.resultSet(any())) + .thenAnswer( + invocationOnMock -> { + ResultSetAction argument = invocationOnMock.getArgument(0); + String path = (String) argument.getParameters().get("path"); + if (pathList.isEmpty()) { + assertEquals("path1", path); + } + pathList.add(path); + + return new ResultSetResult(); + }); + LinkisSQLConnection linkisSQLConnection = new LinkisSQLConnection(ujesClient, t); + + UJESSQLResultSet ujessqlResultSet = + new UJESSQLResultSet( + new String[] {"path1", "path2"}, new LinkisSQLStatement(linkisSQLConnection), 0, 0); + + // 查询 + ujessqlResultSet.next(); + + // 存在下一个结果集 + UJESSQLResultSet nextResultSet = ujessqlResultSet.getNextResultSet(); + assertNotNull(nextResultSet); + nextResultSet.next(); + + // 不存在第三个结果集 + assertNull(nextResultSet.getNextResultSet()); + } + + /** + * multiple result set with multi result switch not Y check point 1: queryResult has two path, + * return last path. check point 2: the next result set is null + */ + @Test + public void nultiQueryWithNoMoreResultSet() { + Properties t = new Properties(); + t.put("user", "hiveUser"); + UJESClient ujesClient = Mockito.mock(UJESClient.class); + Mockito.when(ujesClient.resultSet(any())) + .thenAnswer( + invocationOnMock -> { + ResultSetAction argument = invocationOnMock.getArgument(0); + String path = (String) argument.getParameters().get("path"); + assertEquals("path4", path); + + return new ResultSetResult(); + }); + + LinkisSQLConnection linkisSQLConnection = new LinkisSQLConnection(ujesClient, t); + + UJESSQLResultSet ujessqlResultSet = + new UJESSQLResultSet( + new String[] {"path1", "path2", "path3", "path4"}, + new LinkisSQLStatement(linkisSQLConnection), + 0, + 0); + + // 查询 + ujessqlResultSet.next(); + + // 即使查询有多个结果集，也不会产生多个结果集返回 + UJESSQLResultSet nextResultSet = ujessqlResultSet.getNextResultSet(); + assertNull(nextResultSet); + } }