insights/insights/models/event/sql.py

from django.conf import settings

from insights.insightsql.escape_sql import escape_datastore_identifier, escape_datastore_string

from insights.datastore.base_sql import COPY_ROWS_BETWEEN_TEAMS_BASE_SQL
from insights.datastore.cluster import ON_CLUSTER_CLAUSE

# The events-JSON schema declarations (table names, JSON subcolumn types, index DDL) live in the
# Django-free insights.datastore.events_json module so the InsightsQL engine can use them without
# booting Django; re-exported here for existing callers.
from insights.datastore.events_json import (
    DISTRIBUTED_EVENTS_JSON_TABLE,
    EVENTS_JSON_DATA_TABLE,
    EVENTS_JSON_DATA_TABLE_INDEXES,
    EVENTS_JSON_INDEXED_PROPERTY_NAMES,  # noqa: F401
    EVENTS_PROPERTIES_JSON_SUBCOLUMNS,
    KAFKA_EVENTS_NATIVE_JSON_TABLE,
    PERSON_PROPERTIES_JSON_SUBCOLUMNS,
    WRITABLE_EVENTS_JSON_TABLE,
    _json_subcolumn_type_supports_nullable,
)
from insights.datastore.indexes import index_by_kafka_timestamp
from insights.datastore.kafka_engine import (
    CONSUMER_GROUP_EVENTS_JSON,
    CONSUMER_GROUP_EVENTS_JSON_NATIVE_JSON,
    CONSUMER_GROUP_EVENTS_JSON_WS,
    KAFKA_COLUMNS,
    STORAGE_POLICY,
    kafka_engine,
    trim_quotes_expr,
)
from insights.datastore.property_groups import property_groups
from insights.datastore.table_engines import Distributed, ReplacingMergeTree, ReplicationScheme
from insights.kafka_client.topics import KAFKA_EVENTS_JSON


def EVENTS_DATA_TABLE():
    return "sharded_events"


def WRITABLE_EVENTS_DATA_TABLE():
    return "writable_events"


def _json_column_type(max_dynamic_types: int, max_dynamic_paths: int, subcolumns: dict[str, str]) -> str:
    explicit_paths = ", ".join(
        f"{escape_datastore_identifier(name)} {column_type}" for name, column_type in subcolumns.items()
    )
    return f"JSON(max_dynamic_types = {max_dynamic_types}, max_dynamic_paths = {max_dynamic_paths}, {explicit_paths})"


def EVENTS_PROPERTIES_JSON_TYPE() -> str:
    return _json_column_type(8, 256, EVENTS_PROPERTIES_JSON_SUBCOLUMNS)


def PERSON_PROPERTIES_JSON_TYPE() -> str:
    return _json_column_type(6, 32, PERSON_PROPERTIES_JSON_SUBCOLUMNS)


def json_property_presence_expr(column: str, prop: str) -> str:
    """SQL predicate testing whether a (possibly dotted) property path is present in a native-JSON
    events column, for deletion/mutation predicates that run against the JSON events tables.

    Mirrors the InsightsQL resolver's JSONHas lowering: a typed subcolumn is present when non-null; a
    dynamic path is present when its scalar read is non-null or its sub-object read is non-empty.
    Datastore's JSONHas() cannot be used directly on a JSON-typed column — it does not see typed
    paths or nested objects there.
    """
    if column not in ("properties", "person_properties"):
        raise ValueError(f"unsupported JSON events column: {column}")
    subcolumns = EVENTS_PROPERTIES_JSON_SUBCOLUMNS if column == "properties" else PERSON_PROPERTIES_JSON_SUBCOLUMNS
    parts = prop.split(".")
    column_sql = escape_datastore_identifier(column)
    path_sql = ".".join(escape_datastore_identifier(part) for part in parts)
    scalar = f"{column_sql}.{path_sql}"
    if len(parts) == 1 and prop in subcolumns:
        if not _json_subcolumn_type_supports_nullable(subcolumns[prop]):
            # Native container paths use the empty value for both missing and explicitly empty values.
            return f"notEmpty({scalar})"
        return f"isNotNull({scalar})"
    if parts[0] in subcolumns:
        head = f"{column_sql}.{escape_datastore_identifier(parts[0])}"
        head_document = head if subcolumns[parts[0]] in ("String", "Nullable(String)") else f"toJSONString({head})"
        tail = ", ".join(escape_datastore_string(part) for part in parts[1:])
        return f"JSONHas(ifNull({head_document}, ''), {tail})"
    sub_object = f"{column_sql}.^{path_sql}"
    return f"(isNotNull({scalar}) OR toJSONString({sub_object}) != '{{}}')"


def TRUNCATE_EVENTS_TABLE_SQL():
    return f"TRUNCATE TABLE IF EXISTS {EVENTS_DATA_TABLE()} {ON_CLUSTER_CLAUSE()}"


