from typing import Dict, Callable

from conftest import (
    Cluster,
    find_routed_pk,
    Instance,
    TarantoolError,
)
import pytest
import requests  # type: ignore
from prometheus_client import Metric
from prometheus_client.samples import Sample
import psycopg


@pytest.mark.webui
def test_metrics_ok(instance: Instance) -> None:
    http_listen = instance.http_listen
    response = requests.get(f"http://{http_listen}/metrics")
    assert response.ok


# Verifies that all picodata metrics are present
# and accessible on /metrics endpoint
@pytest.mark.webui
def test_picodata_metrics(cluster: Cluster) -> None:
    instance = cluster.add_instance(name="i1", init_replication_factor=2)
    i2 = cluster.add_instance(name="i2")

    instance.wait_online()
    instance.sql("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT) DISTRIBUTED GLOBALLY")
    instance.sql("INSERT INTO test VALUES (1, 'one')")

    with pytest.raises(TarantoolError) as e:
        i2.sql("INSERT INTO test VALUES (1, 'one')")
        assert e.value.args[:2] == ("SbroadError", "sbroad: Lua error (IR dispatch)")

    i2.sql("INSERT INTO test VALUES (2, 'two')")
    with pytest.raises(TarantoolError) as e:
        instance.sql("INSERT INTO test VALUES (2, 'two')")
        assert e.value.args[:2] == ("SbroadError", "sbroad: Lua error (IR dispatch)")

    http_listen = instance.http_listen
    url = f"http://{http_listen}/metrics"
    response = requests.get(url)
    assert response.ok, f"Metrics endpoint {url} did not return OK: {response.status_code}"
    metrics_output = response.text
    expected_metrics = [
        "pico_governor_changes_total",
        "pico_sql_query_total",
        "pico_sql_query_errors_total",
        "pico_sql_query_duration",
        "pico_rpc_request_total",
        "pico_rpc_request_errors_total",
        "pico_rpc_request_duration",
        "pico_cas_records_total",
        "pico_cas_errors_total",
        "pico_cas_ops_duration",
        "pico_instance_state",
        "pico_raft_applied_index",
        "pico_raft_commit_index",
        "pico_raft_term",
        "pico_raft_state",
        "pico_raft_leader_id",
    ]

    for metric in expected_metrics:
        assert metric in metrics_output, f"Metric '{metric}' not found in /metrics output"


def check_metric(
    families: Dict[str, Metric], name: str, value: float | None, aggregate: Callable | None = None, **labels
):
    found_family = families.get(name)
    if value is None:
        if not labels:
            assert found_family is None, "Metric {} found".format(name)
        elif found_family is not None:
            for sample in found_family.samples:
                match_all_labels = all([sample.labels.get(lbl) == val for lbl, val in labels.items()])
                assert not match_all_labels, "Labeled {} metric found".format(name)
    else:
        assert found_family is not None, "Metric {} not found".format(name)
        if not labels:
            assert len(found_family.samples) == 1, "Metric has {} samples instead of 1".format(
                len(found_family.samples)
            )
            sample = found_family.samples[0]
            assert sample.value == value, "Metric has unexpected value"
        else:
            filtered_samples = []
            for sample in found_family.samples:
                if all([sample.labels.get(label) == value for label, value in labels.items()]):
                    filtered_samples.append(sample)
            if aggregate:
                real_value = aggregate(filtered_samples)
            else:
                assert len(filtered_samples) == 1, "Labeled metric has {} samples instead of 1".format(
                    len(filtered_samples)
                )
                sample = filtered_samples[0]
                real_value = sample.value
            assert real_value == value, "Labeled metric has unexpected value"


