def load(self, writer: BufferedWriterWrapper) -> None:
logger = logging.getLogger('snuba.bulk-loader')
clickhouse_ro = ClickhousePool(client_settings={
'readonly': True,
})
clickhouse_tables = clickhouse_ro.execute('show tables')
if (self.__dest_table, ) not in clickhouse_tables:
raise ValueError("Destination table %s does not exists" %
self.__dest_table)
table_content = clickhouse_ro.execute("select count(*) from %s" %
self.__dest_table)
if table_content != [(0, )]:
raise ValueError("Destination Table is not empty")
descriptor = self.__source.get_descriptor()
logger.info("Loading snapshot %s", descriptor.id)
with self.__source.get_table_file(self.__source_table) as table:
logger.info("Loading table %s from file", self.__source_table)
row_count = 0
with writer as buffer_writer:
for row in table:
clickhouse_data = self.__row_processor(row)
buffer_writer.write(clickhouse_data)
row_count += 1
logger.info("Load complete %d records loaded", row_count)
def drop_partitions(
clickhouse: ClickhousePool,
database: str,
table: str,
parts: Sequence[Tuple[datetime, int]],
dry_run: bool = True,
) -> None:
query_template = """\
ALTER TABLE %(database)s.%(table)s DROP PARTITION ('%(date_str)s', %(retention_days)s)
"""
for part_date, retention_days in parts:
args = {
"database": database,
"table": table,
"date_str": part_date.strftime("%Y-%m-%d"),
"retention_days": retention_days,
}
query = (query_template % args).strip()
if dry_run:
logger.info("Dry run: " + query)
else:
logger.info("Dropping partition: " + query)
clickhouse.execute(query)
def test_fallback_logic() -> None:
state.set_config("use_fallback_host_in_native_connection_pool", 1)
network_failure_connection = mock.Mock()
network_failure_connection.execute.side_effect = EOFError()
verification_connection = mock.Mock()
verification_connection.execute.return_value = []
pool = ClickhousePool(CLUSTER_HOST, CLUSTER_PORT, "test", "test",
TEST_DB_NAME)
# The execute method will try to reuse a single slot in the connection
# pool but reestablish new connections with _create_conn if a connection
# fails with a network-related error. It may be cleaner to move connection
# negotation/establishment into another class for separation of concerns.
with mock.patch.object(pool,
"_create_conn",
lambda x, y=False: network_failure_connection):
pool.pool = queue.LifoQueue(1)
pool.pool.put(network_failure_connection, block=False)
pool.fallback_pool = queue.LifoQueue(1)
pool.fallback_pool.put(verification_connection, block=False)
pool.execute("SELECT something")
assert (network_failure_connection.execute.call_count == 3
), "Expected three (failed) attempts with main connection pool"
assert (
verification_connection.execute.call_count == 1
), "Expected one (successful) attempt with fallback connection pool"
def test_add_node() -> None:
host_name = os.environ.get("CLICKHOUSE_HOST", "localhost")
port = int(os.environ.get("CLICKHOUSE_PORT", 9000))
user = "******"
password = ""
database = os.environ.get("CLICKHOUSE_DATABASE", "default")
client = ClickhousePool(
host_name,
port,
user,
password,
database,
)
assert set(client.execute("SHOW TABLES")) == set()
runner.Runner.add_node(
node_type=cluster.ClickhouseNodeType.LOCAL,
storage_sets=[StorageSetKey.OUTCOMES],
host_name=host_name,
port=port,
user=user,
password=password,
database=database,
)
assert set(client.execute("SHOW TABLES")) == {
("outcomes_raw_local", ),
("outcomes_hourly_local", ),
("outcomes_mv_hourly_local", ),
}
def test_reconnect(FakeClient: Client) -> None:
# If the connection NetworkErrors a first time, make sure we call it a second time.
FakeClient.return_value.execute.side_effect = [
errors.NetworkError,
'{"data": "to my face"}',
]
cp = ClickhousePool("0:0:0:0", 9000, "default", "", "default")
cp.execute("SHOW TABLES")
assert FakeClient.return_value.execute.mock_calls == [
call(
"SHOW TABLES",
params=None,
with_column_types=False,
query_id=None,
settings=None,
types_check=False,
columnar=False,
),
call(
"SHOW TABLES",
params=None,
with_column_types=False,
query_id=None,
settings=None,
types_check=False,
columnar=False,
),
]
def test_reconnect(self, FakeClient):
# If the connection NetworkErrors a first time, make sure we call it a second time.
FakeClient.return_value.execute.side_effect = [
errors.NetworkError, '{"data": "to my face"}'
]
cp = ClickhousePool()
cp.execute("SHOW TABLES")
assert FakeClient.return_value.execute.mock_calls == [
call("SHOW TABLES"), call("SHOW TABLES")
]
def test_reconnect(FakeClient) -> None:
# If the connection NetworkErrors a first time, make sure we call it a second time.
FakeClient.return_value.execute.side_effect = [
errors.NetworkError,
'{"data": "to my face"}',
]
cp = ClickhousePool("0:0:0:0", 9000, "default", "", "default")
cp.execute("SHOW TABLES")
assert FakeClient.return_value.execute.mock_calls == [
call("SHOW TABLES"),
call("SHOW TABLES"),
]
def get_active_partitions(clickhouse: ClickhousePool,
storage: WritableTableStorage, database: str,
table: str) -> Sequence[util.Part]:
response = clickhouse.execute(
"""
SELECT DISTINCT partition
FROM system.parts
WHERE database = %(database)s
AND table = %(table)s
AND active = 1
""",
{
"database": database,
"table": table
},
)
schema = storage.get_schema()
assert isinstance(schema, TableSchema)
part_format = schema.get_part_format()
assert part_format is not None
return [
util.decode_part_str(part, part_format) for part, in response.results
]
def perform_select_query(
columns: Sequence[str],
table: str,
where: Optional[Dict[str, str]],
limit: Optional[str],
connection: ClickhousePool,
) -> Sequence[Any]:
""" Performs a SELECT query, with optional WHERE and LIMIT clauses
Arguments:
columns -- a list of columns to be SELECTed
table -- the name of the table, upon which query is being run
where -- a dict of WHERE conditions, (str, str) key-value pairs
limit -- LIMIT argument, passed in as str
connection -- ClickHouse connection object for query execution
"""
select_clause = "SELECT " + (", ".join(columns))
from_clause = " FROM " + table
where_clause = ""
if where:
where_elems = [(key + " = " + "'" + where[key] + "'") for key in where]
where_clause = " WHERE " + (" AND ".join(where_elems))
limit_clause = (" LIMIT " + limit) if limit else ""
full_query = select_clause + from_clause + where_clause + limit_clause
return connection.execute(full_query)
def add_node(
self,
node_type: ClickhouseNodeType,
storage_sets: Sequence[StorageSetKey],
host_name: str,
port: int,
user: str,
password: str,
database: str,
) -> None:
client_settings = ClickhouseClientSettings.MIGRATE.value
clickhouse = ClickhousePool(
host_name,
port,
user,
password,
database,
client_settings=client_settings.settings,
send_receive_timeout=client_settings.timeout,
)
migrations: List[Migration] = []
for group in get_active_migration_groups():
group_loader = get_group_loader(group)
for migration_id in group_loader.get_migrations():
migration = group_loader.load_migration(migration_id)
migrations.append(migration)
for migration in migrations:
if isinstance(migration, ClickhouseNodeMigration):
operations = (
migration.forwards_local()
if node_type == ClickhouseNodeType.LOCAL
else migration.forwards_dist()
)
for sql_op in operations:
if isinstance(sql_op, SqlOperation):
if sql_op._storage_set in storage_sets:
sql = sql_op.format_sql()
print(f"Executing {sql}")
clickhouse.execute(sql)
elif isinstance(migration, CodeMigration):
for python_op in migration.forwards_global():
python_op.execute_new_node(storage_sets)
def optimize_partitions(
clickhouse: ClickhousePool,
database: str,
table: str,
parts: Sequence[util.Part],
ignore_cutoff: bool,
clickhouse_host: Optional[str] = None,
) -> None:
query_template = """\
OPTIMIZE TABLE %(database)s.%(table)s
PARTITION %(partition)s FINAL
"""
# Adding 10 minutes to the current time before finding the midnight time
# to ensure this keeps working even if the system clock of the host that
# starts the pod is slightly ahead of the system clock of the host running
# the job. This prevents us from getting the wrong midnight.
last_midnight = (datetime.now() + timedelta(minutes=10)).replace(
hour=0, minute=0, second=0, microsecond=0)
if not ignore_cutoff:
cutoff_time: Optional[datetime] = (last_midnight +
settings.OPTIMIZE_JOB_CUTOFF_TIME)
logger.info("Cutoff time: %s", str(cutoff_time))
else:
cutoff_time = None
logger.info("Ignoring cutoff time")
for part in parts:
if cutoff_time is not None and datetime.now() > cutoff_time:
raise JobTimeoutException(
"Optimize job is running past the cutoff time. Abandoning.")
args = {
"database": database,
"table": table,
"partition": part.name,
}
query = (query_template % args).strip()
logger.info(f"Optimizing partition: {part.name}")
start = time.time()
clickhouse.execute(query)
metrics.timing(
"optimized_part",
time.time() - start,
tags=_get_metrics_tags(table, clickhouse_host),
)
def pytest_configure() -> None:
"""
Set up the Sentry SDK to avoid errors hidden by configuration.
Ensure the snuba_test database exists
"""
assert (
settings.TESTING
), "settings.TESTING is False, try `SNUBA_SETTINGS=test` or `make test`"
setup_sentry()
for cluster in settings.CLUSTERS:
connection = ClickhousePool(
cluster["host"], cluster["port"], "default", "", "default",
)
database_name = cluster["database"]
connection.execute(f"DROP DATABASE IF EXISTS {database_name};")
connection.execute(f"CREATE DATABASE {database_name};")
def check_clickhouse(clickhouse: ClickhousePool) -> None:
ver = clickhouse.execute("SELECT version()").results[0][0]
# The newer versions of altinity on arm add this to the version
# and it breaks this check
ver = ver.replace(".testingarm", "")
ver = ver.replace(".altinitystable", "")
if version.parse(ver) < version.parse(CLICKHOUSE_SERVER_MIN_VERSION):
raise InvalidClickhouseVersion(
f"Snuba requires Clickhouse version {CLICKHOUSE_SERVER_MIN_VERSION} ({clickhouse.host}:{clickhouse.port} - {ver})"
)
def test_concurrency_limit() -> None:
connection = mock.Mock()
connection.execute.side_effect = TestError("some error")
state.set_config("simultaneous_queries_sleep_seconds", 0.5)
pool = ClickhousePool("host", 100, "test", "test", "test")
pool.pool = queue.LifoQueue(1)
pool.pool.put(connection, block=False)
with pytest.raises(ClickhouseError):
pool.execute("SELECT something")
connection.execute.assert_called_once()
connection.reset_mock(side_effect=True)
connection.execute.side_effect = TestConcurrentError("some error")
with pytest.raises(ClickhouseError):
pool.execute("SELECT something")
assert connection.execute.call_count == 2, "Expected two attempts"
def pytest_configure() -> None:
"""
Set up the Sentry SDK to avoid errors hidden by configuration.
Ensure the snuba_test database exists
"""
setup_sentry()
# There is only one cluster in test, so fetch the host from there.
cluster = settings.CLUSTERS[0]
connection = ClickhousePool(
cluster["host"],
cluster["port"],
"default",
"",
"default",
)
database_name = cluster["database"]
connection.execute(f"DROP DATABASE IF EXISTS {database_name};")
connection.execute(f"CREATE DATABASE {database_name};")
class BaseTest(object):
def setup_method(self, test_method, dataset_name=None):
assert settings.TESTING, "settings.TESTING is False, try `SNUBA_SETTINGS=test` or `make test`"
self.database = 'default'
self.dataset_name = dataset_name
if self.dataset_name:
self.dataset = get_dataset(self.dataset_name)
self.clickhouse = ClickhousePool()
for statement in self.dataset.get_dataset_schemas(
).get_drop_statements():
self.clickhouse.execute(statement)
for statement in self.dataset.get_dataset_schemas(
).get_create_statements():
self.clickhouse.execute(statement)
redis_client.flushdb()
def teardown_method(self, test_method):
if self.dataset_name:
for statement in self.dataset.get_dataset_schemas(
).get_drop_statements():
self.clickhouse.execute(statement)
redis_client.flushdb()
def optimize_partitions(
clickhouse: ClickhousePool,
database: str,
table: str,
parts: Sequence[util.Part],
) -> None:
query_template = """\
OPTIMIZE TABLE %(database)s.%(table)s
PARTITION %(partition)s FINAL
"""
for part in parts:
args = {
"database": database,
"table": table,
"partition": part.name,
}
query = (query_template % args).strip()
logger.info(f"Optimizing partition: {part.name}")
clickhouse.execute(query)
def drop_partitions(
clickhouse: ClickhousePool,
database: str,
table: str,
parts: Sequence[util.Part],
dry_run: bool = True,
) -> None:
query_template = """\
ALTER TABLE %(database)s.%(table)s DROP PARTITION %(partition)s
"""
for part in parts:
args = {
"database": database,
"table": table,
"partition": part.name,
}
query = (query_template % args).strip()
if dry_run:
logger.info("Dry run: " + query)
else:
logger.info("Dropping partition: " + query)
clickhouse.execute(query)
def get_active_partitions(
clickhouse: ClickhousePool, database: str, table: str
) -> Sequence[util.Part]:
response = clickhouse.execute(
"""
SELECT DISTINCT partition
FROM system.parts
WHERE database = %(database)s
AND table = %(table)s
AND active = 1
""",
{"database": database, "table": table},
)
return [util.decode_part_str(part) for part, in response]
def test_messages(self):
processor = GroupAssigneeProcessor('sentry_groupasignee')
metadata = KafkaMessageMetadata(
offset=42,
partition=0,
)
begin_msg = json.loads(self.BEGIN_MSG)
ret = processor.process_message(begin_msg, metadata)
assert ret is None
commit_msg = json.loads(self.COMMIT_MSG)
ret = processor.process_message(commit_msg, metadata)
assert ret is None
insert_msg = json.loads(self.INSERT_MSG)
ret = processor.process_message(insert_msg, metadata)
assert ret.data == [self.PROCESSED]
self.write_processed_records(ret.data)
cp = ClickhousePool()
ret = cp.execute("SELECT * FROM test_groupassignee_local;")
assert ret[0] == (
42, # offset
0, # deleted
2, # project_id
1359, # group_id
datetime(2019, 9, 19, 0, 17, 55),
1, # user_id
None, # team_id
)
update_msg = json.loads(self.UPDATE_MSG_NO_KEY_CHANGE)
ret = processor.process_message(update_msg, metadata)
assert ret.data == [self.PROCESSED]
# Tests an update with key change which becomes a two inserts:
# one deletion and the insertion of the new row.
update_msg = json.loads(self.UPDATE_MSG_WITH_KEY_CHANGE)
ret = processor.process_message(update_msg, metadata)
assert ret.data == [self.DELETED, self.PROCESSED_UPDATE]
delete_msg = json.loads(self.DELETE_MSG)
ret = processor.process_message(delete_msg, metadata)
assert ret.data == [self.DELETED]
def test_bulk_load(self):
row = GroupAssigneeRow.from_bulk({
'project_id': '2',
'group_id': '1359',
'date_added': '2019-09-19 00:17:55+00',
'user_id': '1',
'team_id': '',
})
self.write_processed_records(row.to_clickhouse())
cp = ClickhousePool()
ret = cp.execute("SELECT * FROM test_groupassignee_local;")
assert ret[0] == (
0, # offset
0, # deleted
2, # project_id
1359, # group_id
datetime(2019, 9, 19, 0, 17, 55),
1, # user_id
None, # team_id
)
def test_messages(self):
processor = GroupedMessageProcessor('sentry_groupedmessage')
metadata = KafkaMessageMetadata(
offset=42,
partition=0,
)
begin_msg = json.loads(self.BEGIN_MSG)
ret = processor.process_message(begin_msg, metadata)
assert ret is None
commit_msg = json.loads(self.COMMIT_MSG)
ret = processor.process_message(commit_msg, metadata)
assert ret is None
insert_msg = json.loads(self.INSERT_MSG)
ret = processor.process_message(insert_msg, metadata)
assert ret[1] == self.PROCESSED
self.write_processed_records(ret[1])
cp = ClickhousePool()
ret = cp.execute("SELECT * FROM test_groupedmessage_local;")
assert ret[0] == (
42, # offset
0, # deleted
2, # project_id
74, # id
0, # status
datetime(2019, 6, 19, 6, 46, 28),
datetime(2019, 6, 19, 6, 45, 32),
datetime(2019, 6, 19, 6, 45, 32),
None,
)
update_msg = json.loads(self.UPDATE_MSG)
ret = processor.process_message(update_msg, metadata)
assert ret[1] == self.PROCESSED
delete_msg = json.loads(self.DELETE_MSG)
ret = processor.process_message(delete_msg, metadata)
assert ret[1] == self.DELETED
def test_bulk_load(self):
row = GroupedMessageRow.from_bulk({
'project_id': '2',
'id': '10',
'status': '0',
'last_seen': '2019-06-28 17:57:32+00',
'first_seen': '2019-06-28 06:40:17+00',
'active_at': '2019-06-28 06:40:17+00',
'first_release_id': '26',
})
self.write_processed_records(row.to_clickhouse())
cp = ClickhousePool()
ret = cp.execute("SELECT * FROM test_groupedmessage_local;")
assert ret[0] == (
0, # offset
0, # deleted
2, # project_id
10, # id
0, # status
datetime(2019, 6, 28, 17, 57, 32),
datetime(2019, 6, 28, 6, 40, 17),
datetime(2019, 6, 28, 6, 40, 17),
26,
)
def check_clickhouse(clickhouse: ClickhousePool) -> None:
ver = clickhouse.execute("SELECT version()")[0][0]
if version.parse(ver) < version.parse(CLICKHOUSE_SERVER_MIN_VERSION):
raise InvalidClickhouseVersion(
f"Snuba requires Clickhouse version {CLICKHOUSE_SERVER_MIN_VERSION}"
)
def get_partitions_to_optimize(
clickhouse: ClickhousePool,
storage: ReadableTableStorage,
database: str,
table: str,
before: Optional[datetime] = None,
) -> Sequence[util.Part]:
engine = clickhouse.execute(
"""
SELECT engine
FROM system.tables
WHERE (database = %(database)s) AND (name = %(table)s)
""",
{
"database": database,
"table": table
},
)
if not engine:
logger.warning("Table %s.%s doesn't exist on %s:%s" %
(database, table, clickhouse.host, clickhouse.port))
return []
if engine[0][0].startswith("Replicated"):
is_leader = clickhouse.execute(
"""
SELECT is_leader
FROM system.replicas
WHERE (database = %(database)s) AND (table = %(table)s)
""",
{
"database": database,
"table": table
},
)
# response: [(0,)] for non-leader or [(1,)] for leader
if not (len(is_leader) == 1 and is_leader[0][0]):
return []
active_parts = clickhouse.execute(
"""
SELECT
partition,
count() AS c
FROM system.parts
WHERE active
AND database = %(database)s
AND table = %(table)s
GROUP BY partition
HAVING c > 1
ORDER BY c DESC, partition
""",
{
"database": database,
"table": table
},
)
schema = storage.get_schema()
assert isinstance(schema, TableSchema)
part_format = schema.get_part_format()
assert part_format is not None
parts = [
util.decode_part_str(part, part_format) for part, count in active_parts
]
if before:
parts = [
p for p in parts
if (p.date + timedelta(days=6 - p.date.weekday())) < before
]
return parts