You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

110 lines
3.5 KiB

4 years ago
  1. #!/home/alpcentaur/ProjektA/PrototypeWebApp/venv/bin/python3.5
  2. import argparse
  3. import os
  4. import boto
  5. from boto.compat import json
  6. from boto.compat import six
  7. from boto.dynamodb.schema import Schema
  8. DESCRIPTION = """Load data into one or more DynamoDB tables.
  9. For each table, data is read from two files:
  10. - {table_name}.metadata for the table's name, schema and provisioned
  11. throughput (only required if creating the table).
  12. - {table_name}.data for the table's actual contents.
  13. Both files are searched for in the current directory. To read them from
  14. somewhere else, use the --in-dir parameter.
  15. This program does not wipe the tables prior to loading data. However, any
  16. items present in the data files will overwrite the table's contents.
  17. """
  18. def _json_iterload(fd):
  19. """Lazily load newline-separated JSON objects from a file-like object."""
  20. buffer = ""
  21. eof = False
  22. while not eof:
  23. try:
  24. # Add a line to the buffer
  25. buffer += fd.next()
  26. except StopIteration:
  27. # We can't let that exception bubble up, otherwise the last
  28. # object in the file will never be decoded.
  29. eof = True
  30. try:
  31. # Try to decode a JSON object.
  32. json_object = json.loads(buffer.strip())
  33. # Success: clear the buffer (everything was decoded).
  34. buffer = ""
  35. except ValueError:
  36. if eof and buffer.strip():
  37. # No more lines to load and the buffer contains something other
  38. # than whitespace: the file is, in fact, malformed.
  39. raise
  40. # We couldn't decode a complete JSON object: load more lines.
  41. continue
  42. yield json_object
  43. def create_table(metadata_fd):
  44. """Create a table from a metadata file-like object."""
  45. def load_table(table, in_fd):
  46. """Load items into a table from a file-like object."""
  47. for i in _json_iterload(in_fd):
  48. # Convert lists back to sets.
  49. data = {}
  50. for k, v in six.iteritems(i):
  51. if isinstance(v, list):
  52. data[k] = set(v)
  53. else:
  54. data[k] = v
  55. table.new_item(attrs=data).put()
  56. def dynamodb_load(tables, in_dir, create_tables):
  57. conn = boto.connect_dynamodb()
  58. for t in tables:
  59. metadata_file = os.path.join(in_dir, "%s.metadata" % t)
  60. data_file = os.path.join(in_dir, "%s.data" % t)
  61. if create_tables:
  62. with open(metadata_file) as meta_fd:
  63. metadata = json.load(meta_fd)
  64. table = conn.create_table(
  65. name=t,
  66. schema=Schema(metadata["schema"]),
  67. read_units=metadata["read_units"],
  68. write_units=metadata["write_units"],
  69. )
  70. table.refresh(wait_for_active=True)
  71. else:
  72. table = conn.get_table(t)
  73. with open(data_file) as in_fd:
  74. load_table(table, in_fd)
  75. if __name__ == "__main__":
  76. parser = argparse.ArgumentParser(
  77. prog="dynamodb_load",
  78. description=DESCRIPTION
  79. )
  80. parser.add_argument(
  81. "--create-tables",
  82. action="store_true",
  83. help="Create the tables if they don't exist already (without this flag, attempts to load data into non-existing tables fail)."
  84. )
  85. parser.add_argument("--in-dir", default=".")
  86. parser.add_argument("tables", metavar="TABLES", nargs="+")
  87. namespace = parser.parse_args()
  88. dynamodb_load(namespace.tables, namespace.in_dir, namespace.create_tables)