def TRUNCATE_EVENTS_JSON_TABLE_SQL():
    return f"TRUNCATE TABLE IF EXISTS {EVENTS_JSON_DATA_TABLE} {ON_CLUSTER_CLAUSE()}"


def DROP_EVENTS_TABLE_SQL():
    return f"DROP TABLE IF EXISTS {EVENTS_DATA_TABLE()} {ON_CLUSTER_CLAUSE()}"


DROP_DISTRIBUTED_EVENTS_TABLE_SQL = f"DROP TABLE IF EXISTS events {ON_CLUSTER_CLAUSE()}"

INSERTED_AT_COLUMN = ", inserted_at Nullable(DateTime64(6, 'UTC')) DEFAULT NOW64()"
INSERTED_AT_NOT_NULLABLE_COLUMN = ", inserted_at DateTime64(6, 'UTC') DEFAULT NOW64()"
KAFKA_CONSUMER_BREADCRUMBS_COLUMN = ", consumer_breadcrumbs Array(String)"

EVENTS_TABLE_BASE_SQL = """
CREATE TABLE IF NOT EXISTS {table_name} {on_cluster_clause}
(
    uuid UUID,
    event VARCHAR,
    properties VARCHAR CODEC(ZSTD(3)),
    timestamp DateTime64(6, 'UTC'),
    team_id Int64,
    distinct_id VARCHAR,
    elements_chain VARCHAR,
    created_at DateTime64(6, 'UTC'),
    person_id UUID,
    person_created_at DateTime64,
    person_properties VARCHAR Codec(ZSTD(3)),
    group0_properties VARCHAR Codec(ZSTD(3)),
    group1_properties VARCHAR Codec(ZSTD(3)),
    group2_properties VARCHAR Codec(ZSTD(3)),
    group3_properties VARCHAR Codec(ZSTD(3)),
    group4_properties VARCHAR Codec(ZSTD(3)),
    group0_created_at DateTime64,
    group1_created_at DateTime64,
    group2_created_at DateTime64,
    group3_created_at DateTime64,
    group4_created_at DateTime64,
    person_mode Enum8('full' = 0, 'propertyless' = 1, 'force_upgrade' = 2),
    historical_migration Bool
    {dynamically_materialized_columns}
    {materialized_columns}
    {extra_fields}
    {indexes}
) ENGINE = {engine}
"""


# Shared per-team allocation: each team picks its slots independently from this range, and
# (team_id, slot_index) → property_name is resolved at write/read time via the dmat dictionary.
# Cap matches MAX_SLOTS_PER_TEAM so every team can fully saturate its slots.
DMAT_STRING_COLUMN_COUNT = 10


def EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS() -> str:
    s = [f"`dmat_string_{i}` Nullable(String)" for i in range(DMAT_STRING_COLUMN_COUNT)]
    return f"    , {'\n    , '.join(s)}"


def ALTER_TABLE_ADD_DYNAMICALLY_MATERIALIZED_COLUMNS(table: str) -> str:
    return ALTER_TABLE_ADD_DMAT_STRING_COLUMNS(table, 0, DMAT_STRING_COLUMN_COUNT)


def ALTER_TABLE_ADD_DMAT_STRING_COLUMNS(table: str, start: int, end_exclusive: int) -> str:
    """ALTER TABLE statement adding dmat_string columns in a half-open index range."""
    pieces = [f"ADD COLUMN IF NOT EXISTS `dmat_string_{i}` Nullable(String)" for i in range(start, end_exclusive)]
    return f"ALTER TABLE {table} \n {',\n'.join(pieces)}"


def MV_DYNAMICALLY_MATERIALIZED_COLUMNS() -> str:
    return ",\n".join(f"dmat_string_{i}" for i in range(DMAT_STRING_COLUMN_COUNT))


EVENTS_TABLE_MATERIALIZED_COLUMNS = f"""
    , $group_0 VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$group_0')")} COMMENT 'column_materializer::$group_0'
    , $group_1 VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$group_1')")} COMMENT 'column_materializer::$group_1'
    , $group_2 VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$group_2')")} COMMENT 'column_materializer::$group_2'
    , $group_3 VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$group_3')")} COMMENT 'column_materializer::$group_3'
    , $group_4 VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$group_4')")} COMMENT 'column_materializer::$group_4'
    , $window_id VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$window_id')")} COMMENT 'column_materializer::$window_id'
    , $session_id VARCHAR MATERIALIZED {trim_quotes_expr("JSONExtractRaw(properties, '$session_id')")} COMMENT 'column_materializer::$session_id'
    , $session_id_uuid Nullable(UInt128) MATERIALIZED toUInt128(JSONExtract(properties, '$session_id', 'Nullable(UUID)'))
    , elements_chain_href String MATERIALIZED extract(elements_chain, '(?::|\")href="(.*?)"')
    , elements_chain_texts Array(String) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?::|\")text="(.*?)"'))
    , elements_chain_ids Array(String) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?::|\")attr_id="(.*?)"'))
    , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label')) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?:^|;)(a|button|form|input|select|textarea|label)(?:\\.|$|:)'))
    , INDEX `minmax_$group_0` `$group_0` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$group_1` `$group_1` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$group_2` `$group_2` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$group_3` `$group_3` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$group_4` `$group_4` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$window_id` `$window_id` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$session_id` `$session_id` TYPE minmax GRANULARITY 1
    , INDEX `minmax_$session_id_uuid` `$session_id_uuid` TYPE minmax GRANULARITY 1
    , {", ".join(property_groups.get_create_table_pieces("sharded_events"))}
"""

