Correct SchemaReferences in APIs that return Schemas (confluentinc#1304)

RickTalken · rtalken · anchitj · web-flow · commit f3055be4d2aa · 2023-02-28T16:53:27.000+05:30
* Correct SchemaReferences in APIs that return Schemas

* Add CHANGELOG

---------

Co-authored-by: rtalken &lt;rtalken@farmobile.com&gt;
Co-authored-by: Anchit Jain &lt;anjain@confluent.io&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@
 - Wheels for Linux / arm64 (#1496).
 - Added support for Default num_partitions in CreateTopics Admin API.
 - Added support for password protected private key in CachedSchemaRegistryClient.
+- Add reference support in Schema Registry client. (@RickTalken, #1304)
 
 ## v2.0.2
 
diff --git a/src/confluent_kafka/schema_registry/schema_registry_client.py b/src/confluent_kafka/schema_registry/schema_registry_client.py
@@ -370,12 +370,10 @@ def get_schema(self, schema_id):
         schema = Schema(schema_str=response['schema'],
                         schema_type=response.get('schemaType', 'AVRO'))
 
-        refs = []
-        for ref in response.get('references', []):
-            refs.append(SchemaReference(name=ref['name'],
-                                        subject=ref['subject'],
-                                        version=ref['version']))
-        schema.references = refs
+        schema.references = [
+            SchemaReference(name=ref['name'], subject=ref['subject'], version=ref['version'])
+            for ref in response.get('references', [])
+        ]
 
         self._cache.set(schema_id, schema)
 
@@ -419,7 +417,12 @@ def lookup_schema(self, subject_name, schema, normalize_schemas=False):
         return RegisteredSchema(schema_id=response['id'],
                                 schema=Schema(response['schema'],
                                               schema_type,
-                                              response.get('references', [])),
+                                              [
+                                                  SchemaReference(name=ref['name'],
+                                                                  subject=ref['subject'],
+                                                                  version=ref['version'])
+                                                  for ref in response.get('references', [])
+                                              ]),
                                 subject=response['subject'],
                                 version=response['version'])
 
@@ -493,7 +496,12 @@ def get_latest_version(self, subject_name):
         return RegisteredSchema(schema_id=response['id'],
                                 schema=Schema(response['schema'],
                                               schema_type,
-                                              response.get('references', [])),
+                                              [
+                                                  SchemaReference(name=ref['name'],
+                                                                  subject=ref['subject'],
+                                                                  version=ref['version'])
+                                                  for ref in response.get('references', [])
+                                              ]),
                                 subject=response['subject'],
                                 version=response['version'])
 
@@ -524,7 +532,12 @@ def get_version(self, subject_name, version):
         return RegisteredSchema(schema_id=response['id'],
                                 schema=Schema(response['schema'],
                                               schema_type,
-                                              response.get('references', [])),
+                                              [
+                                                  SchemaReference(name=ref['name'],
+                                                                  subject=ref['subject'],
+                                                                  version=ref['version'])
+                                                  for ref in response.get('references', [])
+                                              ]),
                                 subject=response['subject'],
                                 version=response['version'])
 
