diff --git a/README.md b/README.md index 0ff23f7e..d07af1d2 100644 --- a/README.md +++ b/README.md @@ -1635,6 +1635,49 @@ The output looks like this: +---+-------+----------+----------+----------+ ``` +### Custom AST transformers +When parsed copybooks need to be transformed in a regular way before processing data, a custom AST transformer can be +used. This allows modifying data types on the fly, or adding/removing fields when necessary. + +In order to use a custom AST transformer, you need to implement the `AstTransformer` trait and provide your +transformation logic in the `transform` method. Then, you can specify your custom transformer class in the `ast_transformers` +option when reading data. + +Example transformer: + +```scala +package com.example + +import za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer +import za.co.absa.cobrix.cobol.reader.parameters.ReaderParameters +import za.co.absa.cobrix.cobol.parser.CopybookParser.CopybookAST + +class MyTransformer1(readerParameters: ReaderParameters) extends AstTransformer { + override def transform(ast: CopybookAST): CopybookAST = { + // implementation goes here + } +} +``` + +Example options that uses the transformer: +```scala +spark.read + .format("cobol") + .option("copybook", someCopybook) + .option("ast_transformers", "com.example.MyTransformer1") + .load("/some/path") +``` + +You can chain multiple transformers by providing transformer classes as a comma-separated list. + +**Notes:** + +1. Transformers are applied before built-in transformers, so the AST might not be complete at that stage. +2. Since the decoder and optionally an encoder are set at the AST building stage, if you want to insert new fields, + you need to define them as well. Use `DecoderSelector.getDecoder()` and [optionally] `EncoderSelector.getEncoder()`. +3. Custom transformers run before binary properties are calculated. So inserting or removing fields changes the copybook + layout. + ## Summary of all available options ##### File reading options @@ -1688,6 +1731,7 @@ The output looks like this: | .option("record_limit", "1000") | Caps decoded output to N rows globally across all input paths/files (zero returns no rows). This is a source-level global cap: `df.limit(N)` does not push down into Cobrix. While enabled Cobrix uses a single scan task and can avoid opening later input partitions/files once satisfied. Indexed reads may still build indexes first. | | .option("with_input_file_name_col", "file_name") | Generates a column containing input file name for each record (Similar to Spark SQL `input_file_name()` function). The column name is specified by the value of the option. This option only works for variable record length files. For fixed record length and ASCII files use `input_file_name()`. | | .option("metadata", "basic") | Specifies wat kind of metadata to include in the Spark schema: `false`, `basic`(default), or `extended` (PIC, usage, etc). | +| .option("ast_transformers", "...") | Specifies a comma-separated list of custom AST transformers to be used to process AST of copybooks just after they are parsed. | | .option("debug", "hex") | If specified, each primitive field will be accompanied by a debug field containing raw bytes from the source file. Possible values: `none` (default), `hex`, `binary`, `string` (ASCII only). The legacy value `true` is supported and will generate debug fields in HEX. | For example, to cap a Cobrix source read at 1,000 decoded rows: @@ -1815,7 +1859,7 @@ val df = spark `common_extended`, `cp037_extended` are code pages supporting non-printable characters that converts to ASCII codes below 32. -## EBCDIC Processor (experimental) +## EBCDIC Processor The EBCDIC processor allows processing files by replacing value of fields without changing the underlying format (`CobolProcessingStrategy.InPlace`) or with conversion of the input format to variable-record-length format with big-endian RDWs (`CobolProcessingStrategy.ToVariableLength`). @@ -1871,7 +1915,7 @@ val count = CobolProcessor.builder ``` -## EBCDIC Spark Processor (experimental) +## EBCDIC Spark Processor This allows in-place processing of data retaining original format in parallel uring RDDs under the hood. Here is an example usage: @@ -1898,7 +1942,7 @@ SparkCobolProcessor.builder .save(outputPath) ``` -## EBCDIC Spark raw record RDD generator (experimental) +## EBCDIC Spark raw record RDD generator You can process raw records of a mainframe file as an `RDD[Array[Byte]]`. This can be useful for custom processing without converting to Spark data types. You can still access fields via parsed copybooks. @@ -1927,7 +1971,7 @@ val segmentRdds RDD[String] = recordsRdd.flatMap { record => segmentRdds.distinct.collect.sorted.foreach(println) ``` -## EBCDIC Writer (experimental) +## EBCDIC Writer Cobrix's EBCDIC writer is an experimental feature that allows writing Spark DataFrames as EBCDIC mainframe files. diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/parser/CopybookParser.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/parser/CopybookParser.scala index 4b29c88d..9a40af56 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/parser/CopybookParser.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/parser/CopybookParser.scala @@ -127,6 +127,7 @@ object CopybookParser extends Logging { * @param nonTerminals A list of non-terminals that should be extracted as strings. * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records. * @param debugFieldsPolicy Specifies if debugging fields need to be added and what should they contain (false, hex, raw). + * @param customAstTransformers A list of custom AST transformers to be applied during parsing. * @return Seq[Group] where a group is a record inside the copybook. */ def parse(copyBookContents: String, @@ -151,7 +152,8 @@ object CopybookParser extends Logging { occursHandlers: Map[String, Map[String, Int]] = Map(), redefineRuleExpressions: Map[String, ExpressionEvaluator] = Map.empty, debugFieldsPolicy: DebugFieldsPolicy = DebugFieldsPolicy.NoDebug, - fieldCodePageMap: Map[String, String] = Map.empty[String, String]): Copybook = { + fieldCodePageMap: Map[String, String] = Map.empty[String, String], + customAstTransformers: Seq[AstTransformer] = Seq.empty): Copybook = { parseTree(dataEncoding, copyBookContents, dropGroupFillers, @@ -174,7 +176,8 @@ object CopybookParser extends Logging { occursHandlers, redefineRuleExpressions, debugFieldsPolicy, - fieldCodePageMap) + fieldCodePageMap, + customAstTransformers) } /** @@ -198,6 +201,7 @@ object CopybookParser extends Logging { * @param nonTerminals A list of non-terminals that should be extracted as strings * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records. * @param debugFieldsPolicy Specifies if debugging fields need to be added and what should they contain (false, hex, raw). + * @param customAstTransformers A list of custom AST transformers to be applied during parsing. * @return Seq[Group] where a group is a record inside the copybook */ def parseTree(copyBookContents: String, @@ -221,7 +225,8 @@ object CopybookParser extends Logging { occursHandlers: Map[String, Map[String, Int]] = Map(), redefineRuleExpressions: Map[String, ExpressionEvaluator] = Map.empty, debugFieldsPolicy: DebugFieldsPolicy = DebugFieldsPolicy.NoDebug, - fieldCodePageMap: Map[String, String] = Map.empty[String, String]): Copybook = { + fieldCodePageMap: Map[String, String] = Map.empty[String, String], + customAstTransformers: Seq[AstTransformer] = Seq.empty): Copybook = { parseTree(EBCDIC, copyBookContents, dropGroupFillers, @@ -244,7 +249,8 @@ object CopybookParser extends Logging { occursHandlers, redefineRuleExpressions, debugFieldsPolicy, - fieldCodePageMap) + fieldCodePageMap, + customAstTransformers) } /** @@ -268,6 +274,7 @@ object CopybookParser extends Logging { * @param nonTerminals A list of non-terminals that should be extracted as strings * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records. * @param debugFieldsPolicy Specifies if debugging fields need to be added and what should they contain (false, hex, raw). + * @param customAstTransformers A list of custom AST transformers to be applied during parsing. * @return Seq[Group] where a group is a record inside the copybook */ @throws(classOf[SyntaxErrorException]) @@ -293,7 +300,8 @@ object CopybookParser extends Logging { occursHandlers: Map[String, Map[String, Int]], redefineRuleExpressions: Map[String, ExpressionEvaluator], debugFieldsPolicy: DebugFieldsPolicy, - fieldCodePageMap: Map[String, String]): Copybook = { + fieldCodePageMap: Map[String, String], + customAstTransformers: Seq[AstTransformer]): Copybook = { val schemaANTLR: CopybookAST = ANTLRParser.parse(copyBookContents, enc, stringTrimmingPolicy, isDisplayAlwaysString, commentPolicy, strictSignOverpunch, improvedNullDetection, strictIntegralPrecision, decodeBinaryAsHex, ebcdicCodePage, asciiCharset, isUtf16BigEndian, floatingPointFormat, fieldCodePageMap) @@ -304,7 +312,7 @@ object CopybookParser extends Logging { val correctedFieldParentMap = transformIdentifierMap(fieldParentMap) validateFieldParentMap(correctedFieldParentMap) - val transformers = Seq( + val transformers = customAstTransformers ++ Seq( // Calculate sizes of fields and their positions from the beginning of a record BinaryPropertiesAdder(), // Adds virtual primitive fields for GROUPs that can be parsed as concatenation of their children. diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala index 012639d4..ce73e107 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala @@ -69,6 +69,7 @@ import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaReten * @param writerParameters Parameters set only for the writer of files in mainfraame format * @param options Options passed to 'spark-cobol'. * @param recordLimit Maximum number of decoded output rows to return. Zero returns no rows; if not specified, all rows are returned. + * @param astTransformerClasses A comma-separated list of fully qualified classes extending [[za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer]] */ case class CobolParameters( copybookPath: Option[String], @@ -118,6 +119,7 @@ case class CobolParameters( fileHeaderField: Option[String] = None, fileTrailerField: Option[String] = None, recordLimit: Option[Int] = None, + astTransformerClasses: Seq[String], writerParameters: Option[WriterParameters], options: Map[String, String] ) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala index 3e4919b6..d48332a4 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala @@ -80,6 +80,7 @@ object CobolParametersParser extends Logging { val PARAM_STRICT_INTEGRAL_PRECISION = "strict_integral_precision" val PARAM_DISPLAY_PIC_ALWAYS_STRING = "display_pic_always_string" val PARAM_STRICT_SCHEMA = "strict_schema" + val PARAM_AST_TRANSFORMERS = "ast_transformers" val PARAM_GROUP_NOT_TERMINALS = "non_terminals" val PARAM_OCCURS_MAPPINGS = "occurs_mappings" @@ -309,6 +310,9 @@ object CobolParametersParser extends Logging { None } + val astTransformerClasses = params.get(PARAM_AST_TRANSFORMERS) + .map(_.split(',').map(_.trim).filter(_.nonEmpty).toSeq).getOrElse(Seq.empty) + val cobolParameters = CobolParameters( copybookPathOpt, copybookPaths, @@ -357,6 +361,7 @@ object CobolParametersParser extends Logging { params.get(PARAM_RECORD_HEADER_NAME).orElse(params.get(PARAM_RECORD_HEADER_NAME2)) .map(_.trim).filter(_.nonEmpty), params.get(PARAM_RECORD_TRAILER_NAME).orElse(params.get(PARAM_RECORD_TRAILER_NAME2)) .map(_.trim).filter(_.nonEmpty), recordLimit, + astTransformerClasses, writerParameters, params.getMap ) @@ -540,6 +545,7 @@ object CobolParametersParser extends Logging { parameters.metadataPolicy, recordsToExclude, parameters.recordLimit, + parameters.astTransformerClasses, parameters.writerParameters, parameters.options ) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala index 75b39056..402aa7a6 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala @@ -82,6 +82,7 @@ import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaReten * @param metadataPolicy Specifies the policy of metadat fields to be added to the Spark schema * @param recordsToExclude A set of fields to exclude from decoding in normal (not header or footer) records. * @param recordLimit Maximum number of decoded output rows to return. Zero returns no rows; if not specified, all rows are returned. + * @param astTransformerClasses A comma-separated list of fully qualified classes extending [[za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer]] * @param writerParameters If specified, contains parameters for the writer. * @param options Options passed to spark-cobol */ @@ -146,6 +147,7 @@ case class ReaderParameters( metadataPolicy: MetadataPolicy = MetadataPolicy.Basic, recordsToExclude: Set[String] = Set.empty, recordLimit: Option[Int] = None, + astTransformerClasses: Seq[String] = Seq.empty, writerParameters: Option[WriterParameters] = None, options: Map[String, String] = Map.empty ) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/schema/CobolSchema.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/schema/CobolSchema.scala index e29381e9..0359ad70 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/schema/CobolSchema.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/schema/CobolSchema.scala @@ -22,6 +22,7 @@ import za.co.absa.cobrix.cobol.parser.policies.MetadataPolicy import za.co.absa.cobrix.cobol.parser.{Copybook, CopybookParser} import za.co.absa.cobrix.cobol.reader.parameters.{CorruptFieldsPolicy, ReaderParameters} import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaRetentionPolicy +import za.co.absa.cobrix.cobol.utils.AstTransformerUtils import java.nio.charset.{Charset, StandardCharsets} import java.time.ZonedDateTime @@ -98,6 +99,8 @@ object CobolSchema { case None => StandardCharsets.UTF_8 } + val customAstTransformers = readerParameters.astTransformerClasses.map(className => AstTransformerUtils.loadAstTransformer(className, readerParameters)) + val schema = if (copyBookContents.size == 1) CopybookParser.parseTree(encoding, copyBookContents.head, @@ -121,7 +124,8 @@ object CobolSchema { readerParameters.occursMappings, readerParameters.redefineRuleExpressions, readerParameters.debugFieldsPolicy, - readerParameters.fieldCodePage) + readerParameters.fieldCodePage, + customAstTransformers) else Copybook.merge(copyBookContents.map(cpb => CopybookParser.parseTree(encoding, @@ -146,7 +150,8 @@ object CobolSchema { readerParameters.occursMappings, readerParameters.redefineRuleExpressions, readerParameters.debugFieldsPolicy, - readerParameters.fieldCodePage) + readerParameters.fieldCodePage, + customAstTransformers) )) val segIdFieldCount = readerParameters.multisegment.map(p => p.segmentLevelIds.size).getOrElse(0) val segmentIdPrefix = readerParameters.multisegment.map(p => p.segmentIdPrefix).getOrElse("") diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/AstTransformerUtils.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/AstTransformerUtils.scala new file mode 100644 index 00000000..bc780840 --- /dev/null +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/AstTransformerUtils.scala @@ -0,0 +1,55 @@ +/* + * Copyright 2018 ABSA Group Limited + * + * Licensed 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 za.co.absa.cobrix.cobol.utils + +import za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer +import za.co.absa.cobrix.cobol.reader.parameters.ReaderParameters + +import java.lang.reflect.Constructor +import scala.util.Try + +object AstTransformerUtils { + /** + * Loads and instantiates an AstTransformer by its fully qualified class name using reflection. + * + * The method attempts to find and invoke a constructor that accepts a ReaderParameters argument. + * If no such constructor exists, it falls back to using the default no-argument constructor. + * + * @param fullyQualifiedName the fully qualified class name of the AstTransformer implementation to load + * @param readerParameters the reader parameters to pass to the constructor if a matching constructor is available + * @return an instance of the specified AstTransformer class + * @throws java.lang.ClassNotFoundException if the class with the given fully qualified name cannot be found + * @throws java.lang.ClassCastException if the instantiated object cannot be cast to AstTransformer + * @throws java.lang.NoSuchMethodException if neither a constructor accepting ReaderParameters nor a default constructor is found + */ + @throws[ClassNotFoundException] + @throws[ClassCastException] + @throws[NoSuchMethodException] + def loadAstTransformer(fullyQualifiedName: String, readerParameters: ReaderParameters): AstTransformer = { + // There are 2 types of constructors supported. If there is a one that takes a ReaderParameters object - use it. + // Otherwise, try the default constructor. + val confCtor = Try[Constructor[_]](Class.forName(fullyQualifiedName).getConstructor(classOf[ReaderParameters])) + + confCtor + .map(ctor => ctor.newInstance(readerParameters).asInstanceOf[AstTransformer]) + .recoverWith { + case _: NoSuchMethodException => + val defCtor = Try[Constructor[_]](Class.forName(fullyQualifiedName).getConstructor()) + defCtor.map(ctor => ctor.newInstance().asInstanceOf[AstTransformer]) + }.get + } +} diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/AstTransformerSpy.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/AstTransformerSpy.scala new file mode 100644 index 00000000..7c0a3fda --- /dev/null +++ b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/AstTransformerSpy.scala @@ -0,0 +1,40 @@ +/* + * Copyright 2018 ABSA Group Limited + * + * Licensed 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 za.co.absa.cobrix.cobol.parser.asttransform + +import za.co.absa.cobrix.cobol.parser.CopybookParser.CopybookAST +import za.co.absa.cobrix.cobol.parser.ast.{Group, Primitive, Statement} + +import scala.collection.mutable.ArrayBuffer + +class AstTransformerSpy extends AstTransformer { + @volatile var transformCaller = 0 + + override def transform(ast: CopybookAST): CopybookAST = { + transformCaller += 1 + val fields = ast.children.head.asInstanceOf[Group].children + val newChildren: ArrayBuffer[Statement] = fields.map { + case p: Primitive => p.copy(name = p.name + "_transformed")(p.parent) + case g: Group => g.copy(name = g.name + "_transformed")(g.parent) + } + + val originalGroup = ast.children.head.asInstanceOf[Group] + val newGroup = originalGroup.copy(children = newChildren)(originalGroup.parent) + ast.children.update(0, newGroup) + ast + } +} diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/ParsingWithAstTransformersSuite.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/ParsingWithAstTransformersSuite.scala new file mode 100644 index 00000000..df68470c --- /dev/null +++ b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/parser/asttransform/ParsingWithAstTransformersSuite.scala @@ -0,0 +1,46 @@ +/* + * Copyright 2018 ABSA Group Limited + * + * Licensed 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 za.co.absa.cobrix.cobol.parser.asttransform + +import org.scalatest.wordspec.AnyWordSpec +import za.co.absa.cobrix.cobol.parser.CopybookParser +import za.co.absa.cobrix.cobol.parser.ast.Group +import za.co.absa.cobrix.cobol.parser.exceptions.RuleExpressionParsingException +import za.co.absa.cobrix.cobol.parser.expression.ExpressionEvaluator + +class ParsingWithAstTransformersSuite extends AnyWordSpec { + "parse" should { + "parse the copybook with a custom transformer" in { + val copybook = + """ 01 RECORD. + | 05 FIELD-A PIC 9(5). + | 05 FIELD-B PIC 9(5). + | 05 FIELD-C PIC X(10). + |""".stripMargin + + val transformers = Seq(new AstTransformerSpy) + val schema = CopybookParser.parseTree(copybook, customAstTransformers = transformers) + + val fields = schema.ast.children.head.asInstanceOf[Group].children + assert(fields.head.name == "FIELD_A_transformed") + assert(fields(1).name == "FIELD_B_transformed") + assert(fields(2).name == "FIELD_C_transformed") + + assert(transformers.head.transformCaller == 1) + } + } +} diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/mocks/AstTransformerSpy.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/mocks/AstTransformerSpy.scala new file mode 100644 index 00000000..dfe7a611 --- /dev/null +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/mocks/AstTransformerSpy.scala @@ -0,0 +1,54 @@ +/* + * Copyright 2018 ABSA Group Limited + * + * Licensed 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 za.co.absa.cobrix.spark.cobol.mocks + +import za.co.absa.cobrix.cobol.parser.CopybookParser.CopybookAST +import za.co.absa.cobrix.cobol.parser.ast.datatype.AlphaNumeric +import za.co.absa.cobrix.cobol.parser.ast.{Group, Primitive, Statement} +import za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer +import za.co.absa.cobrix.cobol.parser.decoders.DecoderSelector +import za.co.absa.cobrix.cobol.parser.decoders.DecoderSelector.getStringDecoder +import za.co.absa.cobrix.cobol.parser.encoding.EBCDIC + +import scala.collection.mutable.ArrayBuffer + +class AstTransformerSpy extends AstTransformer { + import AstTransformerSpy._ + + override def transform(ast: CopybookAST): CopybookAST = { + transformCaller += 1 + val root = ast.children.head.asInstanceOf[Group] + val fields = root.children + val newChildren: ArrayBuffer[Statement] = fields.map { + case p: Primitive => p.copy(name = p.name + "_transformed")(p.parent) + case g: Group => g.copy(name = g.name + "_transformed")(g.parent) + } + + val dataType = AlphaNumeric("X(1)", 1, None, None, None, None) + val stringDecoder = DecoderSelector.getDecoder(dataType) + val extraField = Primitive(3, "EXTRA", "EXTRA", 0, dataType, decode = stringDecoder, encode = None)(Some(root)) + newChildren.append(extraField) + + val newGroup = root.copy(children = newChildren)(root.parent) + ast.children.update(0, newGroup) + ast + } +} + +object AstTransformerSpy { + @volatile var transformCaller = 0 +} \ No newline at end of file diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test44CustomAstTransformersSpec.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test44CustomAstTransformersSpec.scala new file mode 100644 index 00000000..bc09e897 --- /dev/null +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test44CustomAstTransformersSpec.scala @@ -0,0 +1,78 @@ +/* + * Copyright 2018 ABSA Group Limited + * + * Licensed 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 za.co.absa.cobrix.spark.cobol.source.integration + +import org.scalatest.wordspec.AnyWordSpec +import org.slf4j.{Logger, LoggerFactory} +import za.co.absa.cobrix.spark.cobol.mocks.AstTransformerSpy +import za.co.absa.cobrix.spark.cobol.source.base.{SimpleComparisonBase, SparkTestBase} +import za.co.absa.cobrix.spark.cobol.source.fixtures.BinaryFileFixture +import za.co.absa.cobrix.spark.cobol.utils.SparkUtils + +class Test44CustomAstTransformersSpec extends AnyWordSpec with SparkTestBase with BinaryFileFixture with SimpleComparisonBase { + private implicit val logger: Logger = LoggerFactory.getLogger(this.getClass) + + private val copybookContent = + """ 01 R. + | 03 A PIC X(1). + |""".stripMargin + + private val rdwData = Array( + 0x00, 0x02, 0x00, 0x00, 0xF0, 0xF1, + 0x00, 0x02, 0x00, 0x00, 0xF2, 0xF3, + 0x00, 0x02, 0x00, 0x00, 0xF4, 0xF5 + ).map(_.toByte) + + "ast_transformers" should { + "transform the ast before processing in Spark" in { + val expectedData = + """[ { + | "A_transformed" : "0", + | "EXTRA" : "1" + |}, { + | "A_transformed" : "2", + | "EXTRA" : "3" + |}, { + | "A_transformed" : "4", + | "EXTRA" : "5" + |} ]""".stripMargin + + withTempBinFile("record_limit_rdw", ".dat", rdwData) { path => + AstTransformerSpy.transformCaller = 0 + val df = spark + .read + .format("cobol") + .option("copybook_contents", copybookContent) + .option("record_format", "V") + .option("is_rdw_big_endian", "true") + .option("ast_transformers", "za.co.absa.cobrix.spark.cobol.mocks.AstTransformerSpy") + .load(path) + + assert(AstTransformerSpy.transformCaller == 1) + + val schema = df.schema + + assert(schema.fieldNames.contains("A_transformed")) + assert(schema.fieldNames.contains("EXTRA")) + + val actualData = SparkUtils.prettyJSON(df.toJSON.collect().mkString("[", ",", "]")) + + assertEqualsMultiline(actualData, expectedData) + } + } + } +}