#332 Add get dependencies and build proto.

sngular · Feb 23, 2023 · 9a197c7 · 9a197c7
1 parent b21e311
commit 9a197c7
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/pom-maven-central.xml b/pom-maven-central.xml
@@ -7,7 +7,7 @@
 
   <artifactId>kloadgen</artifactId>
 
-  <version>5.3.0</version>
+  <version>5.3.1</version>
 
   <name>KLoadGen</name>
   <description>Load Generation Jmeter plugin for Kafka Cluster. Supporting AVRO, JSON Schema and Protobuf schema types. Generate Artificial data based on Data specification</description>

diff --git a/pom.xml b/pom.xml
@@ -7,7 +7,7 @@
 
   <artifactId>kloadgen</artifactId>
 
-  <version>5.3.0</version>
+  <version>5.3.1</version>
 
   <name>KLoadGen</name>
   <description>Load Generation Jmeter plugin for Kafka Cluster. Supporting AVRO, JSON Schema and Protobuf schema types. Generate Artificial data base on Data specification</description>

diff --git a/src/main/java/com/sngular/kloadgen/processor/util/SchemaProcessorUtils.java b/src/main/java/com/sngular/kloadgen/processor/util/SchemaProcessorUtils.java
@@ -25,16 +25,20 @@
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.sngular.kloadgen.model.FieldValueMapping;
+import com.sngular.kloadgen.util.JMeterHelper;
 import com.sngular.kloadgen.util.ProtobufHelper;
 import com.squareup.wire.schema.internal.parser.EnumElement;
 import com.squareup.wire.schema.internal.parser.FieldElement;
 import com.squareup.wire.schema.internal.parser.MessageElement;
 import com.squareup.wire.schema.internal.parser.ProtoFileElement;
 import com.squareup.wire.schema.internal.parser.TypeElement;
+import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException;
+import io.confluent.kafka.schemaregistry.protobuf.ProtobufSchema;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.collections4.Predicate;
 import org.apache.commons.lang3.RandomUtils;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.jmeter.threads.JMeterContextService;
 
 public class SchemaProcessorUtils {
 
@@ -156,7 +160,15 @@ public static Descriptors.Descriptor buildProtoDescriptor(final ProtoFileElement
             schemaBuilder.addDependency(importedSchema.getFileDescriptorSet().getFile(0).getName());
             schemaBuilder.addSchema(importedSchema);
           }
+        } else {
+          final var importSchema = JMeterHelper.getParsedSchema(importedClass.substring(0, importedClass.length() - 6), JMeterContextService.getContext().getProperties());
+          if (!ProtobufHelper.NOT_ACCEPTED_IMPORTS.contains(importedClass)) {
+            schemaBuilder.addDependency(((ProtobufSchema) importSchema).toDescriptor().getFullName());
+            schemaBuilder.addSchema(convertDynamicSchema((ProtobufSchema) importSchema));
+          }
         }
+      } catch (RestClientException e) {
+        throw new RuntimeException(e);
       }
     }
     final MessageElement messageElement = (MessageElement) schema.getTypes().get(0);
@@ -170,6 +182,10 @@ public static Descriptors.Descriptor buildProtoDescriptor(final ProtoFileElement
     return schemaBuilder.build().getMessageDescriptor(messageElement.getName());
   }
 
+  private static DynamicSchema convertDynamicSchema(final ProtobufSchema importSchema) throws DescriptorValidationException {
+    return processImported(Arrays.asList(importSchema.rawSchema().toSchema().split("\\n")));
+  }
+
   private static Predicate<String> isValid() {
     return line -> !line.contains("//") && !line.isEmpty();
   }
@@ -208,9 +224,11 @@ private static MessageDefinition buildMessage(final String messageName, final Li
     while (messageLines.hasNext() && !exit) {
       final var field = messageLines.next().trim().split("\\s");
       if (ProtobufHelper.isValidType(field[0])) {
-        messageDefinition.addField(OPTIONAL, field[0], field[1], Integer.parseInt(checkIfChoppable(field[3])));
+        messageDefinition.addField(OPTIONAL, field[0], field[1], Integer.parseInt(checkIfGreppable(field[3])));
       } else if (ProtobufHelper.LABEL.contains(field[0])) {
-        messageDefinition.addField(field[0], field[1], field[2], Integer.parseInt(checkIfChoppable(field[4])));
+        messageDefinition.addField(field[0], field[1], field[2], Integer.parseInt(checkIfGreppable(field[4])));
+      } else if ("message".equalsIgnoreCase(field[0])) {
+        messageDefinition.addMessageDefinition(buildMessage(field[1], messageLines));
       } else if ("}".equalsIgnoreCase(field[0])) {
         exit = true;
       }
@@ -219,7 +237,7 @@ private static MessageDefinition buildMessage(final String messageName, final Li
     return messageDefinition.build();
   }
 
-  private static String checkIfChoppable(final String field) {
+  private static String checkIfGreppable(final String field) {
     String choppedField = field;
     if (field.endsWith(";")) {
       choppedField = StringUtils.chop(field);