1#!/usr/bin/env python 2# 3# Copyright 2016 Confluent Inc. 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16# 17 18 19# 20# derived from https://github.com/verisign/python-confluent-schemaregistry.git 21# 22import os 23import os.path 24import random 25 26from avro import schema 27from avro.datafile import DataFileWriter 28from avro.io import DatumWriter 29 30NAMES = ['stefan', 'melanie', 'nick', 'darrel', 'kent', 'simon'] 31AGES = list(range(1, 10)) + [None] 32 33 34def get_schema_path(fname): 35 dname = os.path.dirname(os.path.realpath(__file__)) 36 return os.path.join(dname, fname) 37 38 39def load_schema_file(fname): 40 fname = get_schema_path(fname) 41 with open(fname) as f: 42 return f.read() 43 44 45avsc_dir = os.path.dirname(os.path.realpath(__file__)) 46 47BASIC_SCHEMA = load_schema_file(os.path.join(avsc_dir, 'basic_schema.avsc')) 48 49 50def create_basic_item(i): 51 return { 52 'name': random.choice(NAMES) + '-' + str(i), 53 'number': random.choice(AGES) 54 } 55 56 57BASIC_ITEMS = map(create_basic_item, range(1, 20)) 58 59ADVANCED_SCHEMA = load_schema_file(os.path.join(avsc_dir, 'adv_schema.avsc')) 60 61 62def create_adv_item(i): 63 friends = map(create_basic_item, range(1, 3)) 64 family = map(create_basic_item, range(1, 3)) 65 basic = create_basic_item(i) 66 basic['family'] = dict(map(lambda bi: (bi['name'], bi), family)) 67 basic['friends'] = dict(map(lambda bi: (bi['name'], bi), friends)) 68 return basic 69 70 71ADVANCED_ITEMS = map(create_adv_item, range(1, 20)) 72 73 74def _write_items(base_name, schema_str, items): 75 avro_schema = schema.Parse(schema_str) 76 avro_file = base_name + '.avro' 77 with DataFileWriter(open(avro_file, "w"), DatumWriter(), avro_schema) as writer: 78 for i in items: 79 writer.append(i) 80 writer.close 81 return (avro_file) 82 83 84def write_basic_items(base_name): 85 return _write_items(base_name, BASIC_SCHEMA, BASIC_ITEMS) 86 87 88def write_advanced_items(base_name): 89 return _write_items(base_name, ADVANCED_SCHEMA, ADVANCED_ITEMS) 90 91 92def cleanup(files): 93 for f in files: 94 try: 95 os.remove(f) 96 except OSError: 97 pass 98 99 100if __name__ == "__main__": 101 write_advanced_items("advanced") 102