EVENTS_TABLE_PROXY_MATERIALIZED_COLUMNS = f"""
    , $group_0 VARCHAR COMMENT 'column_materializer::$group_0'
    , $group_1 VARCHAR COMMENT 'column_materializer::$group_1'
    , $group_2 VARCHAR COMMENT 'column_materializer::$group_2'
    , $group_3 VARCHAR COMMENT 'column_materializer::$group_3'
    , $group_4 VARCHAR COMMENT 'column_materializer::$group_4'
    , $window_id VARCHAR COMMENT 'column_materializer::$window_id'
    , $session_id VARCHAR COMMENT 'column_materializer::$session_id'
    , $session_id_uuid Nullable(UInt128)
    , elements_chain_href String COMMENT 'column_materializer::elements_chain::href'
    , elements_chain_texts Array(String) COMMENT 'column_materializer::elements_chain::texts'
    , elements_chain_ids Array(String) COMMENT 'column_materializer::elements_chain::ids'
    , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label')) COMMENT 'column_materializer::elements_chain::elements'
    , {", ".join(property_groups.get_create_table_pieces("events"))}
"""


def EVENTS_DATA_TABLE_ENGINE():
    return ReplacingMergeTree("events", ver="_timestamp", replication_scheme=ReplicationScheme.SHARDED)


def EVENTS_JSON_DATA_TABLE_ENGINE():
    return ReplacingMergeTree("events_json", ver="_timestamp", replication_scheme=ReplicationScheme.SHARDED)


def _json_subcolumn(column: str, path: str) -> str:
    return f"{escape_datastore_identifier(column)}.{escape_datastore_identifier(path)}"


EVENTS_JSON_DATA_COMPATIBILITY_COLUMNS = f"""
    , $group_0 String ALIAS ifNull({_json_subcolumn("properties", "$group_0")}, '')
    , $group_1 String ALIAS ifNull({_json_subcolumn("properties", "$group_1")}, '')
    , $group_2 String ALIAS ifNull({_json_subcolumn("properties", "$group_2")}, '')
    , $group_3 String ALIAS ifNull({_json_subcolumn("properties", "$group_3")}, '')
    , $group_4 String ALIAS ifNull({_json_subcolumn("properties", "$group_4")}, '')
    , $window_id String ALIAS ifNull({_json_subcolumn("properties", "$window_id")}, '')
    , $session_id String ALIAS ifNull({_json_subcolumn("properties", "$session_id")}, '')
    , $session_id_uuid Nullable(UInt128) ALIAS toUInt128(toUUIDOrNull({_json_subcolumn("properties", "$session_id")}))
    , elements_chain_href String MATERIALIZED extract(elements_chain, '(?::|\")href="(.*?)"')
    , elements_chain_texts Array(String) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?::|\")text="(.*?)"'))
    , elements_chain_ids Array(String) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?::|\")attr_id="(.*?)"'))
    , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label')) MATERIALIZED arrayDistinct(extractAll(elements_chain, '(?:^|;)(a|button|form|input|select|textarea|label)(?:\\.|$|:)'))
"""


EVENTS_JSON_PROXY_COMPATIBILITY_COLUMNS = """
    , $group_0 String
    , $group_1 String
    , $group_2 String
    , $group_3 String
    , $group_4 String
    , $window_id String
    , $session_id String
    , $session_id_uuid Nullable(UInt128)
    , elements_chain_href String
    , elements_chain_texts Array(String)
    , elements_chain_ids Array(String)
    , elements_chain_elements Array(Enum('a', 'button', 'form', 'input', 'select', 'textarea', 'label'))
"""


