schema_registry_client 0.0.5 → 0.0.6

data/lib/schema_registry_client/output/json_schema.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+class SchemaRegistry
+  module Output
+    module JsonSchema
+      class << self
+        def fetch(message_name)
+          name = message_name.start_with?('.') ? message_name[1..] : message_name
+          Google::Protobuf::DescriptorPool.generated_pool.lookup(name)
+        end
+
+        def output(descriptor, path: nil)
+          properties = {}
+          result = {
+            '$schema': 'https://json-schema.org/draft/2020-12/schema',
+            type: 'object',
+            properties: properties
+          }
+          if path
+            # follow path down
+            parts = path.split('.')
+            field_name = parts.last
+            parts[...-1].each do |part|
+              field = descriptor.field.find { |f| f.name == part }
+              raise "Field #{part} not found in #{descriptor.name}" unless field
+
+              descriptor = fetch(field.type_name)&.to_proto
+            end
+            result[:required] = [field_name]
+            properties[field_name] = field_object(descriptor.field.find { |f| f.name == field_name.to_s })
+          else
+            result[:required] = descriptor.field.reject(&:proto3_optional).map(&:name)
+            descriptor.field.each do |f|
+              properties[f.name] = field_object(f)
+            end
+          end
+          JSON.pretty_generate(result)
+        end
+
+        def field_object(field, ignore_repeated: false)
+          klass = fetch(field.type_name)&.to_proto
+          if field.label == :LABEL_REPEATED && !ignore_repeated
+            if klass&.options.respond_to?(:map_entry) && klass.options.map_entry
+              return {
+                type: 'object',
+                additionalProperties: field_object(klass.field[1])
+              }
+            end
+            return {
+              type: 'array',
+              items: field_object(field, ignore_repeated: true)
+            }
+          end
+          field_type(field, klass)
+        end
+
+        def field_type(field, klass)
+          case field.type
+          when :TYPE_INT32, :TYPE_UINT32, :TYPE_SINT32, :TYPE_FIXED32, :TYPE_SFIXED32
+            { type: 'integer' }
+          when :TYPE_FLOAT, :TYPE_DOUBLE
+            { type: 'number' }
+          when :TYPE_INT64, :TYPE_UINT64, :TYPE_SINT64, :TYPE_FIXED64, :TYPE_SFIXED64, :TYPE_STRING, :TYPE_BYTES
+            { type: 'string' }
+          when :TYPE_BOOL
+            { type: 'boolean' }
+          else
+            if klass.is_a?(Google::Protobuf::EnumDescriptorProto)
+              { enum: klass.to_h[:value].map { |h| h[:name] } }
+            else
+              { type: 'object' }
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/output/proto_text.rb

