split-silo-database 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. #!/usr/bin/env python
  2. from typing import List
  3. import click
  4. import docker
  5. from django.apps import apps
  6. from sentry.runner import configure
  7. from sentry.silo.base import SiloMode
  8. configure()
  9. from django.conf import settings
  10. from sentry.models.organizationmapping import OrganizationMapping
  11. def exec_run(container, command):
  12. wrapped_command = f'sh -c "{" ".join(command)}"'
  13. exit_code, output = container.exec_run(cmd=wrapped_command, stdout=True, stderr=True)
  14. if exit_code:
  15. click.echo("Container operation Failed!")
  16. click.echo(f"Container operation failed with {output}")
  17. def split_database(tables: List[str], source: str, destination: str, reset: bool, verbose: bool):
  18. click.echo(f">> Dumping tables from {source} database")
  19. command = ["pg_dump", "-U", "postgres", "-d", source, "--clean"]
  20. for table in tables:
  21. command.extend(["-t", table])
  22. command.extend([">", f"/tmp/{destination}-tables.sql"])
  23. client = docker.from_env()
  24. postgres = client.containers.get("sentry_postgres")
  25. if verbose:
  26. click.echo(f">> Running {' '.join(command)}")
  27. exec_run(postgres, command)
  28. if reset:
  29. click.echo(f">> Dropping existing {destination} database")
  30. exec_run(postgres, ["dropdb", "-U", "postgres", "--if-exists", destination])
  31. exec_run(postgres, ["createdb", "-U", "postgres", destination])
  32. # Use the dump file to build control silo tables.
  33. click.echo(f">> Building {destination} database from dump file")
  34. import_command = ["psql", "-U", "postgres", destination, "<", f"/tmp/{destination}-tables.sql"]
  35. if verbose:
  36. click.echo(f">> Running {' '.join(import_command)}")
  37. exec_run(postgres, import_command)
  38. def revise_organization_mappings(legacy_region_name: str):
  39. if settings.SENTRY_MONOLITH_REGION == legacy_region_name:
  40. click.echo(
  41. "> No OrganizationMapping have been modified. Set 'SENTRY_MONOLITH_REGION' in sentry.conf.py to update monolith mappings."
  42. )
  43. else:
  44. qs = OrganizationMapping.objects.filter(region_name=legacy_region_name)
  45. record_count = len(qs)
  46. qs.update(region_name=settings.SENTRY_MONOLITH_REGION)
  47. click.echo(
  48. f"> {record_count} OrganizationMapping record(s) have been updated from '{legacy_region_name}' to '{settings.SENTRY_MONOLITH_REGION}'"
  49. )
  50. @click.command()
  51. @click.option(
  52. "--legacy-region-name",
  53. default="--monolith--",
  54. help="Previous value of settings.SENTRY_MONOLITH_REGION to overwrite in organization mappings",
  55. )
  56. @click.option("--verbose", default=False, is_flag=True, help="Enable verbose logging")
  57. @click.option(
  58. "--reset",
  59. default=False,
  60. is_flag=True,
  61. help="Reset the target databases to be empty before loading extracted data and schema.",
  62. )
  63. @click.option("--database", default="sentry", help="Which database to derive splits from")
  64. def main(database: str, reset: bool, verbose: bool, legacy_region_name: str):
  65. """
  66. This is a development tool that can convert a monolith database into
  67. control + region databases by using silo annotations.
  68. This operation will not modify the original source database.
  69. """
  70. # We have a few tables that either need to be in both silos,
  71. # or only in control. These tables don't have silo annotations
  72. # as they are inherited from django and their silo assignments
  73. # need to be manually defined.
  74. region_tables = ["django_migrations", "django_content_type"]
  75. control_tables = [
  76. "django_migrations",
  77. "django_admin_log",
  78. "django_content_type",
  79. "django_site",
  80. "django_session",
  81. "auth_user",
  82. "auth_group",
  83. "auth_permission",
  84. "auth_group_permissions",
  85. "auth_user_groups",
  86. "auth_user_user_permissions",
  87. ]
  88. for model in apps.get_models():
  89. silo_limit = getattr(model._meta, "silo_limit", None)
  90. if not silo_limit:
  91. click.echo(f"> Could not find silo assignment for {model._meta.db_table}")
  92. continue
  93. if SiloMode.CONTROL in silo_limit.modes:
  94. control_tables.append(model._meta.db_table)
  95. if SiloMode.REGION in silo_limit.modes:
  96. region_tables.append(model._meta.db_table)
  97. revise_organization_mappings(legacy_region_name=legacy_region_name)
  98. split_database(control_tables, database, "control", reset=reset, verbose=verbose)
  99. split_database(region_tables, database, "region", reset=reset, verbose=verbose)
  100. if __name__ == "__main__":
  101. main()