EVENTS_JSON_TABLE_BASE_SQL = """
CREATE TABLE IF NOT EXISTS {table_name} {on_cluster_clause}
(
    uuid UUID,
    event String,
    properties {properties_json_type},
    timestamp DateTime64(6, 'UTC'),
    team_id Int64,
    distinct_id String,
    elements_hash String DEFAULT '',
    created_at DateTime64(6, 'UTC'),
    _timestamp DateTime,
    _offset UInt64,
    elements_chain String,
    person_id UUID,
    person_properties {person_properties_json_type},
    group0_properties String CODEC(ZSTD(3)),
    group1_properties String CODEC(ZSTD(3)),
    group2_properties String CODEC(ZSTD(3)),
    group3_properties String CODEC(ZSTD(3)),
    group4_properties String CODEC(ZSTD(3)),
    person_created_at DateTime64(3),
    group0_created_at DateTime64(3),
    group1_created_at DateTime64(3),
    group2_created_at DateTime64(3),
    group3_created_at DateTime64(3),
    group4_created_at DateTime64(3),
    inserted_at Nullable(DateTime64(6, 'UTC')) DEFAULT now64(),
    person_mode Enum8('full' = 0, 'propertyless' = 1, 'force_upgrade' = 2),
    is_deleted Bool DEFAULT false,
    consumer_breadcrumbs Array(String),
    historical_migration Bool DEFAULT false
    {compatibility_columns}
    {indexes}
) ENGINE = {engine}
"""


def EVENTS_JSON_TABLE_SQL(on_cluster: bool = False) -> str:
    return (
        EVENTS_JSON_TABLE_BASE_SQL
        + """PARTITION BY toYYYYMM(timestamp)
PRIMARY KEY (team_id, toDate(timestamp), event, timestamp, cityHash64(distinct_id))
ORDER BY (team_id, toDate(timestamp), event, timestamp, cityHash64(distinct_id), distinct_id, uuid)
SAMPLE BY cityHash64(distinct_id)
SETTINGS index_granularity = 8192, object_serialization_version = 'v3', object_shared_data_serialization_version = 'map_with_buckets', object_shared_data_serialization_version_for_zero_level_parts = 'map', merge_max_block_size = 131072, merge_max_block_size_bytes = 67108864, vertical_merge_algorithm_min_rows_to_activate = 0
"""
    ).format(
        table_name=EVENTS_JSON_DATA_TABLE,
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=EVENTS_JSON_DATA_TABLE_ENGINE(),
        properties_json_type=EVENTS_PROPERTIES_JSON_TYPE(),
        person_properties_json_type=PERSON_PROPERTIES_JSON_TYPE(),
        compatibility_columns=EVENTS_JSON_DATA_COMPATIBILITY_COLUMNS,
        indexes=EVENTS_JSON_DATA_TABLE_INDEXES(),
    )


def WRITABLE_EVENTS_JSON_TABLE_SQL(on_cluster: bool = False) -> str:
    return EVENTS_JSON_TABLE_BASE_SQL.format(
        table_name=WRITABLE_EVENTS_JSON_TABLE,
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(data_table=EVENTS_JSON_DATA_TABLE, sharding_key="sipHash64(distinct_id)"),
        properties_json_type=EVENTS_PROPERTIES_JSON_TYPE(),
        person_properties_json_type=PERSON_PROPERTIES_JSON_TYPE(),
        compatibility_columns="",
        indexes="",
    )


def DISTRIBUTED_EVENTS_JSON_TABLE_SQL(on_cluster: bool = False) -> str:
    return EVENTS_JSON_TABLE_BASE_SQL.format(
        table_name=DISTRIBUTED_EVENTS_JSON_TABLE,
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(data_table=EVENTS_JSON_DATA_TABLE, sharding_key="sipHash64(distinct_id)"),
        properties_json_type=EVENTS_PROPERTIES_JSON_TYPE(),
        person_properties_json_type=PERSON_PROPERTIES_JSON_TYPE(),
        compatibility_columns=EVENTS_JSON_PROXY_COMPATIBILITY_COLUMNS,
        indexes="",
    )


def EVENTS_TABLE_SQL():
    return (
        EVENTS_TABLE_BASE_SQL
        + """PARTITION BY toYYYYMM(timestamp)
ORDER BY (team_id, toDate(timestamp), event, cityHash64(distinct_id), cityHash64(uuid))
{sample_by}
{storage_policy}
"""
    ).format(
        table_name=EVENTS_DATA_TABLE(),
        on_cluster_clause=ON_CLUSTER_CLAUSE(),
        engine=EVENTS_DATA_TABLE_ENGINE(),
        extra_fields=KAFKA_COLUMNS + INSERTED_AT_COLUMN + KAFKA_CONSUMER_BREADCRUMBS_COLUMN,
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns=EVENTS_TABLE_MATERIALIZED_COLUMNS,
        indexes=f"""
    , {index_by_kafka_timestamp(EVENTS_DATA_TABLE())}
    """,
        sample_by="SAMPLE BY cityHash64(distinct_id)",
        storage_policy=STORAGE_POLICY(),
    )


EVENTS_TABLE_INSERTED_AT_INDEX_SQL = """
ALTER TABLE {table_name} ON CLUSTER {cluster}
ADD INDEX IF NOT EXISTS `minmax_inserted_at` COALESCE(`inserted_at`, `_timestamp`)
TYPE minmax
GRANULARITY 1
""".format(table_name=EVENTS_DATA_TABLE(), cluster=settings.DATASTORE_CLUSTER)

