Edit on GitHub

Replace snapshot model_kind_name enum with value.

 1"""Replace snapshot model_kind_name enum with value."""
 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 in engine_adapter.fetchall(
21        exp.select("name", "identifier", "version", "snapshot", "kind_name").from_(snapshots_table),
22        quote_identifiers=True,
23    ):
24        corrected_kind_name = None
25        parsed_snapshot = json.loads(snapshot)
26        if "kind" in parsed_snapshot["node"]:
27            corrected_kind_name = parsed_snapshot["node"]["kind"].get("name")
28
29        new_snapshots.append(
30            {
31                "name": name,
32                "identifier": identifier,
33                "version": version,
34                "snapshot": snapshot,
35                "kind_name": corrected_kind_name,
36            }
37        )
38
39    if new_snapshots:
40        engine_adapter.delete_from(snapshots_table, "TRUE")
41
42        index_type = index_text_type(engine_adapter.dialect)
43
44        engine_adapter.insert_append(
45            snapshots_table,
46            pd.DataFrame(new_snapshots),
47            columns_to_types={
48                "name": exp.DataType.build(index_type),
49                "identifier": exp.DataType.build(index_type),
50                "version": exp.DataType.build(index_type),
51                "snapshot": exp.DataType.build("text"),
52                "kind_name": exp.DataType.build(index_type),
53            },
54        )
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 in engine_adapter.fetchall(
22        exp.select("name", "identifier", "version", "snapshot", "kind_name").from_(snapshots_table),
23        quote_identifiers=True,
24    ):
25        corrected_kind_name = None
26        parsed_snapshot = json.loads(snapshot)
27        if "kind" in parsed_snapshot["node"]:
28            corrected_kind_name = parsed_snapshot["node"]["kind"].get("name")
29
30        new_snapshots.append(
31            {
32                "name": name,
33                "identifier": identifier,
34                "version": version,
35                "snapshot": snapshot,
36                "kind_name": corrected_kind_name,
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            },
55        )