Repository: ImFlog/schema-registry-plugin
Pull request #156 (2024-01-01 20:26:53)

Add ability to download schema references
Added ability to download schema references.

A new configuration `downloadReferences: Boolean` was added to `DownloadSchema`.

When set to true, and the schema has references, the referenced schemas are also downloaded. This will allow POJO generation to succeed from schema with references.

ImFlog: @ferozed the integration tests are failing. There is a special gradle task for them that you should be able to run locally (if you have Docker installed).

Also, after looking a bit in the confluent code I wonder if there are not an integrated way to do this directly when parsing the schema using [`toNormalize](https://github.com/confluentinc/schema-registry/blob/b3657d1817c0306b63b238ab3dd039889870f30a/client/src/main/java/io/confluent/kafka/schemaregistry/avro/AvroSchema.java#L314C1-L314C1) before retrieving the schema ?
ferozed: > Also, after looking a bit in the confluent code I wonder if there are not an integrated way to do this directly when parsing the schema using [`toNormalize](https://github.com/confluentinc/schema-registry/blob/b3657d1817c0306b63b238ab3dd039889870f30a/client/src/main/java/io/confluent/kafka/schemaregistry/avro/AvroSchema.java#L314C1-L314C1) before retrieving the schema ?

Right... but this will only work for AVRO schemas?

ferozed: Yeah... so the integration tests are failing because the new config I added is a required value. Is there any way to make it optional? I tried using `@Optional` annotation but that didnt work.
ferozed: Ok... i kinda worked around it by putting the `downloadReferences` configuration per subject, instead of it being global.
ImFlog: Starts to look good !
Only issue is that you have no way to set the flag to true as no property in the extension is mapped to the subject.
You should probably add an integration test with the flag set to true and verify that 2 files are created.
ferozed: > you have no way to set the flag to true as no property in the extension is mapped to the subject

Sorry... I didnt understand this. Are you saying that this functionality cannot be set by the user?
ImFlog: LGTM !
I think I will try to release this with #160.

I keep you posted when I found the time to do so. 

Thank you for the contribution 🙏 

Patch:
diff --git a/README.md b/README.md
index 3b84507..f02955f 100644
--- a/README.md
+++ b/README.md
@@ -91,10 +91,15 @@ schemaRegistry {
 Here is the list of all the signatures for the `subject` extension:
 
 * `subject(inputSubject: String, outputPath: String)`
+* `subject(inputSubject: String, outputPath: String, downloadReferences: Boolean)`
 * `subject(inputSubject: String, outputPath: String, outputFileName: String)`
+* `subject(inputSubject: String, outputPath: String, outputFileName: String, downloadReferences: Boolean)`
 * `subject(inputSubject: String, outputPath: String, version: Int)`
+* `subject(inputSubject: String, outputPath: String, version: Int, downloadReferences: Boolean)`
 * `subject(inputSubject: String, outputPath: String, version: Int, outputFileName: String)`
+* `subject(inputSubject: String, outputPath: String, version: Int, outputFileName: String, downloadReferences: Boolean)`
 * `subjectPattern(inputPattern: String, outputPath: String)`
+* `subjectPattern(inputPattern: String, outputPath: String, downloadReferences: Boolean)`
 
 You can configure the metadata extension in order to download the schemas metadata in json files.
 It will be saved in files named like the schema file but suffixed by `-metadata.json` in the outputPath you specify
diff --git a/src/integration/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskIT.kt b/src/integration/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskIT.kt
index 4c2c50c..3dfaf52 100644
--- a/src/integration/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskIT.kt
+++ b/src/integration/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskIT.kt
@@ -4,8 +4,10 @@ import com.github.imflog.schema.registry.toSchemaType
 import com.github.imflog.schema.registry.utils.KafkaTestContainersUtils
 import io.confluent.kafka.schemaregistry.ParsedSchema
 import io.confluent.kafka.schemaregistry.avro.AvroSchema
+import io.confluent.kafka.schemaregistry.client.rest.entities.SchemaReference
 import io.confluent.kafka.schemaregistry.json.JsonSchema
 import io.confluent.kafka.schemaregistry.protobuf.ProtobufSchema
+import org.apache.avro.Schema
 import org.assertj.core.api.Assertions
 import org.gradle.internal.impldep.org.junit.rules.TemporaryFolder
 import org.gradle.testkit.runner.BuildResult
@@ -329,6 +331,88 @@ class DownloadTaskIT : KafkaTestContainersUtils() {
         Assertions.assertThat(result?.task(":downloadSchemasTask")?.outcome).isEqualTo(TaskOutcome.SUCCESS)
     }
 
+    @Test
+    fun `Should download schema with references in specific directory`() {
+        // Given
+        val subjectName = "reference_test"
+        val recordName = "UserReference"
+        val subjectNameLib = "reference_test_lib"
+        val recordNameLib = "UserReferenceLib"
+
+        client.register(
+            subjectNameLib,
+            AvroSchema(
+                """{
+                    "type": "record",
+                    "name": "${recordNameLib}",
+                    "fields": [
+                        { "name": "name", "type": "string" }
+                    ]
+                }"""
+            )
+        )
+
+        client.register(
+            subjectName,
+            AvroSchema(
+                """{
+                    "type": "record",
+                    "name": "${recordName}",
+                    "fields": [
+                        { "name": "name", "type": "string" }
+                    ]
+                }""",
+                listOf(
+                    SchemaReference(recordNameLib, subjectNameLib, 1)
+                ),
+                mapOf(),
+                null
+            )
+        )
+
+        buildFile = folderRule.newFile("build.gradle")
+        buildFile.writeText(
+            """
+            plugins {
+                id 'java'
+                id 'com.github.imflog.kafka-schema-registry-gradle-plugin'
+            }
+
+            import com.github.imflog.schema.registry.tasks.download.MetadataExtension
+            schemaRegistry {
+                url = '$schemaRegistryEndpoint'
+                download {
+                    metadata = new MetadataExtension(true, '${folderRule.root.absolutePath}/src/main/avro/metadata')
+
+                    subject('$subjectName', '${folderRule.root.absolutePath}/src/main/avro/test', true)
+                }
+            }
+        """
+        )
+
+        // When
+        val result: BuildResult? = GradleRunner.create()
+            .withGradleVersion("6.7.1")
+            .withProjectDir(folderRule.root)
+            .withArguments(DownloadTask.TASK_NAME)
+            .withPluginClasspath()
+            .withDebug(true)
+            .build()
+
+        // Then
+        val schemaFile = "$subjectName.avsc"
+        val metadataFile = "$subjectName-metadata.json"
+        Assertions.assertThat(File(folderRule.root, "src/main/avro/test/$schemaFile")).exists()
+        Assertions.assertThat(File(folderRule.root, "src/main/avro/metadata/$metadataFile")).exists()
+
+        val libSchemaFile = "$subjectNameLib.avsc"
+        val libMetadataFile = "$subjectNameLib-metadata.json"
+        Assertions.assertThat(File(folderRule.root, "src/main/avro/test/$libSchemaFile")).exists()
+        Assertions.assertThat(File(folderRule.root, "src/main/avro/metadata/$libMetadataFile")).exists()
+
+        Assertions.assertThat(result?.task(":downloadSchemasTask")?.outcome).isEqualTo(TaskOutcome.SUCCESS)
+    }
+
     private class SchemaArgumentProvider : ArgumentsProvider {
         override fun provideArguments(context: ExtensionContext): Stream<out Arguments> =
             Stream.of(
diff --git a/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadSubjectExtension.kt b/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadSubjectExtension.kt
index c0a763e..79d1ffe 100644
--- a/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadSubjectExtension.kt
+++ b/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadSubjectExtension.kt
@@ -22,21 +22,41 @@ open class DownloadSubjectExtension(objects: ObjectFactory) {
         subjects.add(DownloadSubject(inputSubject, outputPath))
     }
 
+    fun subject(inputSubject: String, outputPath: String, downloadReferences: Boolean) {
+        subjects.add(DownloadSubject(inputSubject, outputPath, downloadReferences = downloadReferences))
+    }
+
     fun subject(inputSubject: String, outputPath: String, outputFileName: String) {
         subjects.add(DownloadSubject(inputSubject, outputPath, outputFileName = outputFileName))
     }
 
+    fun subject(inputSubject: String, outputPath: String, outputFileName: String, downloadReferences: Boolean) {
+        subjects.add(DownloadSubject(inputSubject, outputPath, outputFileName = outputFileName, downloadReferences = downloadReferences))
+    }
+
     fun subject(inputSubject: String, outputPath: String, version: Int) {
         subjects.add(DownloadSubject(inputSubject, outputPath, version))
     }
 
+    fun subject(inputSubject: String, outputPath: String, version: Int, downloadReferences: Boolean) {
+        subjects.add(DownloadSubject(inputSubject, outputPath, version, downloadReferences = downloadReferences))
+    }
+
     fun subject(inputSubject: String, outputPath: String, version: Int, outputFileName: String) {
         subjects.add(DownloadSubject(inputSubject, outputPath, version, outputFileName = outputFileName))
     }
 
+    fun subject(inputSubject: String, outputPath: String, version: Int, outputFileName: String, downloadReferences: Boolean) {
+        subjects.add(DownloadSubject(inputSubject, outputPath, version, outputFileName = outputFileName, downloadReferences = downloadReferences))
+    }
+
     fun subjectPattern(inputPattern: String, outputPath: String) {
         subjects.add(DownloadSubject(inputPattern, outputPath, null, true))
     }
+
+    fun subjectPattern(inputPattern: String, outputPath: String, downloadReferences: Boolean) {
+        subjects.add(DownloadSubject(inputPattern, outputPath, null, true, downloadReferences = downloadReferences))
+    }
 }
 
 data class DownloadSubject(
@@ -44,7 +64,8 @@ data class DownloadSubject(
     val outputPath: String,
     val version: Int? = null,
     val regex: Boolean = false,
-    val outputFileName: String? = null
+    val outputFileName: String? = null,
+    val downloadReferences: Boolean = false
 )
 
 data class MetadataExtension(
@@ -52,4 +73,4 @@ data class MetadataExtension(
     val outputPath: String? = null
 ) {
     constructor(enabled: Boolean) : this(enabled, null)
-}
\ No newline at end of file
+}
diff --git a/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskAction.kt b/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskAction.kt
index f9b89ea..d94270d 100644
--- a/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskAction.kt
+++ b/src/main/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskAction.kt
@@ -36,14 +36,37 @@ class DownloadTaskAction(
                 val metadata = getSchemaMetadata(downloadSubject)
                 val outputDir = File(rootDir.toURI()).resolve(downloadSubject.outputPath)
                 outputDir.mkdirs()
-                if (metadataConfiguration.enabled) {
+                val metadataDirectory = if (metadataConfiguration.enabled) {
                     val metadataDirectory = metadataConfiguration.outputPath?.run {
                         File(rootDir.toURI()).resolve(this)
                     } ?: outputDir
                     metadataDirectory.mkdirs()
+                    metadataDirectory
+                } else {
+                    outputDir
+                }
+
+                if (metadataConfiguration.enabled) {
                     writeSchemaMetadata(downloadSubject, metadata, metadataDirectory)
                 }
                 writeSchemaFile(downloadSubject, metadata, outputDir)
+
+                if (downloadSubject.downloadReferences) {
+                    metadata.references.forEach {
+                        logger.infoIfNotQuiet("Start downloading referenced schema ${it.subject}/${it.version} for ${downloadSubject.subject}")
+                        val referenceSubject = DownloadSubject(
+                            subject = it.subject,
+                            outputPath = downloadSubject.outputPath,
+                            version = it.version
+                        )
+                        val referenceMetadata = getSchemaMetadata(referenceSubject)
+
+                        if (metadataConfiguration.enabled) {
+                            writeSchemaMetadata(referenceSubject, referenceMetadata, metadataDirectory)
+                        }
+                        writeSchemaFile(referenceSubject, referenceMetadata, outputDir)
+                    }
+                }
             } catch (e: Exception) {
                 logger.error("Error during schema retrieval for ${downloadSubject.subject}", e)
                 errorCount++


Test patch:
diff --git a/src/test/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskActionTest.kt b/src/test/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskActionTest.kt
index dd5b10f..7fbd913 100644
--- a/src/test/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskActionTest.kt
+++ b/src/test/kotlin/com/github/imflog/schema/registry/tasks/download/DownloadTaskActionTest.kt
@@ -3,6 +3,7 @@ package com.github.imflog.schema.registry.tasks.download
 import io.confluent.kafka.schemaregistry.avro.AvroSchema
 import io.confluent.kafka.schemaregistry.avro.AvroSchemaProvider
 import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient
+import io.confluent.kafka.schemaregistry.client.rest.entities.SchemaReference
 import io.confluent.kafka.schemaregistry.json.JsonSchemaProvider
 import io.confluent.kafka.schemaregistry.protobuf.ProtobufSchemaProvider
 import org.assertj.core.api.Assertions
@@ -63,7 +64,7 @@ class DownloadTaskActionTest {
                 DownloadSubject(testSubject, outputDir),
                 DownloadSubject(fooSubject, outputDir)
             ),
-            MetadataExtension(),
+            MetadataExtension()
         ).run()
 
         // then
@@ -136,7 +137,7 @@ class DownloadTaskActionTest {
             arrayListOf(
                 DownloadSubject("te.*", outputDir, null, true)
             ),
-            MetadataExtension(),
+            MetadataExtension()
         ).run()
 
         // then
@@ -184,7 +185,7 @@ class DownloadTaskActionTest {
             registryClient,
             folderRoot,
             arrayListOf(DownloadSubject(subject, outputDir)),
-            MetadataExtension(),
+            MetadataExtension()
         ).run()
 
         // then
@@ -220,7 +221,7 @@ class DownloadTaskActionTest {
                 DownloadSubject(invalidSubjectPattern, outputDir, null, true),
                 DownloadSubject("test", outputDir)
             ),
-            MetadataExtension(),
+            MetadataExtension()
         ).run()
 
         // then
@@ -274,7 +275,7 @@ class DownloadTaskActionTest {
             arrayListOf(
                 DownloadSubject("test", outputDir, v1Id)
             ),
-            MetadataExtension(),
+            MetadataExtension()
         ).run()
 
         // Then
@@ -317,7 +318,7 @@ class DownloadTaskActionTest {
             registryClient,
             folderRoot,
             arrayListOf(DownloadSubject(testSubject, outputDir)),
-            MetadataExtension(true),
+            MetadataExtension(true)
         ).run()
 
         // then
@@ -365,7 +366,7 @@ class DownloadTaskActionTest {
             registryClient,
             folderRoot,
             arrayListOf(DownloadSubject(testSubject, outputDir)),
-            MetadataExtension(true, metadataDir),
+            MetadataExtension(true, metadataDir)
         ).run()
 
         // then
@@ -380,4 +381,149 @@ class DownloadTaskActionTest {
             .containsIgnoringCase("\"schema\" :")
             .containsIgnoringCase("\"references\" :")
     }
+
+    @Test
+    fun `Should download referenced schemas with metadata in the specified output dir`() {
+        // given
+        val testSubject = "test"
+        val testLibSubject = "test_lib"
+        val outputDir = "src/main/avro/external"
+        val metadataDir = "src/main/avro/metadata"
+
+        val registryClient = MockSchemaRegistryClient(listOf(AvroSchemaProvider()))
+
+        registryClient.register(
+            testLibSubject,
+            registryClient.parseSchema(
+                AvroSchema.TYPE,
+                """{
+                    "type": "record",
+                    "name": "test_lib",
+                    "fields": [
+                        { "name": "name", "type": "string" }
+                    ]
+                }""",
+                listOf()
+            ).get()
+        )
+
+        registryClient.register(
+            testSubject,
+            registryClient.parseSchema(
+                AvroSchema.TYPE,
+                """{
+                    "type": "record",
+                    "name": "test",
+                    "fields": [
+                        { "name": "name", "type": "test_lib" }
+                    ]
+                }""",
+                listOf(
+                    SchemaReference(
+                        "name",
+                        testLibSubject,
+                        1
+                    )
+                )
+            ).get()
+        )
+
+        folderRule.resolve("src/main/avro/external").toFile().mkdir()
+        val folderRoot = folderRule.toFile()
+
+        // when
+        val errorCount = DownloadTaskAction(
+            registryClient,
+            folderRoot,
+            arrayListOf(DownloadSubject(testSubject, outputDir, downloadReferences = true)),
+            MetadataExtension(true, metadataDir)
+        ).run()
+
+        // then
+        Assertions.assertThat(errorCount).isEqualTo(0)
+        Assertions.assertThat(File(folderRoot, "src/main/avro/external/test.avsc")).isNotNull
+        Assertions.assertThat(File(folderRoot, "src/main/avro/metadata/test-metadata.json")).isNotNull
+        // Would be cleaner to use a JSON assertion library but I am not sure this is really required for now
+        Assertions.assertThat(File(folderRoot, "src/main/avro/metadata/test-metadata.json").readText())
+            .containsIgnoringCase("\"id\" :")
+            .containsIgnoringCase("\"version\" :")
+            .containsIgnoringCase("\"schema_type\" :")
+            .containsIgnoringCase("\"schema\" :")
+            .containsIgnoringCase("\"references\" :")
+
+        Assertions.assertThat(File(folderRoot, "src/main/avro/external/test_lib.avsc")).isNotNull
+        Assertions.assertThat(File(folderRoot, "src/main/avro/metadata/test_lib-metadata.json")).isNotNull
+        // Would be cleaner to use a JSON assertion library but I am not sure this is really required for now
+        Assertions.assertThat(File(folderRoot, "src/main/avro/metadata/test_lib-metadata.json").readText())
+            .containsIgnoringCase("\"id\" :")
+            .containsIgnoringCase("\"version\" :")
+            .containsIgnoringCase("\"schema_type\" :")
+            .containsIgnoringCase("\"schema\" :")
+            .containsIgnoringCase("\"references\" :")
+    }
+
+    @Test
+    fun `Should download referenced schemas in the specified output dir`() {
+        // given
+        val testSubject = "test"
+        val testLibSubject = "test_lib"
+        val outputDir = "src/main/avro/external"
+        val metadataDir = "src/main/avro/metadata"
+
+        val registryClient = MockSchemaRegistryClient(listOf(AvroSchemaProvider()))
+
+        registryClient.register(
+            testLibSubject,
+            registryClient.parseSchema(
+                AvroSchema.TYPE,
+                """{
+                    "type": "record",
+                    "name": "test_lib",
+                    "fields": [
+                        { "name": "name", "type": "string" }
+                    ]
+                }""",
+                listOf()
+            ).get()
+        )
+
+        registryClient.register(
+            testSubject,
+            registryClient.parseSchema(
+                AvroSchema.TYPE,
+                """{
+                    "type": "record",
+                    "name": "test",
+                    "fields": [
+                        { "name": "name", "type": "test_lib" }
+                    ]
+                }""",
+                listOf(
+                    SchemaReference(
+                        "name",
+                        testLibSubject,
+                        1
+                    )
+                )
+            ).get()
+        )
+
+        folderRule.resolve("src/main/avro/external").toFile().mkdir()
+        val folderRoot = folderRule.toFile()
+
+        // when
+        val errorCount = DownloadTaskAction(
+            registryClient,
+            folderRoot,
+            arrayListOf(DownloadSubject(testSubject, outputDir, downloadReferences = true)),
+            MetadataExtension()
+        ).run()
+
+        // then
+        Assertions.assertThat(errorCount).isEqualTo(0)
+        Assertions.assertThat(File(folderRoot, "src/main/avro/external/test.avsc")).isNotNull
+
+        Assertions.assertThat(File(folderRoot, "src/main/avro/external/test_lib.avsc")).isNotNull
+        // Would be cleaner to use a JSON assertion library but I am not sure this is really required for now
+    }
 }