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