SOLR-2199: DataImportHandler (DIH) JdbcDataSource supports multiple resultsets per query

This commit is contained in:
Mikhail Khludnev 2016-06-12 23:50:06 +03:00
parent 664e39292b
commit fce10ae109
3 changed files with 116 additions and 8 deletions

View File

@ -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

View File

@ -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);
}
}

View File

@ -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);