EVENTS_TABLE_MATERIALIZE_INSERTED_AT_INDEX_SQL = """
ALTER TABLE {table_name} ON CLUSTER {cluster}
MATERIALIZE INDEX `minmax_inserted_at`
""".format(table_name=EVENTS_DATA_TABLE(), cluster=settings.DATASTORE_CLUSTER)

# we add the settings to prevent poison pills from stopping ingestion
# kafka_skip_broken_messages is an int, not a boolean, so we explicitly set
# the max block size to consume from kafka such that we skip _all_ broken messages
# this is an added safety mechanism given we control payloads to this topic


def KAFKA_EVENTS_TABLE_JSON_SQL():
    return (
        EVENTS_TABLE_BASE_SQL
        + """
    SETTINGS kafka_skip_broken_messages = 100
"""
    ).format(
        table_name="kafka_events_json",
        on_cluster_clause=ON_CLUSTER_CLAUSE(),
        engine=kafka_engine(topic=KAFKA_EVENTS_JSON, group=CONSUMER_GROUP_EVENTS_JSON),
        extra_fields="",
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns="",
        indexes="",
    )


def KAFKA_EVENTS_NATIVE_JSON_TABLE_SQL(on_cluster: bool = False) -> str:
    return (
        EVENTS_TABLE_BASE_SQL
        + """
    SETTINGS kafka_skip_broken_messages = 100
"""
    ).format(
        table_name=KAFKA_EVENTS_NATIVE_JSON_TABLE,
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=kafka_engine(topic=KAFKA_EVENTS_JSON, group=CONSUMER_GROUP_EVENTS_JSON_NATIVE_JSON),
        extra_fields="",
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns="",
        indexes="",
    )


# NOTE: All parameters must have defaults - zero-argument calls must remain valid.
# 8+ frozen migrations and schema.py reference this function without parameters.
#
# Override parameters to create separate pipelines that reuse the events schema:
# - mv_name: unique MV name to avoid conflicts with the main events_json_mv
# - kafka_table: different Kafka table consuming from a different topic
# - target_table: different target table for the MV to write to
# - on_cluster: False when running on specific node roles (e.g., ingestion layer)
#
# Example: error_tracking_events_test uses custom values to create a parallel
# pipeline for validating Node.js ingestion output against the Python pipeline.
def EVENTS_TABLE_JSON_MV_SQL(
    mv_name="events_json_mv",
    kafka_table="kafka_events_json",
    target_table=None,
    on_cluster=True,
):
    if target_table is None:
        target_table = WRITABLE_EVENTS_DATA_TABLE()

    return """
CREATE MATERIALIZED VIEW IF NOT EXISTS {mv_name} {on_cluster_clause}
TO {database}.{target_table}
AS SELECT
uuid,
event,
properties,
timestamp,
team_id,
distinct_id,
elements_chain,
created_at,
person_id,
person_created_at,
person_properties,
group0_properties,
group1_properties,
group2_properties,
group3_properties,
group4_properties,
group0_created_at,
group1_created_at,
group2_created_at,
group3_created_at,
group4_created_at,
person_mode,
historical_migration,
{dynamically_materialized_columns},
_timestamp,
_offset,
arrayMap(
    i -> _headers.value[i],
    arrayFilter(
        i -> _headers.name[i] = 'kafka-consumer-breadcrumbs',
        arrayEnumerate(_headers.name)
    )
) as consumer_breadcrumbs
FROM {database}.{kafka_table}
""".format(
        mv_name=mv_name,
        kafka_table=kafka_table,
        target_table=target_table,
        dynamically_materialized_columns=MV_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        on_cluster_clause=f"ON CLUSTER '{settings.DATASTORE_CLUSTER}'" if on_cluster else "",
        database=settings.DATASTORE_DATABASE,
    )


def _poison_safe_json_cast(column: str, fallback_key: str) -> str:
    """Cast a String properties payload to the JSON column type without ever throwing.

    The JSON type's parser rejects some payloads that are valid JSON — notably integers outside
    [-2^63, 2^64) fail with INCORRECT_DATA. A throwing cast inside a Kafka materialized view is a
    poison message: kafka_skip_broken_messages only covers format parsing, so the consumer would
    retry the same block forever and stall ingestion for the whole table. Unparseable payloads are
    instead preserved verbatim under a single string property so they can be repaired later.
    """
    return (
        f"ifNull(accurateCastOrNull({column}, 'JSON'), "
        f"CAST(concat('{{\"{fallback_key}\":', toJSONString({column}), '}}'), 'JSON')) AS {column}"
    )


