A db migration does not take Query#apply_auto_limit in count about redash HOT 4 OPEN

"""Make case insensitive hash of query text

Revision ID: 1038c2174f5d
Revises: fd4fc850d7ea
Create Date: 2023-07-16 23:10:12.885949

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table

from redash.utils import gen_query_hash

from redash.query_runner import BaseSQLQueryRunner

import json

# revision identifiers, used by Alembic.
revision = '1038c2174f5d'
down_revision = 'fd4fc850d7ea'
branch_labels = None
depends_on = None


def query_hash(query_text, options):
    options_json = json.loads(options)
    apply_auto_limit = options_json.get('apply_auto_limit', False)
    runner = BaseSQLQueryRunner({})
    limit_aware_query = runner.apply_auto_limit(query_text, apply_auto_limit)
    return gen_query_hash(limit_aware_query)


def change_query_hash(conn, table, query_text_to):
    for record in conn.execute(table.select()):
        query_text = query_text_to(record.query)
        options = record.options
        conn.execute(
            table
            .update()
            .where(table.c.id == record.id)
            .values(query_hash=query_hash(query_text, options)))

def upgrade():
    queries = table(
        'queries',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('query', sa.Text),
        sa.Column('query_hash', sa.String(length=10)),
        sa.Column('options', sa.Text))

    conn = op.get_bind()
    change_query_hash(conn, queries, query_text_to=str)


def downgrade():
    queries = table(
        'queries',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('query', sa.Text),
        sa.Column('query_hash', sa.String(length=10)),
        sa.Column('options', sa.Text))


    conn = op.get_bind()
    change_query_hash(conn, queries, query_text_to=str.lower)

def downgrade():
    # This is irreversible mainly because of Snowflake queries which have a grayed out "auto-limit" option in 10.1.0.
    # But in a "preview" version it is available and Snowflake queries suddenly become "auto-limited"
    # and their hashes becomes different.
    # Also it would take a hash correction because BaseSQLQueryRunner's apply_auto_limit in "upgrade branch" has
    # own opinion on semicolons.
    # To sum up - it is easier to just restore a dump than mess with it.
    raise Exception("Irreversible migration")

"""Make case insensitive hash of query text

Revision ID: 1038c2174f5d
Revises: fd4fc850d7ea
Create Date: 2023-07-16 23:10:12.885949

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table

from redash.utils import gen_query_hash

from redash.query_runner import BaseSQLQueryRunner

import json
import sqlparse

# revision identifiers, used by Alembic.
revision = '1038c2174f5d'
down_revision = 'fd4fc850d7ea'
branch_labels = None
depends_on = None

def is_empty_statement(stmt):
    # copy statement object. `copy.deepcopy` fails to do this, so just re-parse it
    st = sqlparse.engine.FilterStack()
    st.stmtprocess.append(sqlparse.filters.StripCommentsFilter())
    stmt = next(st.run(str(stmt)), None)
    if stmt is None:
        return True
    return str(stmt).strip() == ""


def query_hash(query_text, options):
    options_json = json.loads(options)
    apply_auto_limit = options_json.get('apply_auto_limit', False)
    runner = BaseSQLQueryRunner({})
    limit_aware_query = runner.apply_auto_limit(query_text, apply_auto_limit)
    return gen_query_hash(limit_aware_query)


def change_query_hash(conn, table, query_text_to):
    for record in conn.execute(table.select()):
        if is_empty_statement(record.query):
            continue
        query_text = query_text_to(record.query)
        options = record.options
        conn.execute(
            table
            .update()
            .where(table.c.id == record.id)
            .values(query_hash=query_hash(query_text, options)))


def upgrade():
    queries = table(
        'queries',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('query', sa.Text),
        sa.Column('query_hash', sa.String(length=10)),
        sa.Column('options', sa.Text))

    conn = op.get_bind()
    change_query_hash(conn, queries, query_text_to=str)


def downgrade():
    # This is irreversible mainly because of Snowflake queries which have a grayed out "auto-limit" option in 10.1.0.
    # But in a "preview" version it is available and Snowflake queries suddenly become "auto-limited"
    # and their hashes becomes different.
    # Also it would take a hash correction because BaseSQLQueryRunner's apply_auto_limit in "upgrade branch" has
    # own opinion on semicolons.
    # To sum up - it is easier to just restore a dump than mess with it.
    raise Exception("Irreversible migration")