@@ -0,0 +1,320 @@
+# frozen_string_literal: true
+
+class SchemaRegistry
+  module Output
+    module ProtoText
+      ParseInfo = Struct.new(:writer, :package, :message) do
+        %i[write write_indent write_line writenl indent dedent].each do |method|
+          define_method(method) do |*args|
+            writer.send(method, *args)
+          end
+        end
+      end
+
+      class Writer < StringIO
+        def initialize(...)
+          super
+          @indent = 0
+        end
+
+        def write_indent(str)
+          @indent.times { write(' ') }
+          write(str)
+        end
+
+        def write_line(line, newline = 1)
+          write_indent(line)
+          newline.times { writenl }
+        end
+
+        def writenl
+          write("\n")
+        end
+
+        def indent
+          @indent += 2
+        end
+
+        def dedent
+          @indent -= 2
+        end
+      end
+
+      class << self
+        def fetch(message_name)
+          name = message_name.start_with?('.') ? message_name[1..] : message_name
+          Google::Protobuf::DescriptorPool.generated_pool.lookup(name)
+        end
+
+        def output(file_descriptor)
+          writer = Writer.new
+          info = ParseInfo.new(writer, file_descriptor.package)
+          writer.write_line("syntax = \"#{file_descriptor.syntax}\";", 2)
+          writer.write_line("package #{file_descriptor.package};")
+          writer.writenl
+          found = false
+          file_descriptor.options.to_h.each do |name, value|
+            found = true
+            writer.write_line("option #{name} = #{value.to_json};")
+          end
+          writer.writenl if found
+
+          found = false
+          file_descriptor.dependency.each do |dependency|
+            found = true
+            writer.write_line("import \"#{dependency}\";")
+          end
+          writer.writenl if found
+
+          writer.writenl if write_options(info, file_descriptor)
+          writer.writenl if write_extensions(info, file_descriptor)
+
+          file_descriptor.enum_type.each do |enum_type|
+            write_enum(info, enum_type)
+          end
+          file_descriptor.message_type.each do |message_type|
+            write_message(info, message_type)
+          end
+          file_descriptor.service.each do |service|
+            write_service(info, service)
+          end
+          writer.string
+        end
+
+        def write_extensions(info, descriptor)
+          descriptor.extension.each do |extension|
+            info.write_line("extend #{extension.extendee[1..]} {")
+            info.indent
+            write_field(info, extension)
+            info.dedent
+            info.write_line('}')
+          end
+          descriptor.extension.any?
+        end
+
+        def write_reserved(writer, descriptor)
+          reserved = descriptor.reserved_range.map do |range|
+            range.start == range.end - 1 ? range.start.to_s : "#{range.start} to #{range.end - 1}"
+          end
+          found = false
+          if reserved.any?
+            found = true
+            writer.write_line("reserved #{reserved.join(', ')};")
+          end
+          if descriptor.reserved_name.any?
+            found = true
+            writer.write_line("reserved #{descriptor.reserved_name.map(&:to_json).join(', ')};")
+          end
+          writer.writenl if found
+        end
+
+        def write_imports(writer, file_descriptor)
+          writer.writenl
+          file_descriptor.dependency.each do |dependency|
+            writer.write_line("import \"#{dependency}\";")
+          end
+          file_descriptor.public_dependency.each do |public_dependency|
+            writer.write_line("import public \"#{public_dependency}\";")
+          end
+          file_descriptor.option_dependency.each do |option_dependency|
+            writer.write_line("import weak \"#{option_dependency}\";")
+          end
+          writer.writenl
+        end
+
+        def write_message(info, message_type)
+          info.message = message_type
+          info.write_indent('message ')
+          info.write("#{message_type.name} {")
+          info.writenl
+          info.indent
+
+          write_options(info, message_type)
+          write_reserved(info, message_type)
+
+          message_type.enum_type.each do |enum|
+            info.writenl
+            write_enum(info, enum)
+          end
+          message_type.field.each do |field|
+            write_field(info, field)
+          end
+          message_type.extension.each do |extension|
+            write_field(info, extension)
+          end
+          write_oneofs(info, message_type)
+          message_type.nested_type.each do |subtype|
+            next if subtype.options&.map_entry
+
+            info.writenl
+            write_message(info, subtype)
+          end
+          info.dedent
+          info.write_line('}')
+        end
+
+        def field_type(info, field_type)
+          case field_type.type
+          when :TYPE_INT32
+            'int32'
+          when :TYPE_INT64
+            'int64'
+          when :TYPE_UINT32
+            'uint32'
+          when :TYPE_UINT64
+            'uint64'
+          when :TYPE_SINT32
+            'sint32'
+          when :TYPE_SINT64
+            'sint64'
+          when :TYPE_FIXED32
+            'fixed32'
+          when :TYPE_FIXED64
+            'fixed64'
+          when :TYPE_SFIXED32
+            'sfixed32'
+          when :TYPE_SFIXED64
+            'sfixed64'
+          when :TYPE_FLOAT
+            'float'
+          when :TYPE_DOUBLE
+            'double'
+          when :TYPE_BOOL
+            'bool'
+          when :TYPE_STRING
+            'string'
+          when :TYPE_BYTES
+            'bytes'
+          when :TYPE_ENUM, :TYPE_MESSAGE
+            # remove leading .
+            type = fetch(field_type.type_name[1..])
+            name = type.name.sub("#{info.package}.#{info.message.name}.", '')
+            name.sub("#{info.package}.", '')
+          end
+        end
+
+        def write_field(info, field, oneof: false)
+          return if !oneof && field.has_oneof_index?
+
+          info.write_indent('')
+
+          klass = nil
+          klass = fetch(field.type_name).to_proto if field.type_name && field.type_name != ''
+
+          if field.proto3_optional
+            info.write('optional ')
+          elsif field.label == :LABEL_REPEATED && !klass&.options&.map_entry
+            info.write('repeated ')
+          end
+
+          if klass&.options&.map_entry
+            info.write("map<#{field_type(info, klass.field[0])}, #{field_type(info, klass.field[1])}>")
+          else
+            info.write(field_type(info, field).to_s)
+          end
+          info.write(" #{field.name} = #{field.number}")
+
+          write_field_options(info, field)
+          info.write(';')
+          info.writenl
+        end
+
+        def write_field_options(info, field)
+          return unless field.options
+
+          info.write(' [')
+          info.write(field.options.to_h.map { |name, value| "#{name} = #{value}" }.join(', '))
+          write_options(info, field, include_option_label: false)
+          info.write(']')
+        end
+
+        def write_oneofs(info, message)
+          message.oneof_decl.each_with_index do |oneof, i|
+            # synthetic oneof for proto3 optional fields
+            next if oneof.name.start_with?('_') &&
+                    message.field.any? { |f| f.proto3_optional && f.name == oneof.name[1..] }
+
+            info.write_line("oneof #{oneof.name} {")
+            info.indent
+            message.field.select { |f| f.has_oneof_index? && f.oneof_index == i }.each do |field|
+              write_field(info, field, oneof: true)
+            end
+            info.dedent
+            info.write_line('}')
+          end
+        end
+
+        def write_enum(info, enum_type)
+          info.write('enum ')
+          info.write("#{enum_type.name} {")
+          info.writenl
+          info.indent
+          write_reserved(info, enum_type)
+          enum_type.value.each do |value|
+            info.write_line("#{value.name} = #{value.number};")
+          end
+          info.dedent
+          info.write_line('}')
+          info.writenl
+        end
+
+        def method_type(package, name)
+          output = name.sub("#{package}.", '')
+          output = output[1..] if output.start_with?('.')
+          output
+        end
+
+        def write_service(info, service)
+          info.write_line("service #{service.name} {")
+          info.indent
+          service['method'].each do |method|
+            info.write_indent("rpc #{method.name}(#{method_type(info.package, method.input_type)}) ")
+            info.write("returns (#{method_type(info.package, method.output_type)}) {")
+            info.writenl
+            info.indent
+            write_options(info, method) if method.options
+            info.dedent
+            info.write_line('};')
+          end
+          info.dedent
+          info.write_line('}')
+        end
+
+        # @return [Boolean] true if any options were written
+        def write_options(info, descriptor, include_option_label: true)
+          # unfortunately there doesn't seem to be a way to get the full list of options without
+          # resorting to to_json.
+          json = JSON.parse(descriptor.options.to_json)
+          return if !json || json.empty?
+
+          found = false
+          json.each_key do |name|
+            option_name = name.tr('[]', '')
+            ext = fetch(option_name)
+            next if ext.nil?
+
+            found = true
+            options = ext.get(descriptor.options)
+            if include_option_label
+              info.write_indent("option (#{option_name}) =")
+            else
+              info.write("(#{option_name}) = ")
+            end
+            if options.respond_to?(:to_h)
+              lines = JSON.pretty_generate(options.to_h).lines(chomp: true)
+              lines.each_with_index do |line, i|
+                info.write_indent(line)
+                info.writenl if i < lines.length - 1
+              end
+              info.write(';')
+            else
+              info.write(options.to_json)
+            end
+            info.writenl
+          end
+          found
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/schema/avro.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'schema_registry_client/schema/base'
+require 'schema_registry_client/avro_schema_store'
+
+class SchemaRegistry
+  module Schema
+    class Avro < Base
+      DEFAULT_SCHEMAS_PATH = './schemas'
+
+      class << self
+        def schema_type
+          'AVRO'
+        end
+
+        def schema_store
+          @schema_store ||= SchemaRegistry::AvroSchemaStore.new(
+            path: SchemaRegistry.avro_schema_path || DEFAULT_SCHEMAS_PATH
+          )
+          @schema_store.load_schemas!
+          @schema_store
+        end
+
+        def schema_text(_message, schema_name: nil)
+          schema_store.find_text(schema_name)
+        end
+
+        def encode(message, stream, schema_name: nil)
+          validate_options = { recursive: true,
+                               encoded: false,
+                               fail_on_extra_fields: true }
+          schema = schema_store.find(schema_name)
+
+          ::Avro::SchemaValidator.validate!(schema, message, **validate_options)
+
+          writer = ::Avro::IO::DatumWriter.new(schema)
+          encoder = ::Avro::IO::BinaryEncoder.new(stream)
+          writer.write(message, encoder)
+        end
+
+        def decode(stream, schema_text)
+          # Parse the schema text from the registry into an Avro schema object
+          JSON.parse(schema_text)
+          writers_schema = ::Avro::Schema.parse(schema_text)
+
+          decoder = ::Avro::IO::BinaryDecoder.new(stream)
+
+          # Try to find the reader schema locally, fall back to writer schema
+          readers_schema = begin
+            schema_store.find(writers_schema.fullname)
+          rescue StandardError
+            writers_schema
+          end
+
+          reader = ::Avro::IO::DatumReader.new(writers_schema, readers_schema)
+          reader.read(decoder)
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/schema/base.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+class SchemaRegistry
+  module Schema
+    class MissingImplementationError < StandardError; end
+
+    class Base
+      class << self
+        # @param message [Object]
+        # @param schema_name [String]
+        # @return [String]
+        def schema_text(_message, schema_name: nil)
+          raise MissingImplementationError, 'Subclasses must implement schema_text'
+        end
+
+        # @return [String]
+        def schema_type
+          raise MissingImplementationError, 'Subclasses must implement schema_type'
+        end
+
+        # @param message [Object]
+        # @param stream [StringIO]
+        # @param schema_name [String]
+        def encode(_message, _stream, schema_name: nil)
+          raise MissingImplementationError, 'Subclasses must implement encode'
+        end
+
+        # @param stream [StringIO]
+        # @param schema_text [String]
+        # @param registry [Object]
+        # @return [Object]
+        def decode(_stream, _schema_text)
+          raise MissingImplementationError, 'Subclasses must implement decode'
+        end
+
+        # @param message [Object]
+        # @return [Hash<String, String>]
+        def dependencies(_message)
+          {}
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/schema/proto_json_schema.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'schema_registry_client/schema/base'
+require 'schema_registry_client/output/json_schema'
+
+class SchemaRegistry
+  module Schema
+    class ProtoJsonSchema < Base
+      class << self
+        def schema_type
+          'JSON'
+        end
+
+        def schema_text(message, schema_name: nil)
+          SchemaRegistry::Output::JsonSchema.output(message.class.descriptor.to_proto)
+        end
+
+        def encode(message, stream, schema_name: nil)
+          json = message.to_h.sort.to_h.to_json
+          stream.write(json)
+        end
+
+        def decode(stream, _schema_text)
+          json = stream.read
+          JSON.parse(json)
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/schema/protobuf.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+require 'google/protobuf'
+require 'google/protobuf/well_known_types'
+require 'google/protobuf/descriptor_pb'
+require 'schema_registry_client/output/proto_text'
+require 'schema_registry_client/schema/base'
+require 'schema_registry_client/wire'
+
+class SchemaRegistry
+  module Schema
+    class Protobuf < Base
+      class << self
+        def schema_type
+          'PROTOBUF'
+        end
+
+        def schema_text(message, schema_name: nil)
+          file_descriptor = if message.is_a?(Google::Protobuf::FileDescriptor)
+                              message
+                            else
+                              message.class.descriptor.file_descriptor
+                            end
+          SchemaRegistry::Output::ProtoText.output(file_descriptor.to_proto)
+        end
+
+        def encode(message, stream, schema_name: nil)
+          _, indexes = find_index(message.class.descriptor.to_proto,
+                                  message.class.descriptor.file_descriptor.to_proto.message_type)
+
+          if indexes == [0]
+            SchemaRegistry::Wire.write_int(stream, 0)
+          else
+            SchemaRegistry::Wire.write_int(stream, indexes.length)
+            indexes.each { |i| SchemaRegistry::Wire.write_int(stream, i) }
+          end
+
+          # Now we write the actual message.
+          stream.write(message.to_proto)
+        end
+
+        def decode(stream, schema_text)
+          # See https://docs.confluent.io/platform/current/schema-registry/fundamentals/serdes-develop/index.html#wire-format
+          index_length = SchemaRegistry::Wire.read_int(stream)
+          indexes = []
+          if index_length.zero?
+            indexes.push(0)
+          else
+            index_length.times do
+              indexes.push(SchemaRegistry::Wire.read_int(stream))
+            end
+          end
+
+          encoded = stream.read
+          decode_protobuf(schema_text, encoded, indexes)
+        end
+
+        def load_schemas!
+          @all_schemas = {}
+          all_files = ObjectSpace.each_object(Google::Protobuf::FileDescriptor).to_a
+          all_files.each do |file_desc|
+            file_path = file_desc.name
+            next if file_path.start_with?('google/protobuf/') # skip built-in protos
+
+            @all_schemas[file_path] = file_desc
+          end
+        end
+
+        def dependencies(message)
+          load_schemas! unless @all_schemas&.any?
+          file_descriptor = if message.is_a?(Google::Protobuf::FileDescriptor)
+                              message
+                            else
+                              message.class.descriptor.file_descriptor
+                            end
+
+          deps = file_descriptor.to_proto.dependency.to_a
+                                .reject { |d| d.start_with?('google/protobuf/') }
+          deps.to_h do |dep|
+            dependency_schema = @all_schemas[dep]
+            [dependency_schema.name, dependency_schema]
+          end
+        end
+
+        def find_index(descriptor, messages, indexes = [])
+          messages.each_with_index do |sub_descriptor, i|
+            if sub_descriptor == descriptor
+              indexes.push(i)
+              return [true, indexes]
+            else
+              found, found_indexes = find_index(descriptor, sub_descriptor.nested_type, indexes + [i])
+              return [true, found_indexes] if found
+            end
+          end
+          []
+        end
+
+        def find_descriptor(indexes, messages)
+          first_index = indexes.shift
+          message = messages[first_index]
+          path = [message.name]
+          while indexes.length.positive?
+            message = message.nested_type[indexes.shift]
+            path.push(message.name)
+          end
+          path
+        end
+
+        def decode_protobuf(schema, encoded, indexes)
+          # get the package
+          package = schema.match(/package (\S+);/)[1]
+          # get the first message in the protobuf text
+          # TODO - get the correct message based on schema index
+          message_name = schema.match(/message (\w+) {/)[1]
+          # look up the descriptor
+          full_name = "#{package}.#{message_name}"
+          descriptor = Google::Protobuf::DescriptorPool.generated_pool.lookup(full_name)
+          unless descriptor
+            msg = "Could not find schema for #{full_name}. " \
+                  'Make sure the corresponding .proto file has been compiled and loaded.'
+            raise msg
+          end
+
+          path = find_descriptor(indexes, descriptor.file_descriptor.to_proto.message_type)
+          correct_message = Google::Protobuf::DescriptorPool.generated_pool.lookup("#{package}.#{path.join('.')}")
+          correct_message.msgclass.decode(encoded)
+        end
+      end
+    end
+  end
+end

data/lib/schema_registry_client/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class SchemaRegistry
+  VERSION = '0.0.6'
+end

data/lib/schema_registry_client/wire.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+class SchemaRegistry
+  module Wire
+    class << self
+      # Write an int with zig-zag encoding. Copied from Avro.
+      def write_int(stream, num)
+        num = (num << 1) ^ (num >> 63)
+        while (num & ~0x7F) != 0
+          stream.write(((num & 0x7f) | 0x80).chr)
+          num >>= 7
+        end
+        stream.write(num.chr)
+      end
+
+      # Read an int with zig-zag encoding. Copied from Avro.
+      def read_int(stream)
+        b = stream.readbyte
+        n = b & 0x7F
+        shift = 7
+        while (b & 0x80) != 0
+          b = stream.readbyte
+          n |= (b & 0x7F) << shift
+          shift += 7
+        end
+        (n >> 1) ^ -(n & 1)
+      end
+    end
+  end
+end