exceptionfactory commented on code in PR #7496: URL: https://github.com/apache/nifi/pull/7496#discussion_r1267999623
########## nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/calcite/RecordResultSetOutputStreamCallback.java: ########## @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.standard.calcite; + +import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.io.OutputStreamCallback; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.serialization.RecordSetWriter; +import org.apache.nifi.serialization.RecordSetWriterFactory; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.serialization.record.ResultSetRecordSet; + +import java.io.IOException; +import java.io.OutputStream; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class RecordResultSetOutputStreamCallback implements OutputStreamCallback { + private final ComponentLog logger; + private final ResultSet rs; + private final RecordSchema writerSchema; + private final Integer defaultPrecision; + private final Integer defaultScale; + private final RecordSetWriterFactory recordSetWriterFactory; + private final Map<String, String> originalAttributes; + private final AtomicReference<WriteResult> writeResultRef; + private final AtomicReference<String> mimeTypeRef; Review Comment: As discussed in an offline conversation, removing these passed Atomic References and add get methods for the WriteResult and MIME Type properties would help clean up this callback class. ########## nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/calcite/TestRecordResultSetOutputStreamCallback.java: ########## @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.standard.calcite; + +import org.apache.calcite.adapter.java.ReflectiveSchema; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.schema.impl.AbstractSchema; +import org.apache.nifi.csv.CSVRecordSetWriter; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processors.standard.QueryRecord; +import org.apache.nifi.reporting.InitializationException; +import org.apache.nifi.serialization.RecordSetWriterFactory; +import org.apache.nifi.serialization.SimpleRecordSchema; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.RecordField; +import org.apache.nifi.serialization.record.RecordFieldType; +import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.util.TestRunner; +import org.apache.nifi.util.TestRunners; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.concurrent.atomic.AtomicReference; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class TestRecordResultSetOutputStreamCallback { + + @Test + void testResultSetClosed() throws IOException, SQLException, ClassNotFoundException, InitializationException { + final TestRunner runner = TestRunners.newTestRunner(QueryRecord.class); + + final String writerId = "record-writer"; + runner.setProperty(writerId, writerId); + + final RecordSetWriterFactory writerService = new CSVRecordSetWriter(); + runner.addControllerService(writerId, writerService); + runner.setProperty(writerService, "schema-access-strategy", "inherit-record-schema"); + runner.enableControllerService(writerService); + + final ResultSet resultSet = getResultSet(); + + final RecordField fieldFirst = new RecordField("first", RecordFieldType.STRING.getDataType()); + final RecordField fieldLast = new RecordField("last", RecordFieldType.STRING.getDataType()); + final RecordSchema writerSchema = new SimpleRecordSchema(Arrays.asList(fieldFirst, fieldLast)); + + final FlowFile flowFile = mock(FlowFile.class); + when(flowFile.getAttributes()).thenReturn(new LinkedHashMap<>()); + + final AtomicReference<WriteResult> writeResultRef = new AtomicReference<>(); + final AtomicReference<String> mimeTypeRef = new AtomicReference<>(); + + final RecordResultSetOutputStreamCallback writer = new RecordResultSetOutputStreamCallback(runner.getLogger(), + resultSet, writerSchema, 0, 0, writerService, flowFile.getAttributes(), writeResultRef, mimeTypeRef); + + final OutputStream os = new ByteArrayOutputStream(); + writer.process(os); + + Assertions.assertTrue(resultSet.isClosed()); + } + + private ResultSet getResultSet() throws ClassNotFoundException, SQLException { + Class.forName("org.apache.calcite.jdbc.Driver"); Review Comment: Instead of using the string class name, an actual instance of the Driver class could be passed to DriverManager.registerDriver() ########## nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/calcite/RecordResultSetOutputStreamCallback.java: ########## @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.standard.calcite; + +import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.io.OutputStreamCallback; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.serialization.RecordSetWriter; +import org.apache.nifi.serialization.RecordSetWriterFactory; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.serialization.record.ResultSetRecordSet; + +import java.io.IOException; +import java.io.OutputStream; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class RecordResultSetOutputStreamCallback implements OutputStreamCallback { + private final ComponentLog logger; + private final ResultSet rs; + private final RecordSchema writerSchema; + private final Integer defaultPrecision; + private final Integer defaultScale; + private final RecordSetWriterFactory recordSetWriterFactory; + private final Map<String, String> originalAttributes; + private final AtomicReference<WriteResult> writeResultRef; + private final AtomicReference<String> mimeTypeRef; + + public RecordResultSetOutputStreamCallback( + final ComponentLog logger, final ResultSet rs, final RecordSchema writerSchema, + final Integer defaultPrecision, final Integer defaultScale, + final RecordSetWriterFactory recordSetWriterFactory, final Map<String, String> originalAttributes, + final AtomicReference<WriteResult> writeResultRef, final AtomicReference<String> mimeTypeRef) { + this.logger = logger; + this.rs = rs; + this.writerSchema = writerSchema; + this.defaultPrecision = defaultPrecision; + this.defaultScale = defaultScale; + this.recordSetWriterFactory = recordSetWriterFactory; + this.originalAttributes = originalAttributes; + this.writeResultRef = writeResultRef; + this.mimeTypeRef = mimeTypeRef; + } + + @Override + public void process(OutputStream out) throws IOException { + final RecordSchema writeSchema; + + try (final ResultSetRecordSet recordSet = new ResultSetRecordSet(rs, writerSchema, defaultPrecision, defaultScale)) { + final RecordSchema resultSetSchema = recordSet.getSchema(); + writeSchema = recordSetWriterFactory.getSchema(originalAttributes, resultSetSchema); + + try (final RecordSetWriter resultSetWriter = recordSetWriterFactory.createWriter(logger, writeSchema, out, originalAttributes)) { + writeResultRef.set(resultSetWriter.write(recordSet)); + mimeTypeRef.set(resultSetWriter.getMimeType()); + } catch (final Exception e) { + throw new IOException(e); + } + } catch (final SQLException | SchemaNotFoundException e) { + throw new ProcessException(e); Review Comment: Recommend adding a message to this exception, indicating something like, "Reading query result records failed". ########## nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/calcite/RecordResultSetOutputStreamCallback.java: ########## @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.standard.calcite; + +import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.io.OutputStreamCallback; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.serialization.RecordSetWriter; +import org.apache.nifi.serialization.RecordSetWriterFactory; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.serialization.record.ResultSetRecordSet; + +import java.io.IOException; +import java.io.OutputStream; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class RecordResultSetOutputStreamCallback implements OutputStreamCallback { + private final ComponentLog logger; + private final ResultSet rs; + private final RecordSchema writerSchema; + private final Integer defaultPrecision; + private final Integer defaultScale; + private final RecordSetWriterFactory recordSetWriterFactory; + private final Map<String, String> originalAttributes; + private final AtomicReference<WriteResult> writeResultRef; + private final AtomicReference<String> mimeTypeRef; + + public RecordResultSetOutputStreamCallback( + final ComponentLog logger, final ResultSet rs, final RecordSchema writerSchema, + final Integer defaultPrecision, final Integer defaultScale, + final RecordSetWriterFactory recordSetWriterFactory, final Map<String, String> originalAttributes, + final AtomicReference<WriteResult> writeResultRef, final AtomicReference<String> mimeTypeRef) { + this.logger = logger; + this.rs = rs; + this.writerSchema = writerSchema; + this.defaultPrecision = defaultPrecision; + this.defaultScale = defaultScale; + this.recordSetWriterFactory = recordSetWriterFactory; + this.originalAttributes = originalAttributes; + this.writeResultRef = writeResultRef; + this.mimeTypeRef = mimeTypeRef; + } + + @Override + public void process(OutputStream out) throws IOException { + final RecordSchema writeSchema; + + try (final ResultSetRecordSet recordSet = new ResultSetRecordSet(rs, writerSchema, defaultPrecision, defaultScale)) { + final RecordSchema resultSetSchema = recordSet.getSchema(); + writeSchema = recordSetWriterFactory.getSchema(originalAttributes, resultSetSchema); + + try (final RecordSetWriter resultSetWriter = recordSetWriterFactory.createWriter(logger, writeSchema, out, originalAttributes)) { + writeResultRef.set(resultSetWriter.write(recordSet)); + mimeTypeRef.set(resultSetWriter.getMimeType()); + } catch (final Exception e) { + throw new IOException(e); Review Comment: Recommend adding a message along the lines of "Writing result records failed" -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
