[SchemaRegistry] add auto register schemas (Azure#20479)

fixes: Azure#20306

Author: swathipil

sdk/schemaregistry/azure-schemaregistry-avroserializer/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- `auto_register_schemas` keyword argument has been added to `SchemaRegistryAvroSerializer`, which will allow for automatically registering schemas passed in to the `serialize`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/schemaregistry/azure-schemaregistry-avroserializer/azure/schemaregistry/serializer/avroserializer/_schema_registry_avro_serializer.py

@@ -40,14 +40,23 @@ class SchemaRegistryAvroSerializer(object):
      which is used to register schema and retrieve schema from the service.
     :type schema_registry: ~azure.schemaregistry.SchemaRegistryClient
     :param str schema_group: Schema group under which schema should be registered.
+    :keyword bool auto_register_schemas: When true, register new schemas passed to serialize.
+     Otherwise, and by default, fail if it has not been pre-registered in the registry.
     :keyword str codec: The writer codec. If None, let the avro library decides.
 
     """
+
     def __init__(self, schema_registry, schema_group, **kwargs):
         # type: ("SchemaRegistryClient", str, Any) -> None
         self._schema_group = schema_group
         self._avro_serializer = AvroObjectSerializer(codec=kwargs.get("codec"))
         self._schema_registry_client = schema_registry  # type: "SchemaRegistryClient"
+        self._auto_register_schemas = kwargs.get("auto_register_schemas", False)
+        self._auto_register_schema_func = (
+                self._schema_registry_client.register_schema
+                if self._auto_register_schemas
+                else self._schema_registry_client.get_schema_id
+            )
         self._id_to_schema = {}
         self._schema_to_id = {}
         self._user_input_schema_cache = {}
@@ -63,7 +72,7 @@ def __exit__(self, *exc_details):
 
     def close(self):
         # type: () -> None
-        """ This method is to close the sockets opened by the client.
+        """This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
         self._schema_registry_client.close()
@@ -85,12 +94,8 @@ def _get_schema_id(self, schema_name, schema, **kwargs):
         try:
             return self._schema_to_id[schema_str]
         except KeyError:
-            schema_id = self._schema_registry_client.register_schema(
-                self._schema_group,
-                schema_name,
-                "Avro",
-                schema_str,
-                **kwargs
+            schema_id = self._auto_register_schema_func(
+                self._schema_group, schema_name, "Avro", schema_str, **kwargs
             ).schema_id
             self._schema_to_id[schema_str] = schema_id
             self._id_to_schema[schema_id] = schema_str
@@ -108,7 +113,9 @@ def _get_schema(self, schema_id, **kwargs):
         try:
             return self._id_to_schema[schema_id]
         except KeyError:
-            schema_str = self._schema_registry_client.get_schema(schema_id, **kwargs).schema_content
+            schema_str = self._schema_registry_client.get_schema(
+                schema_id, **kwargs
+            ).schema_content
             self._id_to_schema[schema_id] = schema_str
             self._schema_to_id[schema_str] = schema_id
             return schema_str
@@ -133,14 +140,14 @@ def serialize(self, data, schema, **kwargs):
             self._user_input_schema_cache[raw_input_schema] = parsed_schema
             cached_schema = parsed_schema
 
-        record_format_identifier = b'\0\0\0\0'
+        record_format_identifier = b"\0\0\0\0"
         schema_id = self._get_schema_id(cached_schema.fullname, cached_schema, **kwargs)
         data_bytes = self._avro_serializer.serialize(data, cached_schema)
 
         stream = BytesIO()
 
         stream.write(record_format_identifier)
-        stream.write(schema_id.encode('utf-8'))
+        stream.write(schema_id.encode("utf-8"))
         stream.write(data_bytes)
         stream.flush()
 
@@ -157,8 +164,12 @@ def deserialize(self, data, **kwargs):
         :rtype: Dict[str, Any]
         """
         # record_format_identifier = data[0:4]  # The first 4 bytes are retained for future record format identifier.
-        schema_id = data[SCHEMA_ID_START_INDEX:(SCHEMA_ID_START_INDEX + SCHEMA_ID_LENGTH)].decode('utf-8')
+        schema_id = data[
+            SCHEMA_ID_START_INDEX : (SCHEMA_ID_START_INDEX + SCHEMA_ID_LENGTH)
+        ].decode("utf-8")
         schema_content = self._get_schema(schema_id, **kwargs)
 
-        dict_data = self._avro_serializer.deserialize(data[DATA_START_INDEX:], schema_content)
+        dict_data = self._avro_serializer.deserialize(
+            data[DATA_START_INDEX:], schema_content
+        )
         return dict_data

sdk/schemaregistry/azure-schemaregistry-avroserializer/tests/recordings/test_avro_serializer.test_basic_sr_avro_serializer.yaml

@@ -28,7 +28,7 @@ interactions:
       content-type:
       - application/json
       date:
-      - Wed, 18 Aug 2021 15:11:19 GMT
+      - Wed, 01 Sep 2021 17:08:23 GMT
       location:
       - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/fakegroup/schemas/example.avro.User/versions/1?api-version=2017-04
       server:

sdk/schemaregistry/azure-schemaregistry-avroserializer/tests/recordings/test_avro_serializer.test_basic_sr_avro_serializer_with_auto_register_schemas.yaml

@@ -0,0 +1,53 @@
+interactions:
+- request:
+    body: '"{\"type\": \"record\", \"name\": \"User\", \"namespace\": \"example.avro\",
+      \"fields\": [{\"type\": \"string\", \"name\": \"name\"}, {\"type\": [\"int\",
+      \"null\"], \"name\": \"favorite_number\"}, {\"type\": [\"string\", \"null\"],
+      \"name\": \"favorite_color\"}]}"'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '265'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-azureschemaregistry/1.0.0b1 Python/3.9.0 (Windows-10-10.0.19041-SP0)
+      X-Schema-Type:
+      - Avro
+    method: PUT
+    uri: https://fake_resource.servicebus.windows.net/$schemagroups/fakegroup/schemas/example.avro.User?api-version=2017-04
+  response:
+    body:
+      string: '{"id":"576838e0558c43f8b85cdaadbd4561f5"}'
+    headers:
+      content-type:
+      - application/json
+      date:
+      - Wed, 01 Sep 2021 17:06:46 GMT
+      location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/fakegroup/schemas/example.avro.User/versions/1?api-version=2017-04
+      server:
+      - Microsoft-HTTPAPI/2.0
+      strict-transport-security:
+      - max-age=31536000
+      transfer-encoding:
+      - chunked
+      x-schema-id:
+      - 576838e0558c43f8b85cdaadbd4561f5
+      x-schema-id-location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/getschemabyid/576838e0558c43f8b85cdaadbd4561f5?api-version=2017-04
+      x-schema-type:
+      - Avro
+      x-schema-version:
+      - '1'
+      x-schema-versions-location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/swathip-test-schema/schemas/example.avro.User/versions?api-version=2017-04
+    status:
+      code: 200
+      message: OK
+version: 1

sdk/schemaregistry/azure-schemaregistry-avroserializer/tests/recordings/test_avro_serializer.test_basic_sr_avro_serializer_without_auto_register_schemas.yaml

@@ -0,0 +1,53 @@
+interactions:
+- request:
+    body: '"{\"type\": \"record\", \"name\": \"User\", \"namespace\": \"example.avro\",
+      \"fields\": [{\"type\": \"string\", \"name\": \"name\"}, {\"type\": [\"int\",
+      \"null\"], \"name\": \"favorite_number\"}, {\"type\": [\"string\", \"null\"],
+      \"name\": \"favorite_color\"}]}"'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '265'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-azureschemaregistry/1.0.0b1 Python/3.9.0 (Windows-10-10.0.19041-SP0)
+      X-Schema-Type:
+      - Avro
+    method: POST
+    uri: https://fake_resource.servicebus.windows.net/$schemagroups/fakegroup/schemas/example.avro.User?api-version=2017-04
+  response:
+    body:
+      string: '{"id":"576838e0558c43f8b85cdaadbd4561f5"}'
+    headers:
+      content-type:
+      - application/json
+      date:
+      - Wed, 01 Sep 2021 17:08:24 GMT
+      location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/fakegroup/schemas/example.avro.User/versions/1?api-version=2017-04
+      server:
+      - Microsoft-HTTPAPI/2.0
+      strict-transport-security:
+      - max-age=31536000
+      transfer-encoding:
+      - chunked
+      x-schema-id:
+      - 576838e0558c43f8b85cdaadbd4561f5
+      x-schema-id-location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/getschemabyid/576838e0558c43f8b85cdaadbd4561f5?api-version=2017-04
+      x-schema-type:
+      - Avro
+      x-schema-version:
+      - '1'
+      x-schema-versions-location:
+      - https://swathip-test-eventhubs.servicebus.windows.net:443/$schemagroups/swathip-test-schema/schemas/example.avro.User/versions?api-version=2017-04
+    status:
+      code: 200
+      message: OK
+version: 1

sdk/schemaregistry/azure-schemaregistry-avroserializer/tests/test_avro_serializer.py

@@ -76,6 +76,33 @@ def test_raw_avro_serializer_negative(self):
 
     @SchemaRegistryPowerShellPreparer()
     def test_basic_sr_avro_serializer(self, schemaregistry_endpoint, schemaregistry_group, **kwargs):
+        sr_client = self.create_basic_client(SchemaRegistryClient, endpoint=schemaregistry_endpoint)
+        sr_avro_serializer = SchemaRegistryAvroSerializer(sr_client, schemaregistry_group, auto_register_schemas=True)
+
+        schema_str = """{"namespace":"example.avro","type":"record","name":"User","fields":[{"name":"name","type":"string"},{"name":"favorite_number","type":["int","null"]},{"name":"favorite_color","type":["string","null"]}]}"""
+        schema = avro.schema.parse(schema_str)
+
+        dict_data = {"name": u"Ben", "favorite_number": 7, "favorite_color": u"red"}
+        encoded_data = sr_avro_serializer.serialize(dict_data, schema_str)
+
+        assert schema_str in sr_avro_serializer._user_input_schema_cache
+        assert str(avro.schema.parse(schema_str)) in sr_avro_serializer._schema_to_id
+
+        assert encoded_data[0:4] == b'\0\0\0\0'
+        schema_id = sr_client.get_schema_id(schemaregistry_group, schema.fullname, "Avro", str(schema)).schema_id
+        assert encoded_data[4:36] == schema_id.encode("utf-8")
+
+        assert schema_id in sr_avro_serializer._id_to_schema
+
+        decoded_data = sr_avro_serializer.deserialize(encoded_data)
+        assert decoded_data["name"] == u"Ben"
+        assert decoded_data["favorite_number"] == 7
+        assert decoded_data["favorite_color"] == u"red"
+
+        sr_avro_serializer.close()
+
+    @SchemaRegistryPowerShellPreparer()
+    def test_basic_sr_avro_serializer_without_auto_register_schemas(self, schemaregistry_endpoint, schemaregistry_group, **kwargs):
         sr_client = self.create_basic_client(SchemaRegistryClient, endpoint=schemaregistry_endpoint)
         sr_avro_serializer = SchemaRegistryAvroSerializer(sr_client, schemaregistry_group)