# Dual-write materialized view that writes events into the native-JSON schema
# (writable_events_json). It reads from a dedicated Kafka consumer group so JSON-table retries do not
# replay legacy writes through events_json_mv. The string properties/person_properties payloads are
# cast to the destination JSON columns via a poison-safe cast (see _poison_safe_json_cast). Unlike the
# legacy MV this does not project the dmat_string_* columns — they don't exist on the JSON table,
# whose property reads come from JSON subcolumns instead.
def EVENTS_JSON_TABLE_MV_SQL(
    mv_name="events_json_table_mv",
    kafka_table=KAFKA_EVENTS_NATIVE_JSON_TABLE,
    target_table=None,
    on_cluster=True,
):
    if target_table is None:
        target_table = WRITABLE_EVENTS_JSON_TABLE

    return """
CREATE MATERIALIZED VIEW IF NOT EXISTS {mv_name} {on_cluster_clause}
TO {database}.{target_table}
AS SELECT
uuid,
event,
{properties_expr},
timestamp,
team_id,
distinct_id,
elements_chain,
created_at,
person_id,
person_created_at,
{person_properties_expr},
group0_properties,
group1_properties,
group2_properties,
group3_properties,
group4_properties,
group0_created_at,
group1_created_at,
group2_created_at,
group3_created_at,
group4_created_at,
person_mode,
historical_migration,
_timestamp,
_offset,
arrayMap(
    i -> _headers.value[i],
    arrayFilter(
        i -> _headers.name[i] = 'kafka-consumer-breadcrumbs',
        arrayEnumerate(_headers.name)
    )
) as consumer_breadcrumbs
FROM {database}.{kafka_table}
""".format(
        mv_name=mv_name,
        kafka_table=kafka_table,
        target_table=target_table,
        on_cluster_clause=f"ON CLUSTER '{settings.DATASTORE_CLUSTER}'" if on_cluster else "",
        database=settings.DATASTORE_DATABASE,
        properties_expr=_poison_safe_json_cast("properties", "$unparseable_properties"),
        person_properties_expr=_poison_safe_json_cast("person_properties", "$unparseable_properties"),
    )


# WarpStream Kafka engine tables (coexist alongside MSK tables, same target)

KAFKA_EVENTS_JSON_WS_TABLE = "kafka_events_json_ws"
EVENTS_JSON_WS_MV = "events_json_ws_mv"

DROP_KAFKA_EVENTS_JSON_WS_TABLE_SQL = f"DROP TABLE IF EXISTS {KAFKA_EVENTS_JSON_WS_TABLE}"
DROP_EVENTS_JSON_WS_MV_SQL = f"DROP TABLE IF EXISTS {EVENTS_JSON_WS_MV}"


def KAFKA_EVENTS_TABLE_JSON_WS_SQL():
    return (
        EVENTS_TABLE_BASE_SQL
        + """
    SETTINGS kafka_skip_broken_messages = 100, kafka_thread_per_consumer = 1, kafka_num_consumers = 1
"""
    ).format(
        table_name=KAFKA_EVENTS_JSON_WS_TABLE,
        on_cluster_clause=ON_CLUSTER_CLAUSE(False),
        engine=kafka_engine(
            topic=KAFKA_EVENTS_JSON,
            group=CONSUMER_GROUP_EVENTS_JSON_WS,
            named_collection=settings.DATASTORE_KAFKA_WARPSTREAM_INGESTION_NAMED_COLLECTION,
        ),
        extra_fields="",
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns="",
        indexes="",
    )


def EVENTS_TABLE_JSON_WS_MV_SQL():
    return EVENTS_TABLE_JSON_MV_SQL(
        mv_name=EVENTS_JSON_WS_MV,
        kafka_table=KAFKA_EVENTS_JSON_WS_TABLE,
        on_cluster=False,
    )


# Events recent tables


def WRITABLE_EVENTS_RECENT_TABLE():
    return "writable_events_recent"


def EVENTS_RECENT_DATA_TABLE():
    return "events_recent"


def SHARDED_EVENTS_RECENT_DATA_TABLE():
    return "sharded_events_recent"


def DROP_KAFKA_EVENTS_RECENT_TABLE_SQL():
    return f"DROP TABLE IF EXISTS kafka_events_recent_json"


def DROP_EVENTS_RECENT_MV_TABLE_SQL():
    return f"DROP TABLE IF EXISTS events_recent_json_mv"


def TRUNCATE_EVENTS_RECENT_TABLE_SQL():
    return f"TRUNCATE TABLE IF EXISTS {EVENTS_RECENT_DATA_TABLE()} {ON_CLUSTER_CLAUSE()}"


def EVENTS_RECENT_TABLE_SQL(on_cluster=False):
    return EVENTS_TABLE_BASE_SQL.format(
        table_name=EVENTS_RECENT_DATA_TABLE(),
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(
            data_table=SHARDED_EVENTS_RECENT_DATA_TABLE(),
            sharding_key="sipHash64(distinct_id)",
            cluster=settings.DATASTORE_PRIMARY_REPLICA_CLUSTER,
        ),
        extra_fields=KAFKA_COLUMNS + INSERTED_AT_NOT_NULLABLE_COLUMN,
        dynamically_materialized_columns="",
        materialized_columns="",
        indexes="",
    )


