package groovy.sql;
import groovy.lang.Closure;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.Collection;
import java.util.logging.Level;
import org.teiid.client.plan.PlanNode;
import org.teiid.jdbc.TeiidStatement;
/**
* An extension of Groovy's Sql to support getting {@link TeiidStatement} specific properties.
*/
public final class TeiidSql extends Sql {
private int maxRows;
private PlanNode plan;
private Collection<org.teiid.client.plan.Annotation> annotations;
private String debugLog;
private SQLWarning warnings;
public TeiidSql(Connection connection) {
super(connection);
}
@Override
protected void configure(Statement statement) {
super.configure(statement);
try {
statement.setMaxRows(maxRows);
} catch (SQLException e) {
LOG.log(Level.WARNING, "Caught exception setting max rows: " + e, e);
}
plan = null;
annotations = null;
debugLog = null;
warnings = null;
}
/**
* Overridden to fix passing the resultset to the closure
*/
@Override
public void eachRow(String sql, Closure metaClosure, Closure rowClosure) throws SQLException {
AbstractQueryCommand command = createQueryCommand(sql);
ResultSet results = null;
try {
results = command.execute();
if (metaClosure != null) metaClosure.call(results.getMetaData());
GroovyResultSet groovyRS = new GroovyResultSetProxy(results).getImpl();
while (groovyRS.next()) {
rowClosure.call(groovyRS);
}
} catch (SQLException e) {
LOG.warning("Failed to execute: " + sql + " because: " + e.getMessage());
throw e;
} finally {
command.closeResources();
}
}
@Override
protected void closeResources(Connection connection,
Statement statement) {
getPlanInfo(statement);
super.closeResources(connection, statement);
}
@Override
protected void closeResources(Connection connection,
Statement statement, ResultSet results) {
getPlanInfo(statement);
super.closeResources(connection, statement, results);
}
protected void getPlanInfo(Statement s) {
if (s == null) {
return;
}
TeiidStatement ts;
try {
ts = s.unwrap(TeiidStatement.class);
} catch (SQLException e) {
throw new RuntimeException(e);
}
plan = ts.getPlanDescription();
annotations = ts.getAnnotations();
debugLog = ts.getDebugLog();
try {
warnings = ts.getWarnings();
} catch (SQLException e) {
LOG.log(Level.WARNING, "Caught exception getting warnings: " + e, e);
}
}
public SQLWarning getSQLWarnings() {
return warnings;
}
public int getMaxRows() {
return maxRows;
}
public void setMaxRows(int maxRows) {
this.maxRows = maxRows;
}
public PlanNode getPlan() {
return plan;
}
public Collection<org.teiid.client.plan.Annotation> getAnnotations() {
return annotations;
}
public String getDebugLog() {
return debugLog;
}
}