How to use org.apache.gobblin.kafka.schemareg

/**
 * This call is not cached because we never want to miss out on the latest schema.
 * {@inheritDoc}
 */
@Override
public S getLatestSchema(String name)
  throws IOException, SchemaRegistryException {
 return _kafkaSchemaRegistry.getLatestSchema(name);
}

@Override
public Schema getById(MD5Digest id)
  throws IOException, SchemaRegistryException {
 return fetchSchemaByKey(id);
}

public ConfigDrivenMd5SchemaRegistry(String name, Schema schema)
  throws IOException, SchemaRegistryException {
 this.register(name, schema);
}

@Test
public void testRegisterShouldCacheIds()
  throws IOException, SchemaRegistryException {
 KafkaSchemaRegistry<Integer, String> baseRegistry = mock(KafkaSchemaRegistry.class);
 CachingKafkaSchemaRegistry<Integer, String> cachingReg = new CachingKafkaSchemaRegistry<>(baseRegistry, 2);
 String name = "test";
 String schema1 = new String("schema");
 Integer id1 = 1;
 // first register name, schema1, get back id1
 when(baseRegistry.register(name, schema1)).thenReturn(id1);
 Assert.assertEquals(cachingReg.register(name, schema1), id1);
 // getById should hit the cache and return id1
 when(baseRegistry.getById(id1)).thenReturn(new String("schema2"));
 Assert.assertEquals(cachingReg.getById(id1), schema1);
 verify(baseRegistry, times(0)).getById(anyInt());
}

@Test
public void testRegisterSchemaCaching()
  throws IOException, SchemaRegistryException {
 KafkaSchemaRegistry<Integer, String> baseRegistry = mock(KafkaSchemaRegistry.class);
 String name = "test";
 String schema1 = new String("schema");
 Integer id1 = 1;
 CachingKafkaSchemaRegistry<Integer, String> cachingReg = new CachingKafkaSchemaRegistry<>(baseRegistry, 2);
 when(baseRegistry.register(name, schema1)).thenReturn(id1);
 Assert.assertEquals(cachingReg.register(name, schema1), id1);
 Integer id2 = 2;
 when(baseRegistry.register(name, schema1)).thenReturn(id2);
 // Test that we get back the original id
 Assert.assertEquals(cachingReg.register(name, schema1), id1);
 // Ensure that we only called baseRegistry.register once.
 verify(baseRegistry, times(1)).register(anyString(), anyString());
}

@Test
public void testIdSchemaCaching()
  throws IOException, SchemaRegistryException {
 KafkaSchemaRegistry<Integer, String> baseRegistry = mock(KafkaSchemaRegistry.class);
 String name = "test";
 String schema1 = new String("schema");
 Integer id1 = 1;
 CachingKafkaSchemaRegistry<Integer, String> cachingReg = new CachingKafkaSchemaRegistry<>(baseRegistry, 2);
 when(baseRegistry.getById(id1)).thenReturn(schema1);
 String schemaReturned = cachingReg.getById(id1);
 Assert.assertEquals(schemaReturned, schema1, "Schema returned by id should be the same");
 verify(baseRegistry, times(1)).getById(anyInt());
 when(baseRegistry.getById(id1)).thenReturn(new String("schema2"));
 Assert.assertEquals(cachingReg.getById(id1), schemaReturned);
 verify(baseRegistry, times(1)).getById(anyInt());
}

 public static KafkaSchemaRegistry getSchemaRegistry(Map<String, ?> config) {
  Properties props = new Properties();
  props.putAll(config);
  return KafkaSchemaRegistryFactory.getSchemaRegistry(props);
 }
}

 @SuppressWarnings("unchecked")
 public static KafkaSchemaRegistry getSchemaRegistry(Properties props) {
  Preconditions.checkArgument(props.containsKey(KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CLASS),
    "Missing required property " + KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CLASS);

  boolean tryCache = Boolean.parseBoolean(props.getProperty(KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CACHE,
    DEFAULT_TRY_CACHING));

  Class<?> clazz;
  try {
   clazz =
     (Class<?>) Class.forName(props.getProperty(KafkaSchemaRegistryConfigurationKeys.KAFKA_SCHEMA_REGISTRY_CLASS));
   KafkaSchemaRegistry schemaRegistry = (KafkaSchemaRegistry) ConstructorUtils.invokeConstructor(clazz, props);
   if (tryCache && !schemaRegistry.hasInternalCache())
   {
    schemaRegistry = new CachingKafkaSchemaRegistry(schemaRegistry);
   }
   return schemaRegistry;
  } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException
    | InstantiationException e) {
   log.error("Failed to instantiate " + KafkaSchemaRegistry.class, e);
   throw Throwables.propagate(e);
  }
 }
}

HttpClient httpClient = this.borrowClient();
try {
 statusCode = httpClient.executeMethod(get);
 schemaString = get.getResponseBodyAsString();
} catch (IOException e) {
 throw new SchemaRegistryException(e);
} finally {
 get.releaseConnection();
 throw new SchemaRegistryException(
   String.format("Schema with key %s cannot be retrieved, statusCode = %d", key, statusCode));
 schema = new Schema.Parser().parse(schemaString);
} catch (Throwable t) {
 throw new SchemaRegistryException(String.format("Schema with ID = %s cannot be parsed", key), t);

@Override
synchronized public S getById(K id)
  throws IOException, SchemaRegistryException {
 if (_idBasedCache.containsKey(id))
 {
  return _idBasedCache.get(id);
 }
 else
 {
  S schema = _kafkaSchemaRegistry.getById(id);
  _idBasedCache.put(id, schema);
  return schema;
 }
}

/**
 * Will throw a SchemaRegistryException if it cannot find any schema for the provided name.
 * {@inheritDoc}
 */
@Override
public Schema getLatestSchema(String name)
  throws IOException, SchemaRegistryException {
 if (_topicSchemaMap.containsKey(name))
 {
  return _topicSchemaMap.get(name);
 }
 else
 {
  throw new SchemaRegistryException("Could not find any schema for " + name);
 }
}

/**
 * Register this schema under the provided name
 * @param name
 * @param schema
 * @return
 * @throws IOException
 * @throws SchemaRegistryException
 */
@Override
public synchronized MD5Digest register(String name, Schema schema)
  throws IOException, SchemaRegistryException {
 MD5Digest md5Digest = generateId(schema);
 if (!_schemaHashMap.containsKey(md5Digest)) {
  _schemaHashMap.put(md5Digest, schema);
  _topicSchemaMap.put(name, schema);
 }
 return md5Digest;
}

/**
 * Create a caching schema registry.
 * @param kafkaSchemaRegistry: a schema registry that needs caching
 * @param maxSchemaReferences: the maximum number of unique references that can exist for a given schema.
 */
public CachingKafkaSchemaRegistry(KafkaSchemaRegistry kafkaSchemaRegistry, int maxSchemaReferences)
{
 Preconditions.checkArgument(kafkaSchemaRegistry!=null, "KafkaSchemaRegistry cannot be null");
 Preconditions.checkArgument(!kafkaSchemaRegistry.hasInternalCache(), "SchemaRegistry already has a cache.");
 _kafkaSchemaRegistry = kafkaSchemaRegistry;
 _namedSchemaCache = new HashMap<>();
 _idBasedCache = new HashMap<>();
 _maxSchemaReferences = maxSchemaReferences;
}

/**
 * Register a schema to the Kafka schema registry under the provided input name. This method will change the name
 * of the schema to the provided name if configured to do so. This is useful because certain services (like Gobblin kafka adaptor and
 * Camus) get the schema for a topic by querying for the latest schema with the topic name, requiring the topic
 * name and schema name to match for all topics. If it is not configured to switch names, this is useful for the case
 * where the Kafka topic and Avro schema names do not match. This method registers the schema to the schema registry in such a
 * way that any schema can be written to any topic.
 *
 * @param schema {@link org.apache.avro.Schema} to register.
 * @param name Name of the schema when registerd to the schema registry. This name should match the name
 *                     of the topic where instances will be published.
 * @return schema ID of the registered schema.
 * @throws SchemaRegistryException if registration failed
 */
@Override
public MD5Digest register(String name, Schema schema) throws SchemaRegistryException {
 PostMethod post = new PostMethod(url);
 if (this.switchTopicNames) {
  return register(AvroUtils.switchName(schema, name), post);
 } else {
  post.addParameter("name", name);
  return register(schema, post);
 }
}

@Test
public void testMaxReferences()
  throws IOException, SchemaRegistryException {
 KafkaSchemaRegistry<Integer, String> baseRegistry = mock(KafkaSchemaRegistry.class);
 String name = "test";
 String schema1 = new String("schema");
 String schema2 = new String("schema");
 String schema3 = new String("schema");
 Integer id1 = 1;
 Integer id2 = 2;
 Integer id3 = 3;
 CachingKafkaSchemaRegistry<Integer, String> cachingReg = new CachingKafkaSchemaRegistry<>(baseRegistry, 2);
 when(baseRegistry.register(name, schema1)).thenReturn(id1);
 Assert.assertEquals(cachingReg.register(name, schema1), id1);
 when(baseRegistry.register(name, schema2)).thenReturn(id2);
 Assert.assertEquals(cachingReg.register(name, schema2), id2);
 when(baseRegistry.register(name, schema3)).thenReturn(id3);
 try {
  cachingReg.register(name, schema3);
  Assert.fail("Should have thrown an exception");
 }
 catch (Exception e)
 {
  log.info(e.getMessage());
 }
}

@Override
public Converter<S, Schema, byte[], GenericRecord> init(WorkUnitState workUnit) {
 this.schemaRegistry = KafkaSchemaRegistryFactory.getSchemaRegistry(workUnit.getProperties());
 this.deserializer = new LiAvroDeserializerBase(this.schemaRegistry);
 return this;
}

private HttpClient borrowClient() throws SchemaRegistryException {
 try {
  return this.httpClientPool.borrowObject();
 } catch (Exception e) {
  throw new SchemaRegistryException("Unable to borrow " + HttpClient.class.getSimpleName());
 }
}

@Override
public Schema convertSchema(S schemaIn, WorkUnitState workUnit)
  throws SchemaConversionException {
 Preconditions.checkArgument(workUnit.contains(KafkaSource.TOPIC_NAME), "Must specify topic name.");
 String topic = workUnit.getProp(KafkaSource.TOPIC_NAME);
 try {
  return (Schema) this.schemaRegistry.getLatestSchema(topic);
 } catch (IOException | SchemaRegistryException e) {
  throw new SchemaConversionException(e);
 }
}

public ConfigDrivenMd5SchemaRegistry(Config config)
  throws IOException, SchemaRegistryException {
 if (config.hasPath(ConfigurationKeys.SCHEMA_NAME_KEY)) {
  String name = config.getString(ConfigurationKeys.SCHEMA_NAME_KEY);
  String value = config.getString(ConfigurationKeys.SCHEMA_VALUE_KEY);
  Schema schema = new Schema.Parser().parse(value);
  register(name, schema);
 }
}

/**
 * Get a schema given an id
 * @param id
 * @return
 * @throws IOException
 * @throws SchemaRegistryException
 */
@Override
public Schema getById(MD5Digest id)
  throws IOException, SchemaRegistryException {
  if (_schemaHashMap.containsKey(id))
  {
   return _schemaHashMap.get(id);
  }
  else
  {
   throw new SchemaRegistryException("Could not find schema with id : " + id.asString());
  }
 }