def DISTRIBUTED_EVENTS_RECENT_TABLE_SQL(on_cluster=False):
    return EVENTS_TABLE_BASE_SQL.format(
        table_name="distributed_events_recent",
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(
            data_table=SHARDED_EVENTS_RECENT_DATA_TABLE(),
            sharding_key="sipHash64(distinct_id)",
            cluster=settings.DATASTORE_PRIMARY_REPLICA_CLUSTER,
        ),
        extra_fields=KAFKA_COLUMNS + INSERTED_AT_NOT_NULLABLE_COLUMN,
        dynamically_materialized_columns="",
        materialized_columns="",
        indexes="",
    )


def WRITABLE_EVENTS_RECENT_TABLE_SQL(on_cluster=False):
    return EVENTS_TABLE_BASE_SQL.format(
        table_name=WRITABLE_EVENTS_RECENT_TABLE(),
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(
            data_table=SHARDED_EVENTS_RECENT_DATA_TABLE(),
            sharding_key="sipHash64(distinct_id)",
            cluster=settings.DATASTORE_WRITABLE_CLUSTER,
        ),
        extra_fields=KAFKA_COLUMNS,
        dynamically_materialized_columns="",
        materialized_columns="",
        indexes="",
    )


def SHARDED_EVENTS_RECENT_TABLE_SQL():
    return (
        EVENTS_TABLE_BASE_SQL
        + """PARTITION BY toStartOfDay(inserted_at)
ORDER BY (team_id, toStartOfHour(inserted_at), event, cityHash64(distinct_id), cityHash64(uuid))
TTL toDateTime(inserted_at) + INTERVAL 7 DAY
{storage_policy}
SETTINGS ttl_only_drop_parts = 1
"""
    ).format(
        table_name=SHARDED_EVENTS_RECENT_DATA_TABLE(),
        on_cluster_clause=ON_CLUSTER_CLAUSE(False),
        engine=ReplacingMergeTree(
            "sharded_events_recent", ver="_timestamp", replication_scheme=ReplicationScheme.SHARDED
        ),
        extra_fields=KAFKA_COLUMNS + INSERTED_AT_NOT_NULLABLE_COLUMN,
        dynamically_materialized_columns="",
        materialized_columns="",
        indexes="",
        storage_policy=STORAGE_POLICY(),
    )


def SHARDED_EVENTS_RECENT_MV_SQL():
    return """
CREATE MATERIALIZED VIEW IF NOT EXISTS events_recent_json_mv
TO {database}.{target_table}
AS SELECT
uuid,
event,
properties,
timestamp,
team_id,
distinct_id,
elements_chain,
created_at,
person_id,
person_created_at,
person_properties,
group0_properties,
group1_properties,
group2_properties,
group3_properties,
group4_properties,
group0_created_at,
group1_created_at,
group2_created_at,
group3_created_at,
group4_created_at,
person_mode,
_timestamp,
_offset
FROM {database}.{source_table}
""".format(
        source_table=EVENTS_DATA_TABLE(),
        database=settings.DATASTORE_DATABASE,
        target_table=WRITABLE_EVENTS_RECENT_TABLE(),
    )


# Distributed engine tables are only created if DATASTORE_REPLICATED


# This table is responsible for writing to sharded_events based on a sharding key.
def WRITABLE_EVENTS_TABLE_SQL():
    return EVENTS_TABLE_BASE_SQL.format(
        table_name="writable_events",
        on_cluster_clause=ON_CLUSTER_CLAUSE(),
        engine=Distributed(data_table=EVENTS_DATA_TABLE(), sharding_key="sipHash64(distinct_id)"),
        extra_fields=KAFKA_COLUMNS + KAFKA_CONSUMER_BREADCRUMBS_COLUMN,
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns="",
        indexes="",
    )


# This table is responsible for reading from events on a cluster setting


def DISTRIBUTED_EVENTS_TABLE_SQL(on_cluster=True):
    return EVENTS_TABLE_BASE_SQL.format(
        table_name="events",
        on_cluster_clause=ON_CLUSTER_CLAUSE(on_cluster),
        engine=Distributed(data_table=EVENTS_DATA_TABLE(), sharding_key="sipHash64(distinct_id)"),
        extra_fields=KAFKA_COLUMNS + INSERTED_AT_COLUMN + KAFKA_CONSUMER_BREADCRUMBS_COLUMN,
        dynamically_materialized_columns=EVENTS_TABLE_DYNAMICALLY_MATERIALIZED_COLUMNS(),
        materialized_columns=EVENTS_TABLE_PROXY_MATERIALIZED_COLUMNS,
        indexes="",
    )


