/src/import-sql/generate_sql.py

https://github.com/osm2vectortiles/osm2vectortiles · Python · 128 lines · 87 code · 30 blank · 11 comment · 14 complexity · 4bad715fbe7f6840304edd3610982673 MD5 · raw file

  1. #!/usr/bin/env python
  2. """Generate SQL functions from custom YAML definitions
  3. Usage:
  4. generate_sql.py class <yaml-source>
  5. generate_sql.py changed_tiles <yaml-source>
  6. generate_sql.py tables <yaml-source>
  7. generate_sql.py (-h | --help)
  8. Options:
  9. -h --help Show this screen.
  10. --version Show version.
  11. """
  12. from collections import namedtuple
  13. from docopt import docopt
  14. import yaml
  15. SQL_INDENT = 8 * " "
  16. Class = namedtuple('Class', ['name', 'values'])
  17. def generate_sql_class(source, func_suffix='class'):
  18. def gen_when_stmt(class_name, mapping_values):
  19. in_statements = ["'{}'".format(value) for value in mapping_values]
  20. return " " * 12 + "WHEN type IN ({0}) THEN '{1}'".format(
  21. ','.join(in_statements),
  22. class_name
  23. )
  24. system_name = source['system']['name']
  25. classes = find_classes(source)
  26. when_statements = [gen_when_stmt(cl, val) for cl, val in classes]
  27. return """CREATE OR REPLACE FUNCTION {0}_{1}(type VARCHAR)
  28. RETURNS VARCHAR AS $$
  29. BEGIN
  30. RETURN CASE
  31. {2}
  32. END;
  33. END;
  34. $$ LANGUAGE plpgsql IMMUTABLE;
  35. """.format(system_name, func_suffix, "\n".join(when_statements))
  36. def find_classes(config):
  37. for cl_name, mapped_values in config['system']['classes'].items():
  38. yield Class(cl_name, mapped_values)
  39. Table = namedtuple('Table', ['name', 'buffer', 'min_zoom', 'max_zoom'])
  40. def generate_changed_tiles(
  41. tables,
  42. func_name='changed_tiles',
  43. func_changed_tiles_query='changed_tiles_table'
  44. ):
  45. def gen_select_stmt(table):
  46. return "SELECT * FROM {0}('{1}', ts, {2}, {3}, {4})".format(
  47. func_changed_tiles_query,
  48. table.name,
  49. table.buffer,
  50. table.min_zoom,
  51. table.max_zoom,
  52. )
  53. stmts = [2 * SQL_INDENT + gen_select_stmt(t) for t in tables]
  54. separator = '\n' + 2 * SQL_INDENT + 'UNION\n'
  55. return """CREATE OR REPLACE FUNCTION {0}(ts timestamp)
  56. RETURNS TABLE (x INTEGER, y INTEGER, z INTEGER) AS $$
  57. BEGIN
  58. RETURN QUERY (
  59. {1}
  60. );
  61. END;
  62. $$ language plpgsql;
  63. """.format(func_name, separator.join(stmts))
  64. def generate_static_table_view(tables, view_name='osm_tables'):
  65. def gen_select_stmt(table):
  66. return ("SELECT '{0}' AS table_name,"
  67. '{1} AS buffer_size,'
  68. '{2} AS min_zoom,'
  69. '{3} AS max_zoom').format(table.name, table.buffer,
  70. table.min_zoom, table.max_zoom)
  71. stmts = [SQL_INDENT + gen_select_stmt(t) for t in tables]
  72. sep = '\n' + SQL_INDENT + 'UNION\n'
  73. return 'CREATE OR REPLACE VIEW {0} AS (\n{1}\n);'.format(view_name,
  74. sep.join(stmts))
  75. def find_delete_tables(config, delete_suffix='delete'):
  76. for src_table in find_tables(config):
  77. yield Table(src_table.name + '_' + delete_suffix,
  78. src_table.buffer,
  79. src_table.min_zoom, src_table.max_zoom)
  80. def find_tables_with_deletes(config):
  81. for src_table in find_tables(config):
  82. yield src_table
  83. for delete_table in find_delete_tables(config):
  84. yield delete_table
  85. def find_tables(config, schema_prefix='osm'):
  86. for table_name, config_values in config['tables'].items():
  87. yield Table(schema_prefix + '_' + table_name,config_values['buffer'],
  88. config_values['min_zoom'], config_values['max_zoom'])
  89. if __name__ == '__main__':
  90. args = docopt(__doc__)
  91. with open(args['<yaml-source>'], 'r') as f:
  92. source = yaml.load(f)
  93. if args['class']:
  94. print(generate_sql_class(source))
  95. if args['changed_tiles']:
  96. print(generate_changed_tiles(find_tables_with_deletes(source)))
  97. if args['tables']:
  98. print(generate_static_table_view(find_tables(source)))