def test_schema_equivilence(load_avsc):
schema_str1 = load_avsc('basic_schema.avsc')
schema_str2 = load_avsc('basic_schema.avsc')
schema = Schema(schema_str1, 'AVRO')
schema2 = Schema(schema_str2, 'AVRO')
assert schema.__eq__(schema2)
assert schema == schema2
assert schema_str1.__eq__(schema_str2)
assert schema_str1 == schema_str2
def new_get_latest_version(subject_name: str) -> RegisteredSchema:
return RegisteredSchema(
schema_id="schema_id_1",
schema=Schema(schema_str=schema_str_ref, schema_type="AVRO"),
subject="test",
version=1,
)
def test_register_schema(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
schema = Schema(load_avsc('basic_schema.avsc'), schema_type='AVRO')
result = sr.register_schema('test-key', schema)
assert result == find_schema_id('test-key')
def test_get_schema(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
schema = Schema(load_avsc(mock_schema_registry.SCHEMA), schema_type='AVRO')
schema2 = sr.get_schema(47)
assert cmp_schema(schema, schema2)
def test_register_schema_incompatible(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
schema = Schema(load_avsc('basic_schema.avsc'), schema_type='AVRO')
with pytest.raises(SchemaRegistryError, match="Incompatible Schema") as e:
sr.register_schema('conflict', schema)
assert e.value.http_status_code == 409
assert e.value.error_code == -1
def test_register_schema_invalid(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
schema = Schema(load_avsc('invalid_schema.avsc'), schema_type='AVRO')
with pytest.raises(SchemaRegistryError, match="Invalid Schema") as e:
sr.register_schema('invalid', schema)
assert e.value.http_status_code == 422
assert e.value.error_code == 42201
def test_get_registration_schema_not_found(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
subject = 'schemanotfound'
schema = Schema(load_avsc(mock_schema_registry.SCHEMA), schema_type='AVRO')
with pytest.raises(SchemaRegistryError, match="Schema not found") as e:
sr.lookup_schema(subject, schema)
assert e.value.http_status_code == 404
assert e.value.error_code == 40403
def test_get_version(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
subject = "get_version"
version = 3
schema = Schema(load_avsc(mock_schema_registry.SCHEMA), schema_type='AVRO')
result = sr.get_version(subject, version)
assert result.subject == subject
assert result.version == version
assert cmp_schema(result.schema, schema)
assert result.schema_id == find_schema_id(subject)
def test_get_registration(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
subject = 'get_registration'
schema = Schema(load_avsc(mock_schema_registry.SCHEMA), schema_type='AVRO')
response = sr.lookup_schema(subject, schema)
assert response.subject == subject
assert response.version == mock_schema_registry.VERSION
assert response.schema_id == find_schema_id(subject)
assert cmp_schema(response.schema, schema)
def test_register_schema_cache(mock_schema_registry, load_avsc):
conf = {'url': TEST_URL}
sr = mock_schema_registry(conf)
schema = load_avsc('basic_schema.avsc')
count_before = sr.counter['POST'].get('/subjects/test-cache/versions', 0)
# Caching only starts after the first response is handled.
# A possible improvement would be to add request caching to the http client
# to catch in-flight requests as well.
sr.register_schema('test-cache', Schema(schema, 'AVRO'))
fs = []
with ThreadPoolExecutor(max_workers=10) as executor:
for _ in range(0, 1000):
fs.append(executor.submit(sr.register_schema, 'test-cache',
schema))
wait(fs)
count_after = sr.counter['POST'].get('/subjects/test-cache/versions')
assert count_after - count_before == 1
def test_get_schema_str_replace_confluent_ref_avro(self):
schema_str_orig = """
{
"fields": [
{
"name": "my_field1",
"type": "TestTopic1"
}
],
"name": "TestTopic1Val",
"namespace": "io.acryl",
"type": "record"
}
"""
schema_str_ref = """
{
"doc": "Sample schema to help you get started.",
"fields": [
{
"doc": "The int type is a 32-bit signed integer.",
"name": "my_field1",
"type": "int"
}
],
"name": "TestTopic1",
"namespace": "io.acryl",
"type": "record"
}
"""
schema_str_final = (
"""
{
"fields": [
{
"name": "my_field1",
"type": """
+ schema_str_ref
+ """
}
],
"name": "TestTopic1Val",
"namespace": "io.acryl",
"type": "record"
}
"""
)
kafka_source_config = KafkaSourceConfig.parse_obj(
{
"connection": {
"bootstrap": "localhost:9092",
"schema_registry_url": "http://localhost:8081",
},
}
)
confluent_schema_registry = ConfluentSchemaRegistry.create(
kafka_source_config, KafkaSourceReport()
)
def new_get_latest_version(subject_name: str) -> RegisteredSchema:
return RegisteredSchema(
schema_id="schema_id_1",
schema=Schema(schema_str=schema_str_ref, schema_type="AVRO"),
subject="test",
version=1,
)
with patch.object(
confluent_schema_registry.schema_registry_client,
"get_latest_version",
new_get_latest_version,
):
schema_str = confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
# The external reference would match by name.
schema=Schema(
schema_str=schema_str_orig,
schema_type="AVRO",
references=[
dict(name="TestTopic1", subject="schema_subject_1", version=1)
],
)
)
assert schema_str == ConfluentSchemaRegistry._compact_schema(
schema_str_final
)
with patch.object(
confluent_schema_registry.schema_registry_client,
"get_latest_version",
new_get_latest_version,
):
schema_str = confluent_schema_registry.get_schema_str_replace_confluent_ref_avro(
# The external reference would match by subject.
schema=Schema(
schema_str=schema_str_orig,
schema_type="AVRO",
references=[
dict(name="schema_subject_1", subject="TestTopic1", version=1)
],
)
)
assert schema_str == ConfluentSchemaRegistry._compact_schema(
schema_str_final
)
def test_kafka_source_workunits_schema_registry_subject_name_strategies(
self, mock_kafka_consumer, mock_schema_registry_client):
# Setup the topic to key/value schema mappings for all types of schema registry subject name strategies.
# <key=topic_name, value=(<key_schema>,<value_schema>)
topic_subject_schema_map: Dict[str, Tuple[RegisteredSchema, RegisteredSchema]] = {
# TopicNameStrategy is used for subject
"topic1": (
RegisteredSchema(
schema_id="schema_id_2",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic1Key", "namespace": "test.acryl", "fields": [{"name":"t1key", "type": "string"}]}',
schema_type="AVRO",
),
subject="topic1-key",
version=1,
),
RegisteredSchema(
schema_id="schema_id_1",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic1Value", "namespace": "test.acryl", "fields": [{"name":"t1value", "type": "string"}]}',
schema_type="AVRO",
),
subject=
"topic1-value",
version=
1,
),
),
# RecordNameStrategy is used for subject
"topic2":
(
RegisteredSchema(
schema_id="schema_id_3",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic2Key", "namespace": "test.acryl", "fields": [{"name":"t2key", "type": "string"}]}',
schema_type="AVRO",
),
subject="test.acryl.Topic2Key",
version=1,
),
RegisteredSchema(
schema_id="schema_id_4",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic2Value", "namespace": "test.acryl", "fields": [{"name":"t2value", "type": "string"}]}',
schema_type="AVRO",
),
subject="test.acryl.Topic2Value",
version=1,
),
),
# TopicRecordNameStrategy is used for subject
"topic3": (
RegisteredSchema(
schema_id="schema_id_4",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic3Key", "namespace": "test.acryl", "fields": [{"name":"t3key", "type": "string"}]}',
schema_type="AVRO",
),
subject="topic3-test.acryl.Topic3Key-key",
version=1,
),
RegisteredSchema(
schema_id="schema_id_5",
schema=Schema(
schema_str=
'{"type":"record", "name":"Topic3Value", "namespace": "test.acryl", "fields": [{"name":"t3value", "type": "string"}]}',
schema_type="AVRO",
),
subject="topic3-test.acryl.Topic3Value-value",
version=1,
),
),
}
# Mock the kafka consumer
mock_kafka_instance = mock_kafka_consumer.return_value
mock_cluster_metadata = MagicMock()
mock_cluster_metadata.topics = list(topic_subject_schema_map.keys())
mock_cluster_metadata.topics.append("schema_less_topic")
mock_kafka_instance.list_topics.return_value = mock_cluster_metadata
# Mock the schema registry client
# - mock get_subjects: all subjects in topic_subject_schema_map
mock_schema_registry_client.return_value.get_subjects.return_value = [
v.subject for v in chain(*topic_subject_schema_map.values())
]
# - mock get_latest_version
def mock_get_latest_version(
subject_name: str) -> Optional[RegisteredSchema]:
for registered_schema in chain(*topic_subject_schema_map.values()):
if registered_schema.subject == subject_name:
return registered_schema
return None
mock_schema_registry_client.return_value.get_latest_version = (
mock_get_latest_version)
# Test the kafka source
source_config = {
"connection": {
"bootstrap": "localhost:9092"
},
# Setup the topic_subject_map for topic2 which uses RecordNameStrategy
"topic_subject_map": {
"topic2-key": "test.acryl.Topic2Key",
"topic2-value": "test.acryl.Topic2Value",
},
}
ctx = PipelineContext(run_id="test")
kafka_source = KafkaSource.create(source_config, ctx)
workunits = list(kafka_source.get_workunits())
mock_kafka_consumer.assert_called_once()
mock_kafka_instance.list_topics.assert_called_once()
assert len(workunits) == 8
i: int = -1
for wu in workunits:
assert isinstance(wu, MetadataWorkUnit)
if not isinstance(wu.metadata, MetadataChangeEvent):
continue
mce: MetadataChangeEvent = wu.metadata
i += 1
if i < len(topic_subject_schema_map.keys()):
# First 3 workunits (topics) must have schemaMetadata aspect
assert isinstance(mce.proposedSnapshot.aspects[1],
SchemaMetadataClass)
schemaMetadataAspect: SchemaMetadataClass = (
mce.proposedSnapshot.aspects[1])
assert isinstance(schemaMetadataAspect.platformSchema,
KafkaSchemaClass)
# Make sure the schema name is present in topic_subject_schema_map.
assert schemaMetadataAspect.schemaName in topic_subject_schema_map
# Make sure the schema_str matches for the key schema.
assert (
schemaMetadataAspect.platformSchema.keySchema ==
topic_subject_schema_map[
schemaMetadataAspect.schemaName][0].schema.schema_str)
# Make sure the schema_str matches for the value schema.
assert (
schemaMetadataAspect.platformSchema.documentSchema ==
topic_subject_schema_map[
schemaMetadataAspect.schemaName][1].schema.schema_str)
# Make sure we have 2 fields, one from the key schema & one from the value schema.
assert len(schemaMetadataAspect.fields) == 2
else:
# Last topic('schema_less_topic') has no schema defined in the registry.
# The schemaMetadata aspect should not be present for this.
for aspect in mce.proposedSnapshot.aspects:
assert not isinstance(aspect, SchemaMetadataClass)
def test_get_schema_str_replace_confluent_ref_avro(self):
# References external schema 'TestTopic1' in the definition of 'my_field1' field.
schema_str_orig = """
{
"fields": [
{
"name": "my_field1",
"type": "TestTopic1"
}
],
"name": "TestTopic1Val",
"namespace": "io.acryl",
"type": "record"
}
"""
schema_str_ref = """
{
"doc": "Sample schema to help you get started.",
"fields": [
{
"doc": "The int type is a 32-bit signed integer.",
"name": "my_field1",
"type": "int"
}
],
"name": "TestTopic1",
"namespace": "io.acryl",
"type": "record"
}
"""
schema_str_final = ("""
{
"fields": [
{
"name": "my_field1",
"type": """ + schema_str_ref + """
}
],
"name": "TestTopic1Val",
"namespace": "io.acryl",
"type": "record"
}
""")
ctx = PipelineContext(run_id="test")
kafka_source = KafkaSource.create(
{
"connection": {
"bootstrap": "localhost:9092"
},
},
ctx,
)
def new_get_latest_version(subject_name: str) -> RegisteredSchema:
return RegisteredSchema(
schema_id="schema_id_1",
schema=Schema(schema_str=schema_str_ref, schema_type="AVRO"),
subject="test",
version=1,
)
with patch.object(
kafka_source.schema_registry_client,
"get_latest_version",
new_get_latest_version,
):
schema_str = kafka_source.get_schema_str_replace_confluent_ref_avro(
# The external reference would match by name.
schema=Schema(
schema_str=schema_str_orig,
schema_type="AVRO",
references=[
dict(name="TestTopic1",
subject="schema_subject_1",
version=1)
],
))
assert schema_str == KafkaSource._compact_schema(schema_str_final)
with patch.object(
kafka_source.schema_registry_client,
"get_latest_version",
new_get_latest_version,
):
schema_str = kafka_source.get_schema_str_replace_confluent_ref_avro(
# The external reference would match by subject.
schema=Schema(
schema_str=schema_str_orig,
schema_type="AVRO",
references=[
dict(name="schema_subject_1",
subject="TestTopic1",
version=1)
],
))
assert schema_str == KafkaSource._compact_schema(schema_str_final)