diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/CobolProcessor.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/CobolProcessor.scala index 216097be..6d7a71e0 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/CobolProcessor.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/CobolProcessor.scala @@ -38,10 +38,11 @@ trait CobolProcessor { * @param inputStream the input stream containing raw COBOL records. * @param outputStream the output stream where processed records will be written. * @param rawRecordProcessor the processor that processes each raw record. + * @return The number of records processed. */ def process(inputStream: SimpleStream, outputStream: OutputStream) - (rawRecordProcessor: RawRecordProcessor): Unit + (rawRecordProcessor: RawRecordProcessor): Long } @@ -56,7 +57,7 @@ object CobolProcessor { new CobolProcessor { override def process(inputStream: SimpleStream, outputStream: OutputStream) - (rawRecordProcessor: RawRecordProcessor): Unit = { + (rawRecordProcessor: RawRecordProcessor): Long = { val recordExtractor = getRecordExtractor(readerParameters, inputStream) val dataStream = inputStream.copyStream() @@ -116,7 +117,7 @@ object CobolProcessor { reader.recordExtractor(0, dataStream, headerStream) match { case Some(extractor) => extractor - case None => + case None => throw new IllegalArgumentException(s"Cannot create a record extractor for the given reader parameters. " + "Please check the copybook and the reader parameters." ) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/impl/StreamProcessor.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/impl/StreamProcessor.scala index e9d84347..10d0e170 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/impl/StreamProcessor.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/impl/StreamProcessor.scala @@ -34,14 +34,17 @@ object StreamProcessor { * @param recordExtractor the extractor that extracts raw records from the input stream. * @param recordProcessor the per-record processing logic implementation. * @param outputStream the output stream where the processed records will be written. + * @return The number of records processed. */ def processStream(copybook: Copybook, options: Map[String, String], inputStream: SimpleStream, recordExtractor: RawRecordExtractor, recordProcessor: RawRecordProcessor, - outputStream: OutputStream): Unit = { + outputStream: OutputStream): Long = { + var recordCount = 0L while (recordExtractor.hasNext) { + recordCount += 1 val record = recordExtractor.next() val recordSize = record.length @@ -61,5 +64,6 @@ object StreamProcessor { val footer = inputStream.next(footerSize.toInt) outputStream.write(footer) } + recordCount } } diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/processor/CobolProcessorBuilderSuite.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/processor/CobolProcessorBuilderSuite.scala index b25e237b..7ec9d08b 100644 --- a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/processor/CobolProcessorBuilderSuite.scala +++ b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/processor/CobolProcessorBuilderSuite.scala @@ -42,10 +42,11 @@ class CobolProcessorBuilderSuite extends AnyWordSpec { } } - builder.build().process(is, os)(processor) + val count = builder.build().process(is, os)(processor) val outputArray = os.toByteArray + assert(count == 4) assert(outputArray.head == -16) assert(outputArray(1) == -15) assert(outputArray(2) == -14)