@pytest.mark.webui
def test_pgproto_metrics_collected(instance: Instance) -> None:
    check_metric(instance.get_metrics(), "pico_sql_query", None)

    # make a non-privileged user with iproto
    instance.sql("CREATE USER postgres WITH PASSWORD 'Passw0rd'")

    # creating a user for postgress will increase the metric value, so check it
    check_metric(instance.get_metrics(), "pico_sql_query", 1.0)

    # execute a query through postgresql
    host, port = instance.pg_host, instance.pg_port
    conn = psycopg.connect(f"postgres://postgres:Passw0rd@{host}:{port}")
    conn.autocommit = True  # do not make a transaction

    conn.execute("SELECT 1").fetchall()

    # the select above should've increase the metric value too
    check_metric(instance.get_metrics(), "pico_sql_query", 2.0)


@pytest.mark.webui
def test_router_and_storage_cache_metrics(instance: Instance):
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_misses", 0)

    ################
    # Run ACL query
    ################

    # Create a postgres user to test caches via pgproto
    instance.sql("CREATE USER postgres WITH PASSWORD 'Passw0rd'")
    host, port = instance.pg_host, instance.pg_port
    pgproto = psycopg.connect(f"postgres://postgres:Passw0rd@{host}:{port}")
    pgproto.autocommit = True  # do not make a transaction

    metrics = instance.get_metrics()

    # ACL is not cached
    check_metric(metrics, "pico_router_cache_statements_added", 0)
    check_metric(metrics, "pico_storage_cache_statements_added", 0)

    # Every query (even ACL) is looked in the router cache so miss is reported
    check_metric(metrics, "pico_router_cache_misses", 1)
    # ACL is not executed on storage so no requests are reported
    check_metric(metrics, "pico_storage_1st_requests", None)
    check_metric(metrics, "pico_storage_2nd_requests", None)

    ################
    # Run DDL query
    ################

    instance.sql("CREATE TABLE t (a INT PRIMARY KEY, b INT)")

    metrics = instance.get_metrics()

    # Every query (even DDL) is looked in the router cache so miss is reported
    check_metric(metrics, "pico_router_cache_misses", 2)
    # ACL is not executed on storage so no requests are reported
    check_metric(metrics, "pico_storage_1st_requests", None)
    check_metric(metrics, "pico_storage_2nd_requests", None)

    ################
    # Run DML query
    ################

    instance.sql("INSERT INTO t VALUES (1,2)")

    metrics = instance.get_metrics()

    # Every query is looked in the router cache so miss is reported
    check_metric(metrics, "pico_router_cache_misses", 3)
    # DML is cached on the router
    check_metric(metrics, "pico_router_cache_statements_added", 1)
    # DML is executed via local fast-path on single-rs buckets.
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dml", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None)

    ##########################
    # Execute first DQL query
    ##########################

    instance.sql("SELECT * FROM t")

    # DQL statement is cached
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_statements_added", 2)
    check_metric(metrics, "pico_storage_cache_statements_added", 1)

    # Router cache miss is reported
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_misses", 4)

    # DQL on all buckets is executed via local fast-path on a single-rs cluster.
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dql", result="ok")

    #####################################
    # Run the same DQL query via pgproto
    #####################################

    instance.sql("SELECT * FROM t")

    # Statement wasn't cached again
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_statements_added", 2)
    check_metric(metrics, "pico_storage_cache_statements_added", 1)

    # Cache hit is reported on the router
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 1)

    # Repeated all-buckets DQL is still local, so storage request counters stay untouched.
    check_metric(metrics, "pico_sql_local_query", 2, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dql", result="ok")

    ##############################################
    # Ensure storage cache evictions are reported
    ##############################################

    # Make storage cache full
    instance.sql("ALTER SYSTEM SET sql_storage_cache_count_max = 1")
    # Nothing evicted yet
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_storage_cache_statements_evicted", 0)
    check_metric(metrics, "pico_storage_cache_statements_added", 1)
    # Overflow the cache
    instance.sql("SELECT a + b FROM t")
    # Ensure eviction is reported
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_storage_cache_statements_evicted", 1)
    check_metric(metrics, "pico_storage_cache_statements_added", 2)

    #############################################
    # Ensure router cache evictions are reported
    #############################################

    # There is already 3 entries in the cache
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_statements_added", 3)
    # Make router cache full
    for i in range(0, 47):
        instance.sql(f"SELECT a + {i} FROM t")
    # Nothing evicted yet
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_statements_added", 50)
    check_metric(metrics, "pico_router_cache_statements_evicted", 0)
    # Overflow the cache
    instance.sql("SELECT a + b + 1 FROM t")
    # Ensure eviction is reported
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_statements_added", 51)
    check_metric(metrics, "pico_router_cache_statements_evicted", 1)

    ##########################################################
    # Ensure cache size is correctly reported after shrinking
    ##########################################################

    # Pollute storage cache
    instance.sql("ALTER SYSTEM SET sql_storage_cache_count_max = 4")
    instance.sql("SELECT (1) FROM t")
    instance.sql("SELECT (1), (2) FROM t")
    instance.sql("SELECT (1), (2), (3) FROM t")
    instance.sql("SELECT (1), (2), (3), (4) FROM t")

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_storage_cache_statements_added", 8)
    check_metric(metrics, "pico_storage_cache_statements_evicted", 4)

    # Shrink the cache using count_max parameter
    instance.sql("ALTER SYSTEM SET sql_storage_cache_count_max = 2")
    # size = added - evicted
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_storage_cache_statements_added", 8)
    check_metric(metrics, "pico_storage_cache_statements_evicted", 6)

    # Shrink the cache using size_max parameter
    instance.sql("ALTER SYSTEM SET sql_storage_cache_size_max = 1")
    # size = added - evicted
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_storage_cache_statements_added", 8)
    check_metric(metrics, "pico_storage_cache_statements_evicted", 8)

    # Still can do some queries
    instance.sql("SELECT * FROM t JOIN t t2 on t.a = t2.a")


@pytest.mark.webui
def test_temp_table_lock_metrics(instance: Instance) -> None:
    instance.sql("CREATE TABLE temp_metrics (id INTEGER PRIMARY KEY, name TEXT)")
    instance.sql("INSERT INTO temp_metrics VALUES (1, 'one')")

    metrics = instance.get_metrics()
    lease_metric = metrics.get("pico_sql_temp_table_leases")
    assert lease_metric is not None, "pico_sql_temp_table_leases not found"
    assert len(lease_metric.samples) == 1
    initial_leases = lease_metric.samples[0].value

    waits_metric = metrics.get("pico_sql_temp_table_lock_waits")
    assert waits_metric is not None, "pico_sql_temp_table_lock_waits not found"
    assert len(waits_metric.samples) == 1
    initial_waits = waits_metric.samples[0].value

    instance.sql("SELECT * FROM temp_metrics")

    metrics = instance.get_metrics()
    lease_metric = metrics.get("pico_sql_temp_table_leases")
    assert lease_metric is not None, "pico_sql_temp_table_leases not found"
    assert lease_metric.samples[0].value >= initial_leases + 1

    waits_metric = metrics.get("pico_sql_temp_table_lock_waits")
    assert waits_metric is not None, "pico_sql_temp_table_lock_waits not found"
    assert waits_metric.samples[0].value == initial_waits


def get_instance_state_metric(instance: Instance) -> Dict[str, Sample]:
    metric = instance.get_metrics()["pico_instance_state"]
    samples_by_instance_name = {}
    for sample in metric.samples:
        samples_by_instance_name[sample.labels["instance"]] = sample

    return samples_by_instance_name


def all_but_first_report_online(cluster: Cluster, offline_name):
    for instance_to_fetch_metric in cluster.instances[1:]:
        samples_by_instance_name = get_instance_state_metric(instance_to_fetch_metric)

        for instance_to_check_state in cluster.instances:
            assert instance_to_check_state.name is not None  # mypy
            sample = samples_by_instance_name[instance_to_check_state.name]
            if instance_to_check_state.name == offline_name:
                assert sample.labels["state"] == "Offline"
            else:
                assert sample.labels["state"] == "Online"


def test_instance_state_metric(cluster: Cluster):
    i1, i2, i3 = cluster.deploy(instance_count=3)

    # check after startup all instances report all other instances as Online
    for instance in cluster.instances:
        samples_by_instance_name = get_instance_state_metric(instance)

        for instance in cluster.instances:
            assert instance.name is not None  # mypy
            sample = samples_by_instance_name.get(instance.name)
            assert sample is not None
            assert sample.labels["tier"] == "default"
            assert sample.labels["state"] == "Online"

    # shut down one instance, ensure others see it as offline.
    # With check_quorum, advance the raft clock to trigger immediate election.
    i1.terminate()
    for i in [i2, i3]:
        i.call("pico.raft_tick", 100)

    cluster.wait_leader_elected([i2, i3])

    all_but_first_report_online(cluster, i1.name)

    i1.start_and_wait()

    # check that after instance became online again we see it everywhere as online
    for instance_with_metrics in cluster.instances:
        samples_by_instance_name = get_instance_state_metric(instance_with_metrics)

        for instance in cluster.instances:
            assert instance.name is not None  # mypy
            sample = samples_by_instance_name.get(instance.name)
            assert sample is not None, f"Missing {instance.name} in {instance_with_metrics} metrics"
            assert sample.labels["tier"] == "default"
            assert sample.labels["state"] == "Online"

    # make it offline again to check in snapshot later
    i1.terminate()

    all_but_first_report_online(cluster, i1.name)

    # resest raft_wal_count_max to trigger snapshot transfer
    i3.sql("ALTER SYSTEM SET raft_wal_count_max = 1")
    assert i3.call("box.space._raft_log:len") == 0

    # add new instance to trigger snapshot transfer, to check that metrics are updated with snapshot
    i4 = cluster.add_instance(peers=[i2.iproto_listen])

    for instance in cluster.instances[1:]:
        samples_by_instance_name = get_instance_state_metric(instance)
        assert i4.name in samples_by_instance_name

        all_but_first_report_online(cluster, i1.name)


@pytest.mark.webui
def test_storage_cache_mestrics(cluster: Cluster):
    cluster.set_config_file(
        yaml="""
cluster:
    name: test
    tier:
        router:
            replication_factor: 1
            can_vote: true
        storage:
            replication_factor: 1
            can_vote: true
"""
    )

    router = cluster.add_instance(wait_online=True, tier="router", enable_http=True)
    storage_1 = cluster.add_instance(wait_online=True, tier="storage", enable_http=True)
    storage_2 = cluster.add_instance(wait_online=True, tier="storage", enable_http=True)

    # All metrics uninitialized on storage nodes.
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")

    # Create a postgres user to test metrics via pgproto
    router.sql("CREATE USER postgres WITH PASSWORD 'Passw0rd'")
    host, port = router.pg_host, router.pg_port
    pgproto = psycopg.connect(f"postgres://postgres:Passw0rd@{host}:{port}")
    pgproto.autocommit = True  # do not make a transaction
    router.sql("GRANT CREATE TABLE TO postgres", sudo=True)

    # ACL does not access the storage cache.
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")

    # Create and fill table
    pgproto.execute("CREATE TABLE t (a INT PRIMARY KEY, b INT) DISTRIBUTED BY (a) IN TIER storage")
    pgproto.execute("INSERT INTO t VALUES (1,2)")
    pgproto.execute("INSERT INTO t VALUES (2,3)")

    # DDL and DML do not access the storage cache.
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")

    # DQL true cache misses are reported
    pgproto.execute("SELECT a FROM t")
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="1st", miss_type="true")
        check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="2nd", miss_type="true")

    # DQL cache hits are reported
    pgproto.execute("SELECT a FROM t")
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", 1, query_type="dql", rpc_type="1st")

    # DQL stale cache misses are reported after schema changes.
    pgproto.execute("CREATE TABLE t2 (a INT PRIMARY KEY, b INT) DISTRIBUTED BY (a) IN TIER storage")
    pgproto.execute("SELECT a FROM t")
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", 1, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="1st", miss_type="stale")

    # Recompilation due to schema changes is reported.
    pgproto.execute("CREATE INDEX ib ON t(b)")
    pgproto.execute("SELECT a FROM t")
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", 1, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="1st", miss_type="stale")

    # ALTER TABLE changes picodata's schema version leading to
    # complete recompilation of the plan and the statement from scratch.
    pgproto.execute("ALTER TABLE t ADD COLUMN c INT")
    pgproto.execute("SELECT a FROM t")
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_hits", 1, query_type="dql", rpc_type="1st")
        check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="1st", miss_type="true")
        check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="2nd", miss_type="true")
        check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="local")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="true")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")

    # `pico_storage_cache_busy_misses_total` cannot be reported at the moment
    # due to the lock on the storage cache.
    for instance in (storage_1, storage_2):
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="busy")
        check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="2nd", miss_type="busy")


