def _get_migration_status(self) -> Mapping[MigrationKey, Status]:
data: MutableMapping[MigrationKey, Status] = {}
migration_groups = (
"("
+ (
", ".join(
[
escape_string(group.value)
for group in get_active_migration_groups()
]
)
)
+ ")"
)
try:
for row in self.__connection.execute(
f"SELECT group, migration_id, status FROM {self.__table_name} FINAL WHERE group IN {migration_groups}"
):
group_name, migration_id, status_name = row
data[MigrationKey(MigrationGroup(group_name), migration_id)] = Status(
status_name
)
except ClickhouseError as e:
# If the table wasn't created yet, no migrations have started.
if e.code != errors.ErrorCodes.UNKNOWN_TABLE:
raise e
return data
def visitLiteral(self, exp: Literal) -> str:
if exp.value is None:
return "NULL"
elif exp.value is True:
return "true"
elif exp.value is False:
return "false"
elif isinstance(exp.value, str):
return escape_string(exp.value)
elif isinstance(exp.value, (int, float)):
return str(exp.value)
elif isinstance(exp.value, datetime):
value = exp.value.replace(tzinfo=None, microsecond=0)
return "toDateTime('{}')".format(value.isoformat())
elif isinstance(exp.value, date):
return "toDate('{}')".format(exp.value.isoformat())
else:
raise ValueError(f"Unexpected literal type {type(exp.value)}")
def escape_literal(
value: Optional[Union[str, datetime, date, List[Any], Tuple[Any], numbers.Number]]
) -> str:
"""
Escape a literal value for use in a SQL clause.
"""
if isinstance(value, str):
return escape_string(value)
elif isinstance(value, datetime):
value = value.replace(tzinfo=None, microsecond=0)
return "toDateTime('{}', 'Universal')".format(value.isoformat())
elif isinstance(value, date):
return "toDate('{}', 'Universal')".format(value.isoformat())
elif isinstance(value, (list, tuple)):
return "({})".format(", ".join(escape_literal(v) for v in value))
elif isinstance(value, numbers.Number):
return str(value)
elif value is None:
return ""
else:
raise ValueError("Do not know how to escape {} for SQL".format(type(value)))
def get_sql(self, cluster: ClickhouseCluster, table_name: str) -> str:
re = escape_string(self.__table_name_regex)
return f"Merge(currentDatabase(), {re})"
def process_delete_tag(
message: Mapping[str, Any],
schema: TableSchema,
tag_column_map: Mapping[str, Mapping[str, str]],
promoted_tags: Mapping[str, Sequence[str]],
) -> Optional[Replacement]:
tag = message["tag"]
if not tag:
return None
assert isinstance(tag, str)
timestamp = datetime.strptime(message["datetime"], settings.PAYLOAD_DATETIME_FORMAT)
tag_column_name = tag_column_map["tags"].get(tag, tag)
is_promoted = tag in promoted_tags["tags"]
where = """\
WHERE project_id = %(project_id)s
AND received <= CAST('%(timestamp)s' AS DateTime)
AND NOT deleted
"""
if is_promoted:
prewhere = " PREWHERE %(tag_column)s IS NOT NULL "
else:
prewhere = " PREWHERE has(`tags.key`, %(tag_str)s) "
insert_query_template = (
"""\
INSERT INTO %(dist_write_table_name)s (%(all_columns)s)
SELECT %(select_columns)s
FROM %(dist_read_table_name)s FINAL
"""
+ prewhere + where
)
all_columns = [
col
for col in schema.get_columns()
if Materialized not in col.type.get_all_modifiers()
]
select_columns = []
for col in all_columns:
if is_promoted and col.flattened == tag_column_name:
select_columns.append("NULL")
elif col.flattened == "tags.key":
select_columns.append(
"arrayFilter(x -> (indexOf(`tags.key`, x) != indexOf(`tags.key`, %s)), `tags.key`)"
% escape_string(tag)
)
elif col.flattened == "tags.value":
select_columns.append(
"arrayMap(x -> arrayElement(`tags.value`, x), arrayFilter(x -> x != indexOf(`tags.key`, %s), arrayEnumerate(`tags.value`)))"
% escape_string(tag)
)
elif col.flattened == "_tags_flattened":
select_columns.append(FLATTENED_COLUMN_TEMPLATE % escape_string(tag))
else:
select_columns.append(col.escaped)
all_column_names = [col.escaped for col in all_columns]
query_args = {
"all_columns": ", ".join(all_column_names),
"select_columns": ", ".join(select_columns),
"project_id": message["project_id"],
"tag_str": escape_string(tag),
"tag_column": escape_identifier(tag_column_name),
"timestamp": timestamp.strftime(DATETIME_FORMAT),
}
count_query_template = (
"""\
SELECT count()
FROM %(dist_read_table_name)s FINAL
"""
+ prewhere + where
)
query_time_flags = (NEEDS_FINAL, message["project_id"])
return Replacement(
count_query_template, insert_query_template, query_args, query_time_flags
)
def process_delete_tag(
message: ReplacementMessage,
all_columns: Sequence[FlattenedColumn],
tag_column_map: Mapping[str, Mapping[str, str]],
promoted_tags: Mapping[str, Sequence[str]],
use_promoted_prewhere: bool,
schema: WritableTableSchema,
) -> Optional[Replacement]:
tag = message.data["tag"]
if not tag:
return None
assert isinstance(tag, str)
timestamp = datetime.strptime(
message.data["datetime"], settings.PAYLOAD_DATETIME_FORMAT
)
tag_column_name = tag_column_map["tags"].get(tag, tag)
is_promoted = tag in promoted_tags["tags"]
where = """\
WHERE project_id = %(project_id)s
AND received <= CAST('%(timestamp)s' AS DateTime)
AND NOT deleted
"""
if is_promoted and use_promoted_prewhere:
prewhere = " PREWHERE %(tag_column)s IS NOT NULL "
else:
prewhere = " PREWHERE has(`tags.key`, %(tag_str)s) "
insert_query_template = (
"""\
INSERT INTO %(table_name)s (%(all_columns)s)
SELECT %(select_columns)s
FROM %(table_name)s FINAL
"""
+ prewhere
+ where
)
select_columns = []
for col in all_columns:
if is_promoted and col.flattened == tag_column_name:
# The promoted tag columns of events are non nullable, but those of
# errors are non nullable. We check the column against the schema
# to determine whether to write an empty string or NULL.
column_type = schema.get_data_source().get_columns().get(tag_column_name)
assert column_type is not None
is_nullable = column_type.type.has_modifier(Nullable)
if is_nullable:
select_columns.append("NULL")
else:
select_columns.append("''")
elif col.flattened == "tags.key":
select_columns.append(
"arrayFilter(x -> (indexOf(`tags.key`, x) != indexOf(`tags.key`, %s)), `tags.key`)"
% escape_string(tag)
)
elif col.flattened == "tags.value":
select_columns.append(
"arrayMap(x -> arrayElement(`tags.value`, x), arrayFilter(x -> x != indexOf(`tags.key`, %s), arrayEnumerate(`tags.value`)))"
% escape_string(tag)
)
else:
select_columns.append(col.escaped)
all_column_names = [col.escaped for col in all_columns]
query_args = {
"all_columns": ", ".join(all_column_names),
"select_columns": ", ".join(select_columns),
"project_id": message.data["project_id"],
"tag_str": escape_string(tag),
"tag_column": escape_identifier(tag_column_name),
"timestamp": timestamp.strftime(DATETIME_FORMAT),
}
count_query_template = (
"""\
SELECT count()
FROM %(table_name)s FINAL
"""
+ prewhere
+ where
)
query_time_flags = (NEEDS_FINAL, message.data["project_id"])
return LegacyReplacement(
count_query_template,
insert_query_template,
query_args,
query_time_flags,
replacement_type=message.action_type,
replacement_message_metadata=message.metadata,
)
def _format_string_literal(self, exp: Literal) -> str:
return self._alias(escape_string(cast(str, exp.value)), exp.alias)
