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)