@pytest.mark.webui
def test_local_sql_collisions_gl_2367(cluster: Cluster):
    i1 = cluster.add_instance(name="i1", enable_http=True)
    i2 = cluster.add_instance(name="i2", enable_http=True)

    # Create a postgres user to test metrics via pgproto
    i1.sql("CREATE USER postgres WITH PASSWORD 'Passw0rd'")
    host, port = i1.pg_host, i1.pg_port
    pgproto = psycopg.connect(f"postgres://postgres:Passw0rd@{host}:{port}")
    pgproto.autocommit = True  # do not make a transaction
    i1.sql("GRANT CREATE TABLE TO postgres", sudo=True)

    # Create and fill table
    pgproto.execute("CREATE TABLE t (a INT PRIMARY KEY, b INT)")
    pgproto.execute("INSERT INTO t VALUES (1,2)")

    # Two instances create two replica sets, so the query uses both protocol and local paths.
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")

    # Assert local SQL collides.
    raw1 = pgproto.execute("EXPLAIN (raw) SELECT DISTINCT * FROM t ORDER BY 1").fetchall()
    raw2 = pgproto.execute("EXPLAIN (raw) SELECT * FROM t GROUP BY a, b ORDER BY 1").fetchall()
    assert raw1[1] == raw2[1]

    # EXPLAIN shouldn't be reported.
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="true")

    # Cache statements.
    assert pgproto.execute("SELECT DISTINCT * FROM t ORDER BY 1").fetchall() == [(1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="1st", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="local", miss_type="true")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="1st", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", 1, query_type="dql", rpc_type="2nd", miss_type="true")

    # Cache another statements with the same local SQL.
    assert pgproto.execute("SELECT * FROM t GROUP BY a, b ORDER BY 1").fetchall() == [(1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")
    check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="local", miss_type="true")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 1, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")

    # Run cached statements, temporary spaces creation does not matter.
    assert pgproto.execute("SELECT DISTINCT * FROM t ORDER BY 1").fetchall() == [(1, 2)]
    assert pgproto.execute("SELECT * FROM t GROUP BY a, b ORDER BY 1").fetchall() == [(1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 3, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_hits", 2, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 3, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")

    # Run cached statements again.
    assert pgproto.execute("SELECT DISTINCT * FROM t ORDER BY 1").fetchall() == [(1, 2)]
    assert pgproto.execute("SELECT * FROM t GROUP BY a, b ORDER BY 1").fetchall() == [(1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 5, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_hits", 4, query_type="dql", rpc_type="local")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 5, query_type="dql", rpc_type="1st")

    # Run cached statements one more time.
    assert pgproto.execute("SELECT DISTINCT * FROM t ORDER BY 1").fetchall() == [(1, 2)]
    assert pgproto.execute("SELECT * FROM t GROUP BY a, b ORDER BY 1").fetchall() == [(1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 7, query_type="dql", rpc_type="1st")
    check_metric(metrics, "pico_storage_cache_hits", 6, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", 7, query_type="dql", rpc_type="1st")

    assert pgproto.execute(
        "SELECT * FROM t WHERE a < %s UNION ALL SELECT * FROM t WHERE a < %s and b < %s", (2, 3, 4)
    ).fetchall() == [(1, 2), (1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="1st", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="local", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="2nd", miss_type="true")

    assert pgproto.execute(
        "SELECT * FROM t WHERE a < %s and b < %s UNION ALL SELECT * FROM t WHERE a < %s", (2, 3, 4)
    ).fetchall() == [(1, 2), (1, 2)]
    metrics = i1.get_metrics()
    check_metric(metrics, "pico_storage_cache_misses", 3, query_type="dql", rpc_type="1st", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", 2, query_type="dql", rpc_type="local", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")
    metrics = i2.get_metrics()
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="1st", miss_type="stale")
    check_metric(metrics, "pico_storage_cache_misses", 3, query_type="dql", rpc_type="2nd", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="true")
    check_metric(metrics, "pico_storage_cache_misses", None, query_type="dql", rpc_type="local", miss_type="stale")


@pytest.mark.webui
def test_plan_id_sql_collisions(instance: Instance):
    def sum_samples(samples: list[Sample]) -> float:
        return sum(map(lambda s: s.value, samples))

    instance.sql("CREATE TABLE t (a INT PRIMARY KEY, b INT)")
    instance.sql("CREATE TABLE g (a INT PRIMARY KEY, b INT) DISTRIBUTED GLOBALLY")

    dml = instance.sql("INSERT INTO t VALUES (1,2),(2,3),(3,4)")
    assert dml["row_count"] == 3
    dml = instance.sql("INSERT INTO g VALUES (3,1),(2,2),(1,4)")
    assert dml["row_count"] == 3

    dql = instance.sql("SELECT * FROM t GROUP BY a, b ORDER BY 1")  # top + motion = 2
    assert dql == [[1, 2], [2, 3], [3, 4]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 0)
    check_metric(metrics, "pico_router_cache_misses", 5)
    check_metric(metrics, "pico_storage_cache_hits", None, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 3, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dql = instance.sql("SELECT * FROM t GROUP BY a, b ORDER BY 1")
    assert dql == [[1, 2], [2, 3], [3, 4]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 1)
    check_metric(metrics, "pico_router_cache_misses", 5)
    check_metric(metrics, "pico_storage_cache_hits", 2, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 3, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dql = instance.sql("SELECT * FROM g GROUP BY a, b ORDER BY 1")  # top = 1
    assert dql == [[1, 4], [2, 2], [3, 1]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 1)
    check_metric(metrics, "pico_router_cache_misses", 6)
    check_metric(metrics, "pico_storage_cache_hits", 2, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 4, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dql = instance.sql("SELECT * FROM g GROUP BY a, b ORDER BY 1")
    assert dql == [[1, 4], [2, 2], [3, 1]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 2)
    check_metric(metrics, "pico_router_cache_misses", 6)
    check_metric(metrics, "pico_storage_cache_hits", 3, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 4, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dml = instance.sql("INSERT INTO t VALUES($1, 2+$1)", 5)  # motion + values = 2
    assert dml["row_count"] == 1

    dql = instance.sql("SELECT * FROM t UNION SELECT $2-$1, 5 ORDER BY b DESC", 5, 9)  # top + motion = 2
    assert dql == [[5, 7], [4, 5], [3, 4], [2, 3], [1, 2]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 2)
    check_metric(metrics, "pico_router_cache_misses", 8)
    check_metric(metrics, "pico_storage_cache_hits", 3, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 7, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dml = instance.sql("UPDATE t SET b = b + 2 WHERE a >= 3+2")  # top = 1
    assert dml["row_count"] == 1
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 2)
    check_metric(metrics, "pico_router_cache_misses", 9)
    check_metric(metrics, "pico_storage_cache_hits", 3, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 8, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dml = instance.sql("UPDATE t SET b = b + 2 WHERE a >= 3+2")  # top = 1
    assert dml["row_count"] == 1
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 3)
    check_metric(metrics, "pico_router_cache_misses", 9)
    check_metric(metrics, "pico_storage_cache_hits", 4, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 8, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dml = instance.sql("DELETE FROM g WHERE b = ?", 1)  # motion = 1
    assert dml["row_count"] == 1
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 3)
    check_metric(metrics, "pico_router_cache_misses", 10)
    check_metric(metrics, "pico_storage_cache_hits", 4, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 9, aggregate=sum_samples, query_type="dql", rpc_type="local")

    dql = instance.sql("SELECT * FROM t JOIN g ON t.a = g.b OR t.a = 5")  # top = 1
    assert dql == [[5, 11, 1, 4], [2, 3, 2, 2], [5, 11, 2, 2]]
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_router_cache_hits", 3)
    check_metric(metrics, "pico_router_cache_misses", 11)
    check_metric(metrics, "pico_storage_cache_hits", 4, query_type="dql", rpc_type="local")
    check_metric(metrics, "pico_storage_cache_misses", 10, aggregate=sum_samples, query_type="dql", rpc_type="local")


@pytest.mark.webui
def test_filtered_local_dql_bypasses_iproto(instance: Instance):
    instance.sql("""CREATE TABLE t (a INT NOT NULL, b INT, PRIMARY KEY (a)) DISTRIBUTED BY (a)""")
    instance.sql("""INSERT INTO t VALUES (1, 2)""")

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", None, query_type="dql", result="ok")

    assert instance.sql("""SELECT * FROM t WHERE a = 1""") == [[1, 2]]

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dql", result="ok")

    assert instance.sql("""SELECT * FROM t WHERE a = 1""") == [[1, 2]]

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 2, query_type="dql", result="ok")


@pytest.mark.webui
def test_all_buckets_local_dql_bypasses_iproto(instance: Instance):
    instance.sql("""CREATE TABLE t (a INT NOT NULL, b INT, PRIMARY KEY (a)) DISTRIBUTED BY (a)""")
    instance.sql("""INSERT INTO t VALUES (1, 2), (2, 3)""")

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dql", result="ok")

    assert instance.sql("""SELECT * FROM t ORDER BY a""") == [[1, 2], [2, 3]]

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dql", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dql", result="ok")


@pytest.mark.webui
def test_replica_local_dql_bypasses_iproto(instance: Instance):
    leader = instance
    assert leader.cluster is not None
    replica = leader.cluster.add_instance(wait_online=True, replicaset_name=leader.replicaset_name, enable_http=True)
    assert replica.replicaset_master_name() == leader.name

    leader.sql("""CREATE TABLE t (a INT NOT NULL, b INT, PRIMARY KEY (a)) DISTRIBUTED BY (a)""")
    leader.sql("""INSERT INTO t VALUES (1, 2), (2, 3)""")

    query = """
        SELECT pico_instance_name(pico_instance_uuid()), a, b
        FROM t
        WHERE a = 1
        OPTION (read_preference = replica)
    """

    leader_metrics = leader.get_metrics()
    replica_metrics = replica.get_metrics()
    check_metric(leader_metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(leader_metrics, "pico_sql_replicas_read", None)
    check_metric(replica_metrics, "pico_sql_replicas_read", None)
    check_metric(replica_metrics, "pico_storage_1st_requests", None, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_storage_2nd_requests", None, query_type="dql", result="ok")

    # A leader cannot use the local fast-path for read_preference=replica.
    assert leader.sql(query) == [[replica.name, 1, 2]]

    leader_metrics = leader.get_metrics()
    replica_metrics = replica.get_metrics()
    check_metric(leader_metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(leader_metrics, "pico_sql_replicas_read", None)
    check_metric(replica_metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_sql_replicas_read", 1)
    check_metric(replica_metrics, "pico_storage_1st_requests", 1, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_storage_2nd_requests", 1, query_type="dql", result="ok")

    # The same query on the replica stays on the current instance and bypasses iproto.
    assert replica.sql(query) == [[replica.name, 1, 2]]

    leader_metrics = leader.get_metrics()
    replica_metrics = replica.get_metrics()
    check_metric(leader_metrics, "pico_sql_local_query", None, query_type="dql", result="ok")
    check_metric(leader_metrics, "pico_sql_replicas_read", None)
    check_metric(replica_metrics, "pico_sql_local_query", 1, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_sql_replicas_read", 1)
    check_metric(replica_metrics, "pico_storage_1st_requests", 1, query_type="dql", result="ok")
    check_metric(replica_metrics, "pico_storage_2nd_requests", 1, query_type="dql", result="ok")


@pytest.mark.webui
def test_filtered_local_dml_bypasses_iproto(instance: Instance):
    instance.sql("""CREATE TABLE t (a INT NOT NULL, b INT, PRIMARY KEY (a)) DISTRIBUTED BY (a)""")

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", None, query_type="dml", result="ok")

    instance.sql("""INSERT INTO t VALUES (1, 2)""")
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dml", result="ok")
    assert instance.sql("""SELECT * FROM t WHERE a = 1""") == [[1, 2]]

    instance.sql("""UPDATE t SET b = 3 WHERE a = 1""")
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 2, query_type="dml", result="ok")
    assert instance.sql("""SELECT * FROM t WHERE a = 1""") == [[1, 3]]

    instance.sql("""DELETE FROM t WHERE a = 1""")
    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 3, query_type="dml", result="ok")
    assert instance.sql("""SELECT * FROM t WHERE a = 1""") == []


@pytest.mark.webui
def test_filtered_local_block_bypasses_iproto(cluster: Cluster):
    leader, *_ = cluster.deploy(instance_count=2, enable_http=True)
    cluster.wait_until_buckets_balanced()

    leader.sql("CREATE USER postgres WITH PASSWORD 'Passw0rd'")
    leader.sql("GRANT CREATE TABLE TO postgres", sudo=True)

    host, port = leader.pg_host, leader.pg_port
    pgproto = psycopg.connect(f"postgres://postgres:Passw0rd@{host}:{port}")
    pgproto.autocommit = True  # do not make a transaction

    local_pk = find_routed_pk(leader, is_local=True)

    with pgproto.cursor() as cur:
        cur.execute("CREATE TABLE t (pk INTEGER PRIMARY KEY, a INTEGER)")
        cur.execute("INSERT INTO t (pk, a) VALUES (%s, %s)", (local_pk, 10))

        for instance in cluster.instances:
            metrics = instance.get_metrics()
            check_metric(metrics, "pico_storage_1st_requests", None, query_type="block", result="ok")
            check_metric(metrics, "pico_storage_2nd_requests", None, query_type="block", result="ok")

        cur.execute(
            """
            DO $$ BEGIN
                RETURN QUERY SELECT a FROM t WHERE pk = %s;
                UPDATE t SET a = a + 1 WHERE pk = %s;
            END $$;
            """,
            (local_pk, local_pk),
        )
        assert cur.fetchall() == [(10,)]

    for instance in cluster.instances:
        metrics = instance.get_metrics()
        check_metric(metrics, "pico_storage_1st_requests", None, query_type="block", result="ok")
        check_metric(metrics, "pico_storage_2nd_requests", None, query_type="block", result="ok")

    assert leader.sql(f"SELECT * FROM t WHERE pk = {local_pk}") == [[local_pk, 11]]


@pytest.mark.webui
def test_all_buckets_local_dml_bypasses_iproto(instance: Instance):
    instance.sql("""CREATE TABLE t (a INT NOT NULL, b INT, PRIMARY KEY (a)) DISTRIBUTED BY (a)""")
    instance.sql("""INSERT INTO t VALUES (1, 2), (2, 3)""")

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 1, query_type="dml", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dml", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dml", result="ok")

    dml = instance.sql("""UPDATE t SET b = b + 10""")
    assert dml["row_count"] == 2

    metrics = instance.get_metrics()
    check_metric(metrics, "pico_sql_local_query", 2, query_type="dml", result="ok")
    check_metric(metrics, "pico_storage_1st_requests", None, query_type="dml", result="ok")
    check_metric(metrics, "pico_storage_2nd_requests", None, query_type="dml", result="ok")
    assert instance.sql("""SELECT * FROM t ORDER BY a""") == [[1, 12], [2, 13]]