def INSERT_EVENT_SQL(table_name: str | None = None) -> str:
    if table_name is None:
        table_name = EVENTS_DATA_TABLE()

    return f"""
INSERT INTO {table_name}
(
    uuid,
    event,
    properties,
    timestamp,
    team_id,
    distinct_id,
    elements_chain,
    person_id,
    person_properties,
    person_created_at,
    group0_properties,
    group1_properties,
    group2_properties,
    group3_properties,
    group4_properties,
    group0_created_at,
    group1_created_at,
    group2_created_at,
    group3_created_at,
    group4_created_at,
    person_mode,
    created_at,
    _timestamp,
    _offset
)
VALUES
(
    %(uuid)s,
    %(event)s,
    %(properties)s,
    %(timestamp)s,
    %(team_id)s,
    %(distinct_id)s,
    %(elements_chain)s,
    %(person_id)s,
    %(person_properties)s,
    %(person_created_at)s,
    %(group0_properties)s,
    %(group1_properties)s,
    %(group2_properties)s,
    %(group3_properties)s,
    %(group4_properties)s,
    %(group0_created_at)s,
    %(group1_created_at)s,
    %(group2_created_at)s,
    %(group3_created_at)s,
    %(group4_created_at)s,
    %(person_mode)s,
    %(created_at)s,
    now(),
    0
)
"""


def BULK_INSERT_EVENT_SQL(table_name: str | None = None) -> str:
    if table_name is None:
        table_name = EVENTS_DATA_TABLE()

    return f"""
INSERT INTO {table_name}
(
    uuid,
    event,
    properties,
    timestamp,
    team_id,
    distinct_id,
    elements_chain,
    person_id,
    person_properties,
    person_created_at,
    group0_properties,
    group1_properties,
    group2_properties,
    group3_properties,
    group4_properties,
    group0_created_at,
    group1_created_at,
    group2_created_at,
    group3_created_at,
    group4_created_at,
    person_mode,
    created_at,
    _timestamp,
    _offset
)
VALUES
"""


NULL_SQL = """
-- Creates zero values for all date axis ticks for the given date_from, date_to range
SELECT toUInt16(0) AS total, {date_to_truncated} - {interval_func}(number) AS day_start

-- Get the number of `intervals` between date_from and date_to.
--
-- NOTE: for week there is some unusual behavior, see:
--       https://github.com/Datastore/Datastore/issues/7322
--
--       This actually aligns with what we want, as they are assuming Sunday week starts,
--       and we'd rather have the relative week num difference. Likewise the same for
--       "month" intervals
--
--       To ensure we get all relevant intervals, we add in the truncated "date_from"
--       value.
--
--       This behaviour of dateDiff is different to our handling of "week" and "month"
--       differences we are performing in python, which just considers seconds between
--       date_from and date_to
--
-- TODO: Ths pattern of generating intervals is repeated in several places. Reuse this
--       `ticks` query elsewhere.
FROM numbers(dateDiff(%(interval)s, {date_from_truncated}, toDateTime(%(date_to)s, %(timezone)s)))

UNION ALL

-- Make sure we capture the interval date_from falls into.
SELECT toUInt16(0) AS total, {date_from_truncated}
"""

EVENT_JOIN_PERSON_SQL = """
INNER JOIN ({GET_TEAM_PERSON_DISTINCT_IDS}) as pdi ON events.distinct_id = pdi.distinct_id
"""

GET_EVENTS_WITH_PROPERTIES = """
SELECT * FROM events WHERE
team_id = %(team_id)s
{filters}
{order_by}
"""

EXTRACT_TAG_REGEX = "extract(elements_chain, '^(.*?)[.|:]')"
EXTRACT_TEXT_REGEX = "extract(elements_chain, 'text=\"(.*?)\"')"

ELEMENT_TAG_COUNT = """
SELECT concat('<', {tag_regex}, '> ', {text_regex}) AS tag_name,
       events.elements_chain,
       count(*) as tag_count
FROM events
WHERE events.team_id = %(team_id)s AND event = '$autocapture'
GROUP BY tag_name, elements_chain
ORDER BY tag_count desc, tag_name
LIMIT %(limit)s
""".format(tag_regex=EXTRACT_TAG_REGEX, text_regex=EXTRACT_TEXT_REGEX)

GET_CUSTOM_EVENTS = """
SELECT DISTINCT event FROM events where team_id = %(team_id)s AND event NOT IN ['$autocapture', '$pageview', '$identify', '$pageleave', '$screen']
"""

#
# Demo data
#

COPY_EVENTS_BETWEEN_TEAMS = COPY_ROWS_BETWEEN_TEAMS_BASE_SQL.format(
    table_name=WRITABLE_EVENTS_DATA_TABLE(),
    columns_except_team_id="""uuid, event, properties, timestamp, distinct_id, elements_chain, created_at, person_id, person_created_at,
    person_properties, group0_properties, group1_properties, group2_properties, group3_properties, group4_properties,
     group0_created_at, group1_created_at, group2_created_at, group3_created_at, group4_created_at, person_mode""",
)