Compare commits

...
Author SHA1 Message Date
Elizabeth Thompson 88cc406040 fix(import): catch KeyError for missing uuid/ssh_tunnel in load_configs
A hand-edited or third-party database export YAML that omits the `uuid`
key (or `ssh_tunnel`) hit `config["uuid"]`/`config["ssh_tunnel"]`
indexing in load_configs() before schema.load() ran, raising a raw
KeyError that escaped the enclosing `except ValidationError` and
surfaced as an opaque 500 from the *//import/ endpoints instead of a
clean validation error. Add a sibling `except KeyError` that logs and
appends a ValidationError, routing the failure into the same aggregated
per-file error path as every other validation failure.
2026-09-05 22:23:03 +00:00
fe1b368bcd feat: add dashed line support for derived series (timeseries comparison) on MixedTimeseries chart (#34794)
Co-authored-by: Evan <evan@preset.io>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-09-05 14:17:10 -07:00
Nguyen Van VietandClaude Fable 5 1358543827 fix(sqllab): preserve ClickHouse LIMIT BY when applying the row limit (#43578)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-09-05 09:51:00 -07:00
6 changed files with 617 additions and 13 deletions
@@ -44,9 +44,14 @@ import {
ValueFormatter,
} from '@superset-ui/core';
import { GenericDataType } from '@apache-superset/core/common';
import { getOriginalSeries } from '@superset-ui/chart-controls';
import {
getOriginalSeries,
getTimeOffset,
isDerivedSeries,
} from '@superset-ui/chart-controls';
import type { EChartsCoreOption } from 'echarts/core';
import type { SeriesOption } from 'echarts';
import type { LineStyleOption } from 'echarts/types/src/util/types';
import {
DEFAULT_FORM_DATA,
EchartsMixedTimeseriesChartTransformedProps,
@@ -100,7 +105,11 @@ import {
transformSeries,
transformTimeseriesAnnotation,
} from '../Timeseries/transformers';
import { TIMEGRAIN_TO_TIMESTAMP, TIMESERIES_CONSTANTS } from '../constants';
import {
TIMEGRAIN_TO_TIMESTAMP,
TIMESERIES_CONSTANTS,
OpacityEnum,
} from '../constants';
import { getDefaultTooltip } from '../utils/tooltip';
import {
createSpacedXAxisFormatter,
@@ -453,6 +462,10 @@ export default function transformProps(
const array = ensureIsArray(chartProps.rawFormData?.time_compare);
const inverted = invert(verboseMap);
// Tracks a stable pattern index per time offset so that derived series
// sharing the same comparison window (across both queries A and B) get
// the same dash pattern, mirroring the regular Timeseries transform.
const offsetPatterns: { [key: string]: number } = {};
// The rendered ECharts series names are display names that can diverge from
// the backend `label_map` keys: the metric display name is prepended when
@@ -467,6 +480,22 @@ export default function transformProps(
rawSeriesA.forEach(entry => {
const entryName = String(entry.name || '');
const seriesName = inverted[entryName] || entryName;
const derivedSeries = isDerivedSeries(
entry,
chartProps.rawFormData,
seriesName,
);
const lineStyle: LineStyleOption = {};
if (derivedSeries && timeShiftColor) {
const offset = getTimeOffset(entry, array) || seriesName;
if (!offsetPatterns[offset]) {
offsetPatterns[offset] = Object.keys(offsetPatterns).length + 1;
}
const patternIndex = offsetPatterns[offset];
// use a combination of dash and dot for the line style
lineStyle.type = [(patternIndex % 5) + 1, (patternIndex % 3) + 1];
lineStyle.opacity = OpacityEnum.DerivedSeries;
}
const colorScaleKey = getOriginalSeries(seriesName, array);
const labelMapValues = rawLabelMap?.[seriesName];
@@ -544,6 +573,7 @@ export default function transformProps(
timeShiftColor,
theme,
labelPosition,
lineStyle,
},
);
@@ -556,6 +586,23 @@ export default function transformProps(
rawSeriesB.forEach(entry => {
const entryName = String(entry.name || '');
const seriesEntry = inverted[entryName] || entryName;
const derivedSeries = isDerivedSeries(
entry,
chartProps.rawFormData,
seriesEntry,
);
const lineStyle: LineStyleOption = {};
if (derivedSeries && timeShiftColor) {
const offset = getTimeOffset(entry, array) || seriesEntry;
if (!offsetPatterns[offset]) {
offsetPatterns[offset] = Object.keys(offsetPatterns).length + 1;
}
const patternIndex = offsetPatterns[offset];
// use a combination of dash and dot for the line style
lineStyle.type = [(patternIndex % 5) + 1, (patternIndex % 3) + 1];
lineStyle.opacity = OpacityEnum.DerivedSeries;
}
const colorScaleKey = getOriginalSeries(seriesEntry, array);
const labelMapValuesB = rawLabelMapB?.[seriesEntry];
@@ -634,6 +681,7 @@ export default function transformProps(
timeShiftColor,
theme,
labelPosition: labelPositionB,
lineStyle,
},
);
@@ -21,6 +21,7 @@ import {
AnnotationType,
AnnotationSourceType,
AxisType,
ComparisonType,
DataRecord,
FormulaAnnotationLayer,
IntervalAnnotationLayer,
@@ -1562,6 +1563,7 @@ test('y-axis title position: non-Left sets nameLocation to end', () => {
expect(yAxis[1].nameGap).toEqual(30);
expect(yAxis[1].nameLocation).toEqual('end');
});
describe('EchartsMixedTimeseries tooltip truncation', () => {
const longSeriesName = 'prod-us-east-1-service-checkout-latency-p99';
const marker = '<span style="background-color:#1f77b4;"></span>';
@@ -1763,3 +1765,162 @@ test('hides the ticks on the x axis and both y axes', () => {
expect(yAxis[0].axisTick.show).toBe(false);
expect(yAxis[1].axisTick.show).toBe(false);
});
test('should apply a dashed lineStyle to derived (time comparison) series only', () => {
const queryAData = createTestQueryData(
[
{
sum__num: 100,
'sum__num__1 week ago': 80,
ds: 599616000000,
},
{
sum__num: 150,
'sum__num__1 week ago': 120,
ds: 599916000000,
},
],
{
label_map: {
ds: ['ds'],
sum__num: ['sum__num'],
'sum__num__1 week ago': ['sum__num__1 week ago'],
},
},
);
const chartProps = createEchartsTimeseriesTestChartProps<
EchartsMixedTimeseriesFormData,
EchartsMixedTimeseriesProps
>({
...MIXED_TIMESERIES_CHART_PROPS_DEFAULTS,
defaultQueriesData: [queryAData, queriesData[1]],
formData: {
...formData,
metrics: ['sum__num'],
groupby: [],
time_compare: ['1 week ago'],
comparison_type: ComparisonType.Values,
timeShiftColor: true,
},
queriesData: [queryAData, queriesData[1]],
});
const transformed = transformProps(chartProps);
const series = (transformed.echartOptions.series as SeriesOption[]) || [];
const mainSeries = series.find(s => s.name === 'sum__num') as
| (SeriesOption & { lineStyle?: { type?: number[] | string } })
| undefined;
const derivedSeries = series.find(s => s.name === 'sum__num__1 week ago') as
| (SeriesOption & { lineStyle?: { type?: number[] | string } })
| undefined;
expect(mainSeries).toBeDefined();
expect(derivedSeries).toBeDefined();
// The primary (non-derived) series should not receive a dash pattern
expect(mainSeries?.lineStyle?.type).toBeUndefined();
// The derived (time comparison) series should receive a dash pattern array
expect(Array.isArray(derivedSeries?.lineStyle?.type)).toBe(true);
});
test('should not apply a dashed lineStyle when comparison_type is not Values', () => {
const queryAData = createTestQueryData(
[
{
sum__num: 100,
'sum__num__1 week ago': 80,
ds: 599616000000,
},
{
sum__num: 150,
'sum__num__1 week ago': 120,
ds: 599916000000,
},
],
{
label_map: {
ds: ['ds'],
sum__num: ['sum__num'],
'sum__num__1 week ago': ['sum__num__1 week ago'],
},
},
);
const chartProps = createEchartsTimeseriesTestChartProps<
EchartsMixedTimeseriesFormData,
EchartsMixedTimeseriesProps
>({
...MIXED_TIMESERIES_CHART_PROPS_DEFAULTS,
defaultQueriesData: [queryAData, queriesData[1]],
formData: {
...formData,
metrics: ['sum__num'],
groupby: [],
time_compare: ['1 week ago'],
comparison_type: ComparisonType.Difference,
},
queriesData: [queryAData, queriesData[1]],
});
const transformed = transformProps(chartProps);
const series = (transformed.echartOptions.series as SeriesOption[]) || [];
const derivedSeries = series.find(s => s.name === 'sum__num__1 week ago') as
| (SeriesOption & { lineStyle?: { type?: number[] | string } })
| undefined;
expect(derivedSeries).toBeDefined();
expect(derivedSeries?.lineStyle?.type).toBeUndefined();
});
test('should not apply a dashed lineStyle when timeShiftColor is disabled', () => {
const queryAData = createTestQueryData(
[
{
sum__num: 100,
'sum__num__1 week ago': 80,
ds: 599616000000,
},
{
sum__num: 150,
'sum__num__1 week ago': 120,
ds: 599916000000,
},
],
{
label_map: {
ds: ['ds'],
sum__num: ['sum__num'],
'sum__num__1 week ago': ['sum__num__1 week ago'],
},
},
);
const chartProps = createEchartsTimeseriesTestChartProps<
EchartsMixedTimeseriesFormData,
EchartsMixedTimeseriesProps
>({
...MIXED_TIMESERIES_CHART_PROPS_DEFAULTS,
defaultQueriesData: [queryAData, queriesData[1]],
formData: {
...formData,
metrics: ['sum__num'],
groupby: [],
time_compare: ['1 week ago'],
comparison_type: ComparisonType.Values,
timeShiftColor: false,
},
queriesData: [queryAData, queriesData[1]],
});
const transformed = transformProps(chartProps);
const series = (transformed.echartOptions.series as SeriesOption[]) || [];
const derivedSeries = series.find(s => s.name === 'sum__num__1 week ago') as
| (SeriesOption & { lineStyle?: { type?: number[] | string } })
| undefined;
expect(derivedSeries).toBeDefined();
expect(derivedSeries?.lineStyle?.type).toBeUndefined();
});
+19
View File
@@ -339,6 +339,25 @@ def load_configs(
exceptions.append(
ValidationError({file_name: {"masked_encrypted_extra": [str(exc)]}})
)
except KeyError as exc:
# Some config fields (e.g. `uuid`, `ssh_tunnel`) are read
# directly from the imported YAML before schema validation runs;
# a config missing one of these keys raises a raw KeyError
# instead of failing validation cleanly like every other
# per-file error. Convert it into a ValidationError so it flows
# into the same aggregated error path.
field = str(exc).strip("'\"")
logger.error(
"Missing required key %s in config for %s (prefix: %s)",
exc,
file_name,
prefix,
)
exceptions.append(
ValidationError(
{file_name: {field: ["Missing data for required field."]}}
)
)
return configs
+49 -4
View File
@@ -1453,10 +1453,33 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
found.add(entry)
return found
def _has_limit_by(self) -> bool:
"""
Check if the statement has a ClickHouse `LIMIT ... BY` clause.
`LIMIT n BY <cols>` keeps `n` rows *per group*, so it is a de-duplication
clause rather than a row cap. sqlglot models the `BY` columns as the
`expressions` of the root `Limit` node, or of the root `Offset` node for
the `LIMIT n OFFSET m BY x` and `LIMIT m, n BY x` spellings.
:return: True if the statement's limit or offset carries `BY` columns.
"""
for arg in ("limit", "offset"):
node = self._parsed.args.get(arg)
if isinstance(node, exp.Expression) and node.expressions:
return True
return False
def get_limit_value(self) -> int | None:
"""
Parse a SQL query and return the `LIMIT` or `TOP` value, if present.
"""
# `LIMIT 2 BY id` bounds each group, not the result set, so reporting 2
# here would make `_set_query_limit()` clamp the whole query to 2 rows.
if self._has_limit_by():
return None
if limit_node := self._parsed.args.get("limit"):
literal = limit_node.args.get("expression") or getattr(
limit_node, "this", None
@@ -1494,18 +1517,40 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
if not isinstance(self._parsed, exp.Query):
return
if method == LimitMethod.FORCE_LIMIT:
# A ClickHouse `LIMIT ... BY` occupies the very `limit`/`offset` slot that
# `FORCE_LIMIT` overwrites, so forcing a row cap in place would drop the
# `BY` grouping and silently change what the query returns. The cap can't
# be appended alongside it either -- sqlglot rejects ClickHouse's native
# `LIMIT n BY x LIMIT m` with "Found multiple 'LIMIT' clauses" -- so it
# goes on a wrapping query instead, exactly as `WRAP_SQL` does.
if method == LimitMethod.FORCE_LIMIT and not self._has_limit_by():
self._parsed.args["limit"] = exp.Limit(
expression=exp.Literal(this=str(limit), is_string=False)
)
elif method == LimitMethod.WRAP_SQL:
self._parsed = exp.Select(
elif method in {LimitMethod.FORCE_LIMIT, LimitMethod.WRAP_SQL}:
inner = self._parsed.copy()
wrapper = exp.Select(
expressions=[exp.Star()],
limit=exp.Limit(
expression=exp.Literal(this=str(limit), is_string=False)
),
from_=exp.From(this=exp.Subquery(this=self._parsed.copy())),
from_=exp.From(this=exp.Subquery(this=inner)),
)
# `FORMAT` and `SETTINGS` configure the query rather than produce
# rows, and only mean what they say at the top level: ClickHouse
# rejects `FORMAT` inside a subquery outright, and a nested
# `SETTINGS` binds to that subquery alone, so top-level-only settings
# such as `extremes` would quietly stop applying. Moving them onto
# the wrapper keeps their original whole-query scope. Row-producing
# modifiers stay in the subquery, where ClickHouse keeps honoring
# them: a wrapped `WITH TOTALS` query still emits its totals block,
# and `WITH ROLLUP`/`WITH CUBE` still emit their extra rows.
for modifier in ("format", "settings"):
if value := inner.args.pop(modifier, None):
wrapper.set(modifier, value)
self._parsed = wrapper
else: # method == LimitMethod.FETCH_MANY
pass
@@ -18,7 +18,7 @@
import gzip
import io
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
@@ -168,12 +168,14 @@ class TestLoadYaml:
class TestLoadConfigs:
"""
load_configs() merges caller-supplied ``encrypted_extra_secrets`` into the
``masked_encrypted_extra`` field of each config, which comes straight from
the imported YAML (before schema validation). A malformed value there used
to raise a raw simplejson.JSONDecodeError that escaped uncaught (opaque
500); it must instead be collected as a ValidationError like every other
per-file failure.
Per-file failures inside load_configs() must be collected as
ValidationErrors rather than propagating as raw exceptions (opaque 500s):
- A malformed ``masked_encrypted_extra`` (into which caller-supplied
``encrypted_extra_secrets`` are merged before schema validation) used to
raise a raw simplejson.JSONDecodeError.
- A config missing its ``uuid`` used to raise a raw KeyError when the
password/ssh-tunnel validation looked up ``config["uuid"]``.
"""
@staticmethod
@@ -186,6 +188,17 @@ class TestLoadConfigs:
return TrivialSchema()
def _database_schemas(self) -> dict[str, object]:
from marshmallow import fields, Schema
class DatabaseSchema(Schema):
uuid = fields.UUID(required=True)
database_name = fields.String(required=True)
sqlalchemy_uri = fields.String(required=True)
password = fields.String(required=False, allow_none=True)
return {"databases/": DatabaseSchema()}
@patch("superset.commands.importers.v1.utils.db")
def test_invalid_json_in_masked_encrypted_extra_is_collected(
self, mock_db: object
@@ -267,6 +280,75 @@ class TestLoadConfigs:
merged = json.loads(configs[file_name]["masked_encrypted_extra"])
assert merged == {"foo": "actual_secret"}
@patch("superset.commands.importers.v1.utils.db")
def test_missing_uuid_appends_validation_error(self, mock_db: MagicMock) -> None:
"""A databases config missing `uuid` must not raise a raw KeyError;
it should be excluded from the returned configs and a ValidationError
appended to the exceptions list instead."""
from marshmallow.exceptions import ValidationError
from superset.commands.importers.v1.utils import load_configs
mock_db.session.query.return_value.all.return_value = []
# No `uuid` and no `password`, so the code reaches
# `config["uuid"] in db_passwords` and would raise KeyError pre-fix.
contents = {
"databases/bad.yaml": (
"database_name: bad\nsqlalchemy_uri: postgres://localhost\n"
),
}
exceptions: list[ValidationError] = []
configs = load_configs(
contents,
self._database_schemas(),
{},
exceptions,
{},
{},
{},
{},
)
assert "databases/bad.yaml" not in configs
assert len(exceptions) == 1
assert isinstance(exceptions[0], ValidationError)
assert "databases/bad.yaml" in exceptions[0].messages
@patch("superset.commands.importers.v1.utils.db")
def test_uuid_present_loads_successfully(self, mock_db: MagicMock) -> None:
"""Control: a well-formed databases config loads with no exceptions."""
from marshmallow.exceptions import ValidationError
from superset.commands.importers.v1.utils import load_configs
mock_db.session.query.return_value.all.return_value = []
contents = {
"databases/good.yaml": (
"uuid: 6ff1d5b3-4b0f-4c6a-9d2f-9c8b7a6e5d4c\n"
"database_name: good\n"
"sqlalchemy_uri: postgres://localhost\n"
"password: secret\n"
),
}
exceptions: list[ValidationError] = []
configs = load_configs(
contents,
self._database_schemas(),
{},
exceptions,
{},
{},
{},
{},
)
assert "databases/good.yaml" in configs
assert exceptions == []
class TestLoadConfigsNonMappingYaml:
"""A syntactically valid YAML document whose top-level value is a
+249
View File
@@ -2502,6 +2502,21 @@ LATERAL generate_series(1, value) AS i;
),
# not really valid SQL, but let's roll with it
("SELECT * FROM my_table LIMIT invalid", "postgresql", None),
# A ClickHouse `LIMIT ... BY` caps rows per group, not overall, so it is
# not a row limit. sqlglot hangs the `BY` columns off the `Limit` node,
# or off the `Offset` node for the `OFFSET` / `m, n` spellings.
("SELECT * FROM t ORDER BY id, val LIMIT 2 BY id", "clickhouse", None),
("SELECT * FROM t ORDER BY id, val LIMIT 2 BY id, val", "clickhouse", None),
(
"SELECT * FROM t ORDER BY id, val LIMIT 2 OFFSET 1 BY id",
"clickhouse",
None,
),
("SELECT * FROM t ORDER BY id, val LIMIT 1, 2 BY id", "clickhouse", None),
# ... while a plain ClickHouse limit, with or without an offset, is.
("SELECT * FROM t ORDER BY c LIMIT 555", "clickhouse", 555),
("SELECT * FROM t LIMIT 5 OFFSET 3", "clickhouse", 5),
("SELECT * FROM t LIMIT 3, 5", "clickhouse", 5),
],
)
def test_get_limit_value(sql: str, engine: str, expected: str) -> None:
@@ -2717,6 +2732,158 @@ LIMIT 1000
LimitMethod.FETCH_MANY,
"SELECT\n *\nFROM birth_names\nLIMIT 555",
),
# A ClickHouse `LIMIT ... BY` shares the `limit`/`offset` slot with the
# row limit, so `FORCE_LIMIT` wraps instead of overwriting it.
(
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 BY id",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
*
FROM limit_by
ORDER BY
id,
val
LIMIT 2 BY id
)
LIMIT 1001
""".strip(),
),
(
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 BY id, val",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
*
FROM limit_by
ORDER BY
id,
val
LIMIT 2 BY id, val
)
LIMIT 1001
""".strip(),
),
# For `LIMIT n OFFSET m BY x` sqlglot hangs the `BY` columns off the
# `Offset` node instead, so the `limit` arg alone doesn't reveal them.
(
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 OFFSET 1 BY id",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
*
FROM limit_by
ORDER BY
id,
val
LIMIT 2
OFFSET 1 BY id
)
LIMIT 1001
""".strip(),
),
(
"SELECT * FROM limit_by ORDER BY id, val LIMIT 1, 2 BY id",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
*
FROM limit_by
ORDER BY
id,
val
LIMIT 2
OFFSET 1 BY id
)
LIMIT 1001
""".strip(),
),
# `WITH TOTALS` rides into the subquery untouched: ClickHouse keeps
# emitting the totals block for a wrapped query, so the cap really is
# the only thing the rewrite adds.
(
"SELECT id, count() AS c FROM limit_by "
"GROUP BY id WITH TOTALS ORDER BY id LIMIT 2 BY id",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
id,
count() AS c
FROM limit_by
GROUP BY
id
WITH TOTALS
ORDER BY
id
LIMIT 2 BY id
)
LIMIT 1001
""".strip(),
),
# `SETTINGS` and `FORMAT` do not survive a demotion into the subquery,
# so they move up onto the wrapper instead.
(
"SELECT * FROM limit_by ORDER BY id LIMIT 2 BY id "
"SETTINGS extremes = 1 FORMAT JSONCompact",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"""
SELECT
*
FROM (
SELECT
*
FROM limit_by
ORDER BY
id
LIMIT 2 BY id
)
LIMIT 1001
SETTINGS extremes = 1
FORMAT JSONCompact
""".strip(),
),
# A ClickHouse limit without a `BY` still takes the in-place path.
(
"SELECT * FROM t ORDER BY c LIMIT 555",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM t\nORDER BY\n c\nLIMIT 1001",
),
(
"SELECT * FROM t LIMIT 5 OFFSET 3",
"clickhouse",
1001,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM t\nLIMIT 1001\nOFFSET 3",
),
],
)
def test_set_limit_value(
@@ -2731,6 +2898,88 @@ def test_set_limit_value(
assert statement.format() == expected
@pytest.mark.parametrize("engine", ["clickhouse", "clickhousedb"])
@pytest.mark.parametrize(
"sql",
[
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 BY id",
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 BY id, val",
"SELECT * FROM limit_by ORDER BY id, val LIMIT 2 OFFSET 1 BY id",
"SELECT * FROM limit_by ORDER BY id, val LIMIT 1, 2 BY id",
],
)
def test_set_limit_value_preserves_clickhouse_limit_by(sql: str, engine: str) -> None:
"""
A row limit must not cannibalize a ClickHouse ``LIMIT ... BY``.
``LIMIT 2 BY id`` keeps 2 rows *per id*; ``FORCE_LIMIT`` used to build a
fresh ``Limit`` node over ``args["limit"]``, dropping the ``BY`` columns and
turning the query into a flat ``LIMIT 1001`` -- a different result set, with
no error to hint at it. ``get_limit_value()`` reported the per-group 2 as a
row cap on top of that, so ``_set_query_limit()`` clamped the query to 2 rows.
The cap can't simply be appended next to the ``BY`` either: sqlglot cannot
parse ClickHouse's own ``LIMIT n BY x LIMIT m`` ("Found multiple 'LIMIT'
clauses"), so the result would not survive a reparse. Wrapping the query is
what keeps both the grouping and the cap.
"""
statement = SQLStatement(sql, engine)
assert statement.get_limit_value() is None
statement.set_limit_value(1001, LimitMethod.FORCE_LIMIT)
limited = statement.format()
assert "BY id" in limited
assert limited.endswith("LIMIT 1001")
# The rewrite has to be valid ClickHouse, not just valid-looking.
assert SQLStatement(limited, engine).format() == limited
def test_set_limit_value_keeps_clickhouse_top_level_modifiers() -> None:
"""
The wrap must not demote clauses that only work at the top level.
ClickHouse rejects `FORMAT` inside a subquery outright, and a `SETTINGS`
attached to a subquery binds to that subquery alone -- top-level-only
settings such as ``extremes`` would silently stop applying. Both therefore
move onto the wrapper, which is where the original query had them.
The row-producing modifiers are left alone, because ClickHouse honors them
inside a `FROM` subquery: a wrapped `WITH TOTALS` query still emits its
totals block, and `WITH ROLLUP`/`WITH CUBE` still emit their extra rows.
Hoisting those would change the result rather than preserve it.
"""
statement = SQLStatement(
"SELECT id, count() AS c FROM limit_by "
"GROUP BY id WITH TOTALS ORDER BY id LIMIT 2 BY id "
"SETTINGS extremes = 1 FORMAT JSONCompact",
"clickhouse",
)
statement.set_limit_value(1001, LimitMethod.FORCE_LIMIT)
limited = statement.format()
assert limited.endswith("LIMIT 1001\nSETTINGS extremes = 1\nFORMAT JSONCompact")
# `WITH TOTALS` stays with the aggregation it belongs to.
assert "WITH TOTALS\n" in limited.split("LIMIT 2 BY id")[0]
assert SQLStatement(limited, "clickhouse").format() == limited
@pytest.mark.parametrize(
"engine", ["clickhouse", "clickhousedb", "postgresql", "mysql"]
)
def test_set_limit_value_without_limit_by_stays_in_place(engine: str) -> None:
"""
Queries with no ``LIMIT ... BY`` keep the cheaper in-place rewrite.
The wrap is reserved for the ``LIMIT ... BY`` case; everything else -- every
non-ClickHouse dialect, and ClickHouse's own plain ``LIMIT`` -- must still
have its limit replaced without gaining a subquery.
"""
statement = SQLStatement("SELECT * FROM t ORDER BY c LIMIT 555", engine)
statement.set_limit_value(1001, LimitMethod.FORCE_LIMIT)
assert statement.format() == "SELECT\n *\nFROM t\nORDER BY\n c\nLIMIT 1001"
@pytest.mark.parametrize(
"method",
[LimitMethod.FORCE_LIMIT, LimitMethod.WRAP_SQL],