Edit on GitHub

Remove hash_raw_query from existing snapshots.

 1"""Remove hash_raw_query from existing snapshots."""
 2
 3import json
 4
 5import pandas as pd
 6from sqlglot import exp
 7
 8from sqlmesh.utils.migration import index_text_type
 9
10
11def migrate(state_sync, **kwargs):  # type: ignore
12    engine_adapter = state_sync.engine_adapter
13    schema = state_sync.schema
14    snapshots_table = "_snapshots"
15    if schema:
16        snapshots_table = f"{schema}.{snapshots_table}"
17
18    new_snapshots = []
19
20    for name, identifier, version, snapshot, kind_name, expiration_ts in engine_adapter.fetchall(
21        exp.select("name", "identifier", "version", "snapshot", "kind_name", "expiration_ts").from_(
22            snapshots_table
23        ),
24        quote_identifiers=True,
25    ):
26        parsed_snapshot = json.loads(snapshot)
27        parsed_snapshot["node"].pop("hash_raw_query", None)
28
29        new_snapshots.append(
30            {
31                "name": name,
32                "identifier": identifier,
33                "version": version,
34                "snapshot": json.dumps(parsed_snapshot),
35                "kind_name": kind_name,
36                "expiration_ts": expiration_ts,
37            }
38        )
39
40    if new_snapshots:
41        engine_adapter.delete_from(snapshots_table, "TRUE")
42
43        index_type = index_text_type(engine_adapter.dialect)
44
45        engine_adapter.insert_append(
46            snapshots_table,
47            pd.DataFrame(new_snapshots),
48            columns_to_types={
49                "name": exp.DataType.build(index_type),
50                "identifier": exp.DataType.build(index_type),
51                "version": exp.DataType.build(index_type),
52                "snapshot": exp.DataType.build("text"),
53                "kind_name": exp.DataType.build(index_type),
54                "expiration_ts": exp.DataType.build("bigint"),
55            },
56        )
def migrate(state_sync, **kwargs): 
12def migrate(state_sync, **kwargs):  # type: ignore
13    engine_adapter = state_sync.engine_adapter
14    schema = state_sync.schema
15    snapshots_table = "_snapshots"
16    if schema:
17        snapshots_table = f"{schema}.{snapshots_table}"
18
19    new_snapshots = []
20
21    for name, identifier, version, snapshot, kind_name, expiration_ts in engine_adapter.fetchall(
22        exp.select("name", "identifier", "version", "snapshot", "kind_name", "expiration_ts").from_(
23            snapshots_table
24        ),
25        quote_identifiers=True,
26    ):
27        parsed_snapshot = json.loads(snapshot)
28        parsed_snapshot["node"].pop("hash_raw_query", None)
29
30        new_snapshots.append(
31            {
32                "name": name,
33                "identifier": identifier,
34                "version": version,
35                "snapshot": json.dumps(parsed_snapshot),
36                "kind_name": kind_name,
37                "expiration_ts": expiration_ts,
38            }
39        )
40
41    if new_snapshots:
42        engine_adapter.delete_from(snapshots_table, "TRUE")
43
44        index_type = index_text_type(engine_adapter.dialect)
45
46        engine_adapter.insert_append(
47            snapshots_table,
48            pd.DataFrame(new_snapshots),
49            columns_to_types={
50                "name": exp.DataType.build(index_type),
51                "identifier": exp.DataType.build(index_type),
52                "version": exp.DataType.build(index_type),
53                "snapshot": exp.DataType.build("text"),
54                "kind_name": exp.DataType.build(index_type),
55                "expiration_ts": exp.DataType.build("bigint"),
56            },
57        )