generate static getClassSchema() method for record classes under avro…

… < 1.7 (#43)

Co-authored-by: Radai Rosenblatt <[email protected]>

helper/helper-common/src/main/java/com/linkedin/avroutil1/compatibility/CodeTransformations.java

@@ -37,6 +37,10 @@ public class CodeTransformations {
   private static final String  ESCAPED_QUOTE = "\\\"";
   private static final String  BACKSLASH = "\\";
   private static final String  DOUBLE_BACKSLASH = "\\\\";
+  private static final Pattern GET_CLASS_SCHEMA_PATTERN = Pattern.compile(Pattern.quote("public static org.apache.avro.Schema getClassSchema()"));
+  private static final Pattern SCHEMA_DOLLAR_DECL_PATTERN = Pattern.compile(Pattern.quote("public static final org.apache.avro.Schema SCHEMA$"));
+  private static final Pattern END_OF_SCHEMA_DOLLAR_DECL_PATTERN = Pattern.compile("}\"\\)(\\.toString\\(\\)\\))?;");
+  private static final String  GET_CLASS_SCHEMA_METHOD = "public static org.apache.avro.Schema getClassSchema() { return SCHEMA$; }";
   private static final Pattern NEW_BUILDER_METHOD_PATTERN = Pattern.compile("public static ([\\w.]+) newBuilder\\(\\)");
   private static final Pattern END_BUILDER_CLASS_PATTERN = Pattern.compile("}\\s+}\\s+}");
   private static final Pattern FROM_BYTEBUFFER_METHOD_END_PATTERN = Pattern.compile("return DECODER.decode\\s*\\(\\s*b\\s*\\)\\s*;\\s*}");
@@ -375,6 +379,50 @@ static String escapeJavaLiteral(String str) {
     return result.toString();
   }
 
+  /**
+   * avro 1.7+ adds a static getClassSchema() method that returns SCHEMA$
+   * @param code avro generated source code which may lack getClassSchema()
+   * @param generatedWith version of avro that generated the original code
+   * @param minSupportedVersion lowest avro version under which the generated code should work
+   * @param maxSupportedVersion highest avro version under which the generated code should work
+   * @return source code that contains getClassSchema()
+   */
+  public static String addGetClassSchemaMethod(
+      String code,
+      AvroVersion generatedWith,
+      AvroVersion minSupportedVersion,
+      AvroVersion maxSupportedVersion
+  ) {
+    Matcher classMatcher = GET_CLASS_SCHEMA_PATTERN.matcher(code);
+    if (classMatcher.find()) {
+      //this code already has the method
+      return code;
+    }
+    int schema$EndPosition = findEndOfSchemaDeclaration(code);
+    return code.substring(0, schema$EndPosition) + "\n  " + GET_CLASS_SCHEMA_METHOD + "\n" + code.substring(schema$EndPosition);
+  }
+
+  /**
+   * find the end of the SCHEMA$ declaration, for either simple declarations or those where the schema is huge and
+   * we had to use a StringBuilder
+   * @param code avro class code
+   * @return location of the ending of the SCHEMA$ declaration i the given code
+   */
+  static int findEndOfSchemaDeclaration(String code) {
+    Matcher schema$Matcher = SCHEMA_DOLLAR_DECL_PATTERN.matcher(code);
+    if (!schema$Matcher.find()) {
+      throw new IllegalStateException("unable to locate SCHEMA$ in " + code);
+    }
+    int schema$Position = schema$Matcher.end();
+    //locate the end of the SCHEMA$ declaration (allow to StringBuilder().toString() code if the schema
+    //was huge and we split is in transformParseCalls())
+    Matcher schema$EndMatcher = END_OF_SCHEMA_DOLLAR_DECL_PATTERN.matcher(code);
+    if (!schema$EndMatcher.find(schema$Position)) {
+      throw new IllegalStateException("unable to locate SCHEMA$ ending in " + code);
+    }
+    return schema$EndMatcher.end();
+  }
+
   /**
    * avro 1.6+ adds builder support to all generated record classes. the problem is that these builders all extend
    * org.apache.avro.specific.SpecificRecordBuilderBase, which doeswnt exist before avro 1.6

helper/helper-common/src/test/java/com/linkedin/avroutil1/compatibility/CodeTransformationsTest.java

@@ -61,4 +61,13 @@ public void testSafeSplit() {
         Arrays.asList("xxxxxxxxxx","x\\u1234xxx"),
         CodeTransformations.safeSplit("xxxxxxxxxxx\\u1234xxx", 10));
   }
+
+  @Test
+  public void testFindEndOfSchemaDeclaration() {
+    String normal = "public static final org.apache.avro.Schema SCHEMA$ = whatever(\"{json}\"); fluff";
+    String huge = "public static final org.apache.avro.Schema SCHEMA$ = whatever(new StringBuilder().append(\"{json}\").append(\"{more json}\").toString());";
+
+    Assert.assertEquals(CodeTransformations.findEndOfSchemaDeclaration(normal), normal.length() - 6);
+    Assert.assertEquals(CodeTransformations.findEndOfSchemaDeclaration(huge), huge.length());
+  }
 }

helper/impls/helper-impl-14/src/main/java/com/linkedin/avroutil1/compatibility/avro14/Avro14Adapter.java

@@ -212,6 +212,7 @@ private Collection<AvroGeneratedSourceCode> transform(List<AvroGeneratedSourceCo
       fixed = CodeTransformations.transformFixedClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformEnumClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformParseCalls(fixed, AvroVersion.AVRO_1_4, minAvro, maxAvro);
+      fixed = CodeTransformations.addGetClassSchemaMethod(fixed, AvroVersion.AVRO_1_4, minAvro, maxAvro);
       fixed = CodeTransformations.removeBuilderSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeBinaryMessageCodecSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeAvroGeneratedAnnotation(fixed, minAvro, maxAvro);

helper/impls/helper-impl-15/src/main/java/com/linkedin/avroutil1/compatibility/avro15/Avro15Adapter.java

@@ -237,6 +237,7 @@ private Collection<AvroGeneratedSourceCode> transform(List<AvroGeneratedSourceCo
       fixed = CodeTransformations.transformFixedClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformEnumClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformParseCalls(fixed, AvroVersion.AVRO_1_5, minAvro, maxAvro);
+      fixed = CodeTransformations.addGetClassSchemaMethod(fixed, AvroVersion.AVRO_1_5, minAvro, maxAvro);
       fixed = CodeTransformations.removeBuilderSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeBinaryMessageCodecSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeAvroGeneratedAnnotation(fixed, minAvro, maxAvro);

helper/impls/helper-impl-16/src/main/java/com/linkedin/avroutil1/compatibility/avro16/Avro16Adapter.java

@@ -246,6 +246,7 @@ private Collection<AvroGeneratedSourceCode> transform(List<AvroGeneratedSourceCo
       fixed = CodeTransformations.transformFixedClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformEnumClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformParseCalls(fixed, AvroVersion.AVRO_1_6, minAvro, maxAvro);
+      fixed = CodeTransformations.addGetClassSchemaMethod(fixed, AvroVersion.AVRO_1_6, minAvro, maxAvro);
       fixed = CodeTransformations.removeBuilderSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeBinaryMessageCodecSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeAvroGeneratedAnnotation(fixed, minAvro, maxAvro);

helper/impls/helper-impl-17/src/main/java/com/linkedin/avroutil1/compatibility/avro17/Avro17Adapter.java

@@ -307,6 +307,7 @@ private Collection<AvroGeneratedSourceCode> transform(List<AvroGeneratedSourceCo
       fixed = CodeTransformations.transformFixedClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformEnumClass(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.transformParseCalls(fixed, AvroVersion.AVRO_1_7, minAvro, maxAvro);
+      fixed = CodeTransformations.addGetClassSchemaMethod(fixed, AvroVersion.AVRO_1_7, minAvro, maxAvro);
       fixed = CodeTransformations.removeBuilderSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeBinaryMessageCodecSupport(fixed, minAvro, maxAvro);
       fixed = CodeTransformations.removeAvroGeneratedAnnotation(fixed, minAvro, maxAvro);

helper/tests/helper-tests-allavro/src/test/java/com/linkedin/avroutil1/compatibility/AvroCompatibilityHelperGeneratedRecordClassesTest.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2020 LinkedIn Corp.
+ * Licensed under the BSD 2-Clause License (the "License").
+ * See License in the project root for license information.
+ */
+
+package com.linkedin.avroutil1.compatibility;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * tests record classes generated by {@link AvroCompatibilityHelper} under different avro versions
+ */
+public class AvroCompatibilityHelperGeneratedRecordClassesTest {
+
+  @Test
+  public void testRecordGeneratedUnderAvro14HasGetClassSchema() throws Exception {
+    Assert.assertEquals(under14.SimpleRecord.getClassSchema(), under14.SimpleRecord.SCHEMA$);
+  }
+
+  @Test
+  public void testRecordGeneratedUnderAvro15HasGetClassSchema() throws Exception {
+    Assert.assertEquals(under15.SimpleRecord.getClassSchema(), under15.SimpleRecord.SCHEMA$);
+  }
+
+  @Test
+  public void testRecordGeneratedUnderAvro16HasGetClassSchema() throws Exception {
+    Assert.assertEquals(under16.SimpleRecord.getClassSchema(), under16.SimpleRecord.SCHEMA$);
+  }
+
+  @Test
+  public void testRecordGeneratedUnderAvro17HasGetClassSchema() throws Exception {
+    Assert.assertEquals(under17.SimpleRecord.getClassSchema(), under17.SimpleRecord.SCHEMA$);
+  }
+}