From fce10ae1097fa7f764516f2b343365e86afc273d Mon Sep 17 00:00:00 2001 From: Mikhail Khludnev Date: Sun, 12 Jun 2016 23:50:06 +0300 Subject: [PATCH] SOLR-2199: DataImportHandler (DIH) JdbcDataSource supports multiple resultsets per query --- solr/CHANGES.txt | 2 + .../handler/dataimport/JdbcDataSource.java | 33 +++++-- .../dataimport/TestJdbcDataSource.java | 89 ++++++++++++++++++- 3 files changed, 116 insertions(+), 8 deletions(-) diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt index 659a1d7dbe4..2146539395a 100644 --- a/solr/CHANGES.txt +++ b/solr/CHANGES.txt @@ -315,6 +315,8 @@ Other Changes * SOLR-8445: fix line separator in log4j.properties files (Ahmet Arslan via Mikhail Khludnev) +* SOLR-2199: DataImportHandler (DIH) JdbcDataSource supports multiple resultsets per query (Kristine Jetzke, Mark Waddle via Mikhail Khludnev) + ================== 6.0.1 ================== Upgrade Notes diff --git a/solr/contrib/dataimporthandler/src/java/org/apache/solr/handler/dataimport/JdbcDataSource.java b/solr/contrib/dataimporthandler/src/java/org/apache/solr/handler/dataimport/JdbcDataSource.java index e1eabeb002b..09ad775a4ec 100644 --- a/solr/contrib/dataimporthandler/src/java/org/apache/solr/handler/dataimport/JdbcDataSource.java +++ b/solr/contrib/dataimporthandler/src/java/org/apache/solr/handler/dataimport/JdbcDataSource.java @@ -342,7 +342,16 @@ public class JdbcDataSource extends } protected ResultSet executeStatement(Statement statement, String query) throws SQLException { - if (statement.execute(query)) { + boolean resultSetReturned = statement.execute(query); + return getNextResultSet(resultSetReturned, statement); + } + + protected ResultSet getNextResultSet(final boolean initialResultSetAvailable, final Statement statement) throws SQLException { + boolean resultSetAvailable = initialResultSetAvailable; + while (!resultSetAvailable && statement.getUpdateCount() != -1) { + resultSetAvailable = statement.getMoreResults(); + } + if (resultSetAvailable) { return statement.getResultSet(); } return null; @@ -441,8 +450,10 @@ public class JdbcDataSource extends if (getResultSet().next()) { return true; } else { - close(); - return false; + closeResultSet(); + setResultSet(getNextResultSet(getStatement().getMoreResults(), getStatement())); + setColNames(getResultSet()); + return hasnext(); } } catch (SQLException e) { close(); @@ -452,16 +463,26 @@ public class JdbcDataSource extends } protected void close() { + closeResultSet(); try { - if (getResultSet() != null) - getResultSet().close(); if (getStatement() != null) getStatement().close(); + } catch (Exception e) { + logError("Exception while closing statement", e); + } finally { + setStatement(null); + } + } + + protected void closeResultSet() { + try { + if (getResultSet() != null) { + getResultSet().close(); + } } catch (Exception e) { logError("Exception while closing result set", e); } finally { setResultSet(null); - setStatement(null); } } diff --git a/solr/contrib/dataimporthandler/src/test/org/apache/solr/handler/dataimport/TestJdbcDataSource.java b/solr/contrib/dataimporthandler/src/test/org/apache/solr/handler/dataimport/TestJdbcDataSource.java index 38ca83ec29f..b6d05c441f1 100644 --- a/solr/contrib/dataimporthandler/src/test/org/apache/solr/handler/dataimport/TestJdbcDataSource.java +++ b/solr/contrib/dataimporthandler/src/test/org/apache/solr/handler/dataimport/TestJdbcDataSource.java @@ -19,7 +19,6 @@ package org.apache.solr.handler.dataimport; import java.io.File; import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.nio.file.Files; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; @@ -27,7 +26,12 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; -import java.util.*; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Properties; import javax.sql.DataSource; @@ -276,6 +280,7 @@ public class TestJdbcDataSource extends AbstractDataImportHandlerTestCase { statement.setFetchSize(500); statement.setMaxRows(0); EasyMock.expect(statement.execute("query")).andReturn(false); + EasyMock.expect(statement.getUpdateCount()).andReturn(-1); statement.close(); mockControl.replay(); @@ -388,6 +393,7 @@ public class TestJdbcDataSource extends AbstractDataImportHandlerTestCase { statement.setFetchSize(500); statement.setMaxRows(0); EasyMock.expect(statement.execute("other query")).andReturn(false); + EasyMock.expect(statement.getUpdateCount()).andReturn(-1); statement.close(); mockControl.replay(); @@ -398,6 +404,85 @@ public class TestJdbcDataSource extends AbstractDataImportHandlerTestCase { mockControl.verify(); } + @Test + public void testMultipleResultsSets_UpdateCountUpdateCountResultSet() throws Exception { + MockInitialContextFactory.bind("java:comp/env/jdbc/JndiDB", dataSource); + props.put(JdbcDataSource.JNDI_NAME, "java:comp/env/jdbc/JndiDB"); + EasyMock.expect(dataSource.getConnection()).andReturn(connection); + jdbcDataSource.init(context, props); + connection.setAutoCommit(false); + + Statement statement = mockControl.createMock(Statement.class); + EasyMock.expect(connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) + .andReturn(statement); + statement.setFetchSize(500); + statement.setMaxRows(0); + EasyMock.expect(statement.execute("query")).andReturn(false); + EasyMock.expect(statement.getUpdateCount()).andReturn(1); + EasyMock.expect(statement.getMoreResults()).andReturn(false); + EasyMock.expect(statement.getUpdateCount()).andReturn(1); + EasyMock.expect(statement.getMoreResults()).andReturn(true); + ResultSet resultSet = mockControl.createMock(ResultSet.class); + EasyMock.expect(statement.getResultSet()).andReturn(resultSet); + ResultSetMetaData metaData = mockControl.createMock(ResultSetMetaData.class); + EasyMock.expect(resultSet.getMetaData()).andReturn(metaData); + EasyMock.expect(metaData.getColumnCount()).andReturn(0); + + mockControl.replay(); + + final ResultSetIterator resultSetIterator = jdbcDataSource.new ResultSetIterator("query"); + assertSame(resultSet, resultSetIterator.getResultSet()); + + mockControl.verify(); + + } + + @Test + public void testMultipleResultsSets_ResultSetResultSet() throws Exception { + MockInitialContextFactory.bind("java:comp/env/jdbc/JndiDB", dataSource); + props.put(JdbcDataSource.JNDI_NAME, "java:comp/env/jdbc/JndiDB"); + EasyMock.expect(dataSource.getConnection()).andReturn(connection); + jdbcDataSource.init(context, props); + connection.setAutoCommit(false); + + Statement statement = mockControl.createMock(Statement.class); + EasyMock.expect(connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) + .andReturn(statement); + statement.setFetchSize(500); + statement.setMaxRows(0); + EasyMock.expect(statement.execute("query")).andReturn(true); + ResultSet resultSet1 = mockControl.createMock(ResultSet.class); + EasyMock.expect(statement.getResultSet()).andReturn(resultSet1); + ResultSetMetaData metaData1 = mockControl.createMock(ResultSetMetaData.class); + EasyMock.expect(resultSet1.getMetaData()).andReturn(metaData1); + EasyMock.expect(metaData1.getColumnCount()).andReturn(0); + EasyMock.expect(resultSet1.next()).andReturn(false); + resultSet1.close(); + EasyMock.expect(statement.getMoreResults()).andReturn(true); + ResultSet resultSet2 = mockControl.createMock(ResultSet.class); + EasyMock.expect(statement.getResultSet()).andReturn(resultSet2); + ResultSetMetaData metaData2 = mockControl.createMock(ResultSetMetaData.class); + EasyMock.expect(resultSet2.getMetaData()).andReturn(metaData2); + EasyMock.expect(metaData2.getColumnCount()).andReturn(0); + EasyMock.expect(resultSet2.next()).andReturn(true); + EasyMock.expect(resultSet2.next()).andReturn(false); + resultSet2.close(); + EasyMock.expect(statement.getMoreResults()).andReturn(false); + EasyMock.expect(statement.getUpdateCount()).andReturn(-1); + statement.close(); + + mockControl.replay(); + + final ResultSetIterator resultSetIterator = jdbcDataSource.new ResultSetIterator("query"); + assertSame(resultSet1, resultSetIterator.getResultSet()); + assertTrue(resultSetIterator.hasnext()); + assertSame(resultSet2, resultSetIterator.getResultSet()); + assertFalse(resultSetIterator.hasnext()); + + mockControl.verify(); + + } + @Test public void testRetrieveFromDriverManager() throws Exception { DriverManager.registerDriver(driver);