def _main() -> None:
"""
This method takes the output of `git diff --name-status master snuba/migrations` and
runs `snuba migrations run -dry-run with the proper parameters`, for a CI action
"""
diff_result = subprocess.run(
[
"git",
"diff",
"--diff-filter=AM",
"--name-only",
"origin/master",
"--",
"snuba/migrations/snuba_migrations/*/[0-9]*.py",
],
stdout=subprocess.PIPE,
text=True,
)
if diff_result.returncode != 0:
raise ExecError(diff_result.stdout)
else:
lines = diff_result.stdout.splitlines()
if len(lines) > 0:
print("-- start migrations")
print()
for line in lines:
migration_filename = os.path.basename(line)
migration_group = MigrationGroup(os.path.basename(os.path.dirname(line)))
migration_id, _ = os.path.splitext(migration_filename)
runner = Runner()
migration_key = MigrationKey(migration_group, migration_id)
print(f"-- migration {migration_group.value} : {migration_id}")
runner.run_migration(migration_key, dry_run=True)
print(f"-- end migration {migration_group.value} : {migration_id}")
def run_prior_migrations(
migration_group: MigrationGroup, stop_migration_id: str, runner: Runner
) -> None:
"""Runs all migrations up to the migration denoted by migration ID
Arguments:
migration_group -- the group of the desired migration
stop_migration_id -- desired migration ID, as a stopping point
runner -- migration runner object
"""
right_migrations = next(
group_migrations
for (group, group_migrations) in runner.show_all()
if group == migration_group
)
# Run migrations up to the desired 'stop' ID
for migration in right_migrations:
if migration.migration_id == stop_migration_id:
break
runner.run_migration(
MigrationKey(migration_group, migration.migration_id), force=True
)
def list() -> None:
"""
Lists migrations and their statuses
"""
check_clickhouse_connections()
runner = Runner()
for group, group_migrations in runner.show_all():
click.echo(group.value)
for migration_id, status, blocking in group_migrations:
symbol = {
Status.COMPLETED: "X",
Status.NOT_STARTED: " ",
Status.IN_PROGRESS: "-",
}[status]
in_progress_text = " (IN PROGRESS)" if status == Status.IN_PROGRESS else ""
blocking_text = ""
if status != Status.COMPLETED and blocking:
blocking_text = " (blocking)"
click.echo(
f"[{symbol}] {migration_id}{in_progress_text}{blocking_text}")
click.echo()
def test_run_migration() -> None:
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
connection = get_cluster(StorageSetKey.MIGRATIONS).get_query_connection(
ClickhouseClientSettings.MIGRATE)
assert connection.execute(
"SELECT group, migration_id, status, version FROM migrations_local;"
) == [("system", "0001_migrations", "completed", 1)]
# Invalid migration ID
with pytest.raises(MigrationError):
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM, "xxx"))
# Run out of order
with pytest.raises(MigrationError):
runner.run_migration(MigrationKey(MigrationGroup.EVENTS,
"0003_errors"))
# Running with --fake
runner.run_migration(MigrationKey(MigrationGroup.EVENTS,
"0001_events_initial"),
fake=True)
assert connection.execute("SHOW TABLES LIKE 'sentry_local'") == []
def migrate(force: bool) -> None:
"""
Runs all migrations. Blocking migrations will not be run unless --force is passed.
"""
runner = Runner()
try:
runner.run_all(force=force)
except MigrationError as e:
raise click.ClickException(str(e))
click.echo("Finished running migrations")
def migrate(force: bool, log_level: Optional[str] = None) -> None:
"""
Runs all migrations. Blocking migrations will not be run unless --force is passed.
"""
setup_logging(log_level)
check_clickhouse_connections()
runner = Runner()
try:
runner.run_all(force=force)
except MigrationError as e:
raise click.ClickException(str(e))
click.echo("Finished running migrations")
def run_migrations() -> Iterator[None]:
from snuba.migrations.runner import Runner
Runner().run_all(force=True)
yield
for storage_key in STORAGES:
storage = get_storage(storage_key)
cluster = storage.get_cluster()
database = cluster.get_database()
schema = storage.get_schema()
if isinstance(schema, WritableTableSchema):
table_name = schema.get_local_table_name()
nodes = [
*cluster.get_local_nodes(), *cluster.get_distributed_nodes()
]
for node in nodes:
connection = cluster.get_node_connection(
ClickhouseClientSettings.MIGRATE, node)
connection.execute(
f"TRUNCATE TABLE IF EXISTS {database}.{table_name}")
redis_client.flushdb()
def test_run_all() -> None:
runner = Runner()
assert len(runner._get_pending_migrations()) == get_total_migration_count()
with pytest.raises(MigrationError):
runner.run_all(force=False)
runner.run_all(force=True)
assert runner._get_pending_migrations() == []
def add_node(
node_type: str,
storage_set_names: Sequence[str],
host_name: str,
port: int,
database: str,
) -> None:
"""
Runs all migrations on a brand new ClickHouse node. This should be performed
before a new node is added to an existing ClickHouse cluster.
All of the SQL operations for the provided storage sets will be run. Any non
SQL (Python) operations will be skipped.
This operation does not change the migration status in the migrations_local
/ migrations_dist tables, since it is designed to bring a new node up to
the same state as existing ones already added to the cluster.
"""
user = os.environ.get("CLICKHOUSE_USER", "default")
password = os.environ.get("CLICKHOUSE_PASSWORD", "")
storage_set_keys = [StorageSetKey(name) for name in storage_set_names]
cluster = next(
(
c
for c in CLUSTERS
if all(ss in c.get_storage_set_keys() for ss in storage_set_keys)
),
None,
)
if not cluster:
raise click.ClickException("Storage sets should be in the same cluster")
if cluster.is_single_node():
raise click.ClickException("You cannot add a node to a single node cluster")
Runner.add_node(
node_type=ClickhouseNodeType(node_type),
storage_sets=storage_set_keys,
host_name=host_name,
port=port,
user=user,
password=password,
database=database,
)
def test_transactions_compatibility() -> None:
cluster = get_cluster(StorageSetKey.TRANSACTIONS)
connection = cluster.get_query_connection(ClickhouseClientSettings.MIGRATE)
def get_sampling_key() -> str:
database = cluster.get_database()
((sampling_key, ), ) = connection.execute(
f"SELECT sampling_key FROM system.tables WHERE name = 'transactions_local' AND database = '{database}'"
)
return sampling_key
# Create old style table without sampling expression and insert data
connection.execute("""
CREATE TABLE transactions_local (`project_id` UInt64, `event_id` UUID,
`trace_id` UUID, `span_id` UInt64, `transaction_name` LowCardinality(String),
`transaction_hash` UInt64 MATERIALIZED CAST(cityHash64(transaction_name), 'UInt64'),
`transaction_op` LowCardinality(String), `transaction_status` UInt8 DEFAULT 2,
`start_ts` DateTime, `start_ms` UInt16, `finish_ts` DateTime, `finish_ms` UInt16,
`duration` UInt32, `platform` LowCardinality(String), `environment` LowCardinality(Nullable(String)),
`release` LowCardinality(Nullable(String)), `dist` LowCardinality(Nullable(String)),
`ip_address_v4` Nullable(IPv4), `ip_address_v6` Nullable(IPv6), `user` String DEFAULT '',
`user_hash` UInt64 MATERIALIZED cityHash64(user), `user_id` Nullable(String),
`user_name` Nullable(String), `user_email` Nullable(String),
`sdk_name` LowCardinality(String) DEFAULT CAST('', 'LowCardinality(String)'),
`sdk_version` LowCardinality(String) DEFAULT CAST('', 'LowCardinality(String)'),
`http_method` LowCardinality(Nullable(String)) DEFAULT CAST('', 'LowCardinality(Nullable(String))'),
`http_referer` Nullable(String),
`tags.key` Array(String), `tags.value` Array(String), `_tags_flattened` String,
`contexts.key` Array(String), `contexts.value` Array(String), `_contexts_flattened` String,
`partition` UInt16, `offset` UInt64, `message_timestamp` DateTime, `retention_days` UInt16,
`deleted` UInt8) ENGINE = ReplacingMergeTree(deleted) PARTITION BY (retention_days, toMonday(finish_ts))
ORDER BY (project_id, toStartOfDay(finish_ts), transaction_name, cityHash64(span_id))
TTL finish_ts + toIntervalDay(retention_days);
""")
assert get_sampling_key() == ""
generate_transactions()
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
runner._update_migration_status(
MigrationKey(MigrationGroup.TRANSACTIONS, "0001_transactions"),
Status.COMPLETED)
runner.run_migration(
MigrationKey(
MigrationGroup.TRANSACTIONS,
"0002_transactions_onpremise_fix_orderby_and_partitionby",
),
force=True,
)
assert get_sampling_key() == "cityHash64(span_id)"
assert connection.execute("SELECT count(*) FROM transactions_local;") == [
(5, )
]
def test_no_schema_differences() -> None:
runner = Runner()
runner.run_all(force=True)
for storage_key in StorageKey:
storage = get_storage(storage_key)
conn = storage.get_cluster().get_query_connection(
ClickhouseClientSettings.MIGRATE)
schema = storage.get_schema()
if not isinstance(schema, TableSchema):
continue
table_name = schema.get_local_table_name()
local_schema = get_local_schema(conn, table_name)
assert (schema.get_column_differences(local_schema) == []
), f"Schema mismatch: {table_name} does not match schema"
def test_get_status() -> None:
runner = Runner()
assert runner.get_status(
MigrationKey(MigrationGroup.EVENTS,
"0001_events_initial")) == (Status.NOT_STARTED, None)
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
assert runner.get_status(
MigrationKey(MigrationGroup.EVENTS,
"0001_events_initial")) == (Status.NOT_STARTED, None)
runner.run_migration(
MigrationKey(MigrationGroup.EVENTS, "0001_events_initial"))
status = runner.get_status(
MigrationKey(MigrationGroup.EVENTS, "0001_events_initial"))
assert status[0] == Status.COMPLETED
assert isinstance(status[1], datetime)
def test_groupedmessages_compatibility() -> None:
cluster = get_cluster(StorageSetKey.EVENTS)
# Ignore the multi node mode because this tests a migration
# for an older table state that only applied to single node
if not cluster.is_single_node():
return
database = cluster.get_database()
connection = cluster.get_query_connection(ClickhouseClientSettings.MIGRATE)
# Create old style table witihout project ID
connection.execute("""
CREATE TABLE groupedmessage_local (`offset` UInt64, `record_deleted` UInt8,
`id` UInt64, `status` Nullable(UInt8), `last_seen` Nullable(DateTime),
`first_seen` Nullable(DateTime), `active_at` Nullable(DateTime),
`first_release_id` Nullable(UInt64)) ENGINE = ReplacingMergeTree(offset)
ORDER BY id SAMPLE BY id SETTINGS index_granularity = 8192
""")
migration_id = "0010_groupedmessages_onpremise_compatibility"
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
events_migrations = get_group_loader(
MigrationGroup.EVENTS).get_migrations()
# Mark prior migrations complete
for migration in events_migrations[:(
events_migrations.index(migration_id))]:
runner._update_migration_status(
MigrationKey(MigrationGroup.EVENTS, migration), Status.COMPLETED)
runner.run_migration(
MigrationKey(MigrationGroup.EVENTS, migration_id),
force=True,
)
outcome = perform_select_query(
["primary_key"],
"system.tables",
{
"name": "groupedmessage_local",
"database": str(database)
},
None,
connection,
)
assert outcome == [("project_id, id", )]
def dataset_manager(name: str) -> Iterator[Dataset]:
from snuba.migrations.runner import Runner
from snuba.web.views import truncate_dataset
Runner().run_all(force=True)
dataset = get_dataset(name)
truncate_dataset(dataset)
try:
yield dataset
finally:
truncate_dataset(dataset)
def reverse(group: str, migration_id: str, force: bool, fake: bool,
dry_run: bool) -> None:
"""
Reverses a single migration.
--force is required to reverse an already completed migration.
--fake marks a migration as reversed without doing anything.
"""
if not dry_run:
check_clickhouse_connections()
runner = Runner()
migration_group = MigrationGroup(group)
migration_key = MigrationKey(migration_group, migration_id)
if dry_run:
runner.reverse_migration(migration_key, dry_run=True)
return
try:
if fake:
click.confirm(
"This will mark the migration as not started without actually reversing it. Your database may be in an invalid state. Are you sure?",
abort=True,
)
runner.reverse_migration(migration_key, force=force, fake=fake)
except MigrationError as e:
raise click.ClickException(str(e))
click.echo(f"Finished reversing migration {migration_key}")
def run(group: str, migration_id: str, force: bool, fake: bool,
dry_run: bool) -> None:
"""
Runs a single migration.
--force must be passed in order to run blocking migrations.
--fake marks a migration as completed without running anything.
Migrations that are already in an in-progress or completed status will not be run.
"""
if not dry_run:
check_clickhouse_connections()
runner = Runner()
migration_group = MigrationGroup(group)
migration_key = MigrationKey(migration_group, migration_id)
if dry_run:
runner.run_migration(migration_key, dry_run=True)
return
try:
if fake:
click.confirm(
"This will mark the migration as completed without actually running it. Your database may be in an invalid state. Are you sure?",
abort=True,
)
runner.run_migration(migration_key, force=force, fake=fake)
except MigrationError as e:
raise click.ClickException(str(e))
click.echo(f"Finished running migration {migration_key}")
def test_get_pending_migrations() -> None:
runner = Runner()
total_migrations = get_total_migration_count()
assert len(runner._get_pending_migrations()) == total_migrations
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
assert len(runner._get_pending_migrations()) == total_migrations - 1
def test_reverse_all() -> None:
runner = Runner()
all_migrations = runner._get_pending_migrations()
runner.run_all(force=True)
for migration in reversed(all_migrations):
runner.reverse_migration(migration, force=True)
connection = get_cluster(StorageSetKey.MIGRATIONS).get_query_connection(
ClickhouseClientSettings.MIGRATE)
assert connection.execute(
"SHOW TABLES") == [], "All tables should be deleted"
def test_reverse_migration() -> None:
runner = Runner()
runner.run_all(force=True)
connection = get_cluster(StorageSetKey.MIGRATIONS).get_query_connection(
ClickhouseClientSettings.MIGRATE)
# Invalid migration ID
with pytest.raises(MigrationError):
runner.reverse_migration(MigrationKey(MigrationGroup.SYSTEM, "xxx"))
with pytest.raises(MigrationError):
runner.reverse_migration(
MigrationKey(MigrationGroup.EVENTS, "0003_errors"))
# Reverse with --fake
for migration_id in reversed(
get_group_loader(MigrationGroup.EVENTS).get_migrations()):
runner.reverse_migration(MigrationKey(MigrationGroup.EVENTS,
migration_id),
fake=True)
assert (len(connection.execute("SHOW TABLES LIKE 'sentry_local'")) == 1
), "Table still exists"
def test_no_schema_differences() -> None:
settings.ENABLE_DEV_FEATURES = True
importlib.reload(factory)
runner = Runner()
runner.run_all(force=True)
for storage_key in STORAGES:
storage = get_storage(storage_key)
conn = storage.get_cluster().get_query_connection(
ClickhouseClientSettings.MIGRATE)
schema = storage.get_schema()
if not isinstance(schema, TableSchema):
continue
table_name = schema.get_local_table_name()
local_schema = get_local_schema(conn, table_name)
assert (schema.get_column_differences(local_schema) == []
), f"Schema mismatch: {table_name} does not match schema"
importlib.reload(settings)
importlib.reload(factory)
def test_show_all() -> None:
runner = Runner()
assert all([
migration.status == Status.NOT_STARTED
for (_, group_migrations) in runner.show_all()
for migration in group_migrations
])
runner.run_all(force=True)
assert all([
migration.status == Status.COMPLETED
for (_, group_migrations) in runner.show_all()
for migration in group_migrations
])
def test_groupedmessages_compatibility() -> None:
cluster = get_cluster(StorageSetKey.EVENTS)
database = cluster.get_database()
connection = cluster.get_query_connection(ClickhouseClientSettings.MIGRATE)
# Create old style table witihout project ID
connection.execute("""
CREATE TABLE groupedmessage_local (`offset` UInt64, `record_deleted` UInt8,
`id` UInt64, `status` Nullable(UInt8), `last_seen` Nullable(DateTime),
`first_seen` Nullable(DateTime), `active_at` Nullable(DateTime),
`first_release_id` Nullable(UInt64)) ENGINE = ReplacingMergeTree(offset)
ORDER BY id SAMPLE BY id SETTINGS index_granularity = 8192
""")
migration_id = "0010_groupedmessages_onpremise_compatibility"
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
events_migrations = get_group_loader(
MigrationGroup.EVENTS).get_migrations()
# Mark prior migrations complete
for migration in events_migrations[:(
events_migrations.index(migration_id))]:
runner._update_migration_status(
MigrationKey(MigrationGroup.EVENTS, migration), Status.COMPLETED)
runner.run_migration(
MigrationKey(MigrationGroup.EVENTS, migration_id),
force=True,
)
assert connection.execute(
f"SELECT primary_key FROM system.tables WHERE name = 'groupedmessage_local' AND database = '{database}'"
) == [("project_id, id", )]
def test_backfill_errors() -> None:
backfill_migration_id = "0014_backfill_errors"
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
run_prior_migrations(MigrationGroup.EVENTS, backfill_migration_id, runner)
errors_storage = get_writable_storage(StorageKey.ERRORS)
clickhouse = errors_storage.get_cluster().get_query_connection(
ClickhouseClientSettings.QUERY)
errors_table_name = errors_storage.get_table_writer().get_schema(
).get_table_name()
raw_events = []
for i in range(10):
event = get_raw_event()
raw_events.append(event)
events_storage = get_writable_storage(StorageKey.EVENTS)
write_unprocessed_events(events_storage, raw_events)
assert get_count_from_storage(errors_table_name, clickhouse) == 0
# Run 0014_backfill_errors
runner.run_migration(MigrationKey(MigrationGroup.EVENTS,
backfill_migration_id),
force=True)
assert get_count_from_storage(errors_table_name, clickhouse) == 10
outcome = perform_select_query(["contexts.key", "contexts.value"],
errors_table_name, None, str(1), clickhouse)
assert outcome[0] == (
[
"device.model_id",
"geo.city",
"geo.country_code",
"geo.region",
"os.kernel_version",
],
["Galaxy", "San Francisco", "US", "CA", "1.1.1"],
)
def bootstrap(
*,
bootstrap_server: Sequence[str],
kafka: bool,
migrate: bool,
force: bool,
log_level: Optional[str] = None,
) -> None:
"""
Warning: Not intended to be used in production yet.
"""
if not force:
raise click.ClickException("Must use --force to run")
setup_logging(log_level)
logger = logging.getLogger("snuba.bootstrap")
import time
if kafka:
logger.debug("Using Kafka with %r", bootstrap_server)
from confluent_kafka.admin import AdminClient, NewTopic
override_params = {
# Same as above: override socket timeout as we expect Kafka
# to not getting ready for a while
"socket.timeout.ms": 1000,
}
if logger.getEffectiveLevel() != logging.DEBUG:
# Override rdkafka loglevel to be critical unless we are
# debugging as we expect failures when trying to connect
# (Kafka may not be up yet)
override_params["log_level"] = LOG_CRIT
attempts = 0
while True:
try:
logger.info("Attempting to connect to Kafka (attempt %d)...", attempts)
client = AdminClient(
get_default_kafka_configuration(
bootstrap_servers=bootstrap_server,
override_params=override_params,
)
)
client.list_topics(timeout=1)
break
except KafkaException as err:
logger.debug(
"Connection to Kafka failed (attempt %d)", attempts, exc_info=err
)
attempts += 1
if attempts == 60:
raise
time.sleep(1)
logger.info("Connected to Kafka on attempt %d", attempts)
topics = {}
for topic in Topic:
topic_spec = KafkaTopicSpec(topic)
logger.debug("Adding topic %s to creation list", topic_spec.topic_name)
topics[topic_spec.topic_name] = NewTopic(
topic_spec.topic_name,
num_partitions=topic_spec.partitions_number,
replication_factor=topic_spec.replication_factor,
config=topic_spec.topic_creation_config,
)
logger.info("Creating Kafka topics...")
for topic, future in client.create_topics(
list(topics.values()), operation_timeout=1
).items():
try:
future.result()
logger.info("Topic %s created", topic)
except KafkaException as err:
if err.args[0].code() != KafkaError.TOPIC_ALREADY_EXISTS:
logger.error("Failed to create topic %s", topic, exc_info=err)
if migrate:
check_clickhouse_connections()
Runner().run_all(force=True)
def test_version() -> None:
runner = Runner()
runner.run_migration(MigrationKey(MigrationGroup.SYSTEM,
"0001_migrations"))
migration_key = MigrationKey(MigrationGroup.EVENTS, "test")
assert runner._get_next_version(migration_key) == 1
runner._update_migration_status(migration_key, Status.IN_PROGRESS)
assert runner._get_next_version(migration_key) == 2
runner._update_migration_status(migration_key, Status.COMPLETED)
assert runner._get_next_version(migration_key) == 3
def bootstrap(
*,
bootstrap_server: Sequence[str],
kafka: bool,
migrate: bool,
force: bool,
log_level: Optional[str] = None,
) -> None:
"""
Warning: Not intended to be used in production yet.
"""
if not force:
raise click.ClickException("Must use --force to run")
setup_logging(log_level)
logger = logging.getLogger("snuba.bootstrap")
import time
if kafka:
logger.debug("Using Kafka with %r", bootstrap_server)
from confluent_kafka.admin import AdminClient
override_params = {
# Same as above: override socket timeout as we expect Kafka
# to not getting ready for a while
"socket.timeout.ms": 1000,
}
if logger.getEffectiveLevel() != logging.DEBUG:
# Override rdkafka loglevel to be critical unless we are
# debugging as we expect failures when trying to connect
# (Kafka may not be up yet)
override_params["log_level"] = LOG_CRIT
attempts = 0
while True:
try:
logger.info("Attempting to connect to Kafka (attempt %d)...",
attempts)
client = AdminClient(
get_default_kafka_configuration(
bootstrap_servers=bootstrap_server,
override_params=override_params,
))
client.list_topics(timeout=1)
break
except KafkaException as err:
logger.debug("Connection to Kafka failed (attempt %d)",
attempts,
exc_info=err)
attempts += 1
if attempts == 60:
raise
time.sleep(1)
logger.info("Connected to Kafka on attempt %d", attempts)
create_topics(client, [t for t in Topic])
if migrate:
check_clickhouse_connections()
Runner().run_all(force=True)
def bootstrap(
*,
bootstrap_server: Sequence[str],
kafka: bool,
migrate: bool,
force: bool,
log_level: Optional[str] = None,
) -> None:
"""
Warning: Not intended to be used in production yet.
"""
if not force:
raise click.ClickException("Must use --force to run")
setup_logging(log_level)
logger = logging.getLogger("snuba.bootstrap")
import time
if kafka:
logger.debug("Using Kafka with %r", bootstrap_server)
from confluent_kafka.admin import AdminClient, NewTopic
attempts = 0
while True:
try:
logger.debug("Attempting to connect to Kafka (attempt %d)",
attempts)
client = AdminClient({
"bootstrap.servers":
",".join(bootstrap_server),
"socket.timeout.ms":
1000,
})
client.list_topics(timeout=1)
break
except Exception as e:
logger.error("Connection to Kafka failed (attempt %d)",
attempts,
exc_info=e)
attempts += 1
if attempts == 60:
raise
time.sleep(1)
topics = {}
for name in ACTIVE_DATASET_NAMES:
dataset = get_dataset(name)
for entity in dataset.get_all_entities():
writable_storage = entity.get_writable_storage()
if writable_storage:
table_writer = writable_storage.get_table_writer()
stream_loader = table_writer.get_stream_loader()
for topic_spec in stream_loader.get_all_topic_specs():
if topic_spec.topic_name in topics:
continue
logger.debug("Adding topic %s to creation list",
topic_spec.topic_name)
topics[topic_spec.topic_name] = NewTopic(
topic_spec.topic_name,
num_partitions=topic_spec.partitions_number,
replication_factor=topic_spec.replication_factor,
)
logger.debug("Initiating topic creation")
for topic, future in client.create_topics(list(topics.values()),
operation_timeout=1).items():
try:
future.result()
logger.info("Topic %s created", topic)
except Exception as e:
logger.error("Failed to create topic %s", topic, exc_info=e)
if migrate:
check_clickhouse_connections()
Runner().run_all(force=True)
def run(events_file,
dataset,
repeat=1,
profile_process=False,
profile_write=False):
"""
Measures the write performance of a dataset
"""
from snuba.consumer import ConsumerWorker
from snuba.migrations.runner import Runner
Runner().run_all(force=True)
writable_storage = dataset.get_writable_storage()
consumer = ConsumerWorker(writable_storage, metrics=DummyMetricsBackend())
messages = get_messages(events_file)
messages = chain(*([messages] * repeat))
processed = []
def process():
with settings_override({"DISCARD_OLD_EVENTS": False}):
for message in messages:
result = consumer.process_message(message)
if result is not None:
processed.append(result)
def write():
consumer.flush_batch(processed)
time_start = time.time()
if profile_process:
filename = tempfile.NamedTemporaryFile(
prefix=os.path.basename(events_file) + ".process.",
suffix=".pstats",
delete=False,
).name
cProfile.runctx("process()", globals(), locals(), filename=filename)
logger.info("Profile Data: %s", filename)
else:
process()
time_write = time.time()
if profile_write:
filename = tempfile.NamedTemporaryFile(
prefix=os.path.basename(events_file) + ".write.",
suffix=".pstats",
delete=False,
).name
cProfile.runctx("write()", globals(), locals(), filename=filename)
logger.info("Profile Data: %s", filename)
else:
write()
time_finish = time.time()
time_to_process = (time_write - time_start) * 1000
time_to_write = (time_finish - time_write) * 1000
time_total = (time_finish - time_start) * 1000
num_events = len(processed)
logger.info("Number of events: %s" % str(num_events).rjust(10, " "))
logger.info("Total: %sms" % format_time(time_total))
logger.info("Total process: %sms" % format_time(time_to_process))
logger.info("Total write: %sms" % format_time(time_to_write))
logger.info("Process event: %sms/ea" %
format_time(time_to_process / num_events))
logger.info("Write event: %sms/ea" %
format_time(time_to_write / num_events))
