Edit on GitHub

sqlmesh.schedulers.airflow.common

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlmesh.core import constants as c
  6from sqlmesh.core._typing import NotificationTarget
  7from sqlmesh.core.environment import Environment
  8from sqlmesh.core.scheduler import Interval
  9from sqlmesh.core.snapshot import (
 10    Snapshot,
 11    SnapshotId,
 12    SnapshotIdLike,
 13    SnapshotInfoLike,
 14    SnapshotTableInfo,
 15)
 16from sqlmesh.core.user import User
 17from sqlmesh.utils.date import TimeLike
 18from sqlmesh.utils.errors import SQLMeshError
 19from sqlmesh.utils.pydantic import PydanticModel
 20
 21JANITOR_DAG_ID = "sqlmesh_janitor_dag"
 22JANITOR_TASK_ID = "janitor_task"
 23
 24SQLMESH_AIRFLOW_TAG = "sqlmesh"
 25SNAPSHOT_AIRFLOW_TAG = "sqlmesh_snapshot"
 26PLAN_AIRFLOW_TAG = "sqlmesh_plan"
 27
 28SNAPSHOT_CLEANUP_COMMAND_XCOM_KEY = "snapshot_cleanup_command"
 29
 30PLAN_DAG_SPEC_KEY_PREFIX = "sqlmesh__plan_dag_spec"
 31SNAPSHOT_PAYLOAD_KEY_PREFIX = "sqlmesh__snapshot_payload"
 32SNAPSHOT_VERSION_KEY_PREFIX = "sqlmesh__snapshot_version_index"
 33ENV_KEY_PREFIX = "sqlmesh__environment"
 34
 35AIRFLOW_LOCAL_URL = "http://localhost:8080/"
 36
 37SQLMESH_API_BASE_PATH: str = f"{c.SQLMESH}/api/v1"
 38
 39
 40class PlanApplicationRequest(PydanticModel):
 41    request_id: str
 42    new_snapshots: t.List[Snapshot]
 43    environment: Environment
 44    no_gaps: bool
 45    skip_backfill: bool
 46    restatements: t.Set[str]
 47    notification_targets: t.List[NotificationTarget]
 48    backfill_concurrent_tasks: int
 49    ddl_concurrent_tasks: int
 50    users: t.List[User]
 51    is_dev: bool
 52
 53
 54class BackfillIntervalsPerSnapshot(PydanticModel):
 55    snapshot_id: SnapshotId
 56    intervals: t.List[Interval]
 57
 58
 59class PlanDagSpec(PydanticModel):
 60    request_id: str
 61    environment_name: str
 62    new_snapshots: t.List[Snapshot]
 63    backfill_intervals_per_snapshot: t.List[BackfillIntervalsPerSnapshot]
 64    promoted_snapshots: t.List[SnapshotTableInfo]
 65    demoted_snapshots: t.List[SnapshotTableInfo]
 66    start: TimeLike
 67    end: t.Optional[TimeLike]
 68    unpaused_dt: t.Optional[TimeLike]
 69    no_gaps: bool
 70    plan_id: str
 71    previous_plan_id: t.Optional[str]
 72    notification_targets: t.List[NotificationTarget]
 73    backfill_concurrent_tasks: int
 74    ddl_concurrent_tasks: int
 75    users: t.List[User]
 76    is_dev: bool
 77    environment_expiration_ts: t.Optional[int]
 78
 79
 80class EnvironmentsResponse(PydanticModel):
 81    environments: t.List[Environment]
 82
 83
 84class SnapshotsResponse(PydanticModel):
 85    snapshots: t.List[Snapshot]
 86
 87
 88class SnapshotIdsResponse(PydanticModel):
 89    snapshot_ids: t.List[SnapshotId]
 90
 91
 92def snapshot_key(snapshot: SnapshotIdLike) -> str:
 93    return snapshot_key_from_name_identifier(snapshot.name, snapshot.identifier)
 94
 95
 96def snapshot_key_from_name_identifier(name: str, identifier: str) -> str:
 97    return f"{SNAPSHOT_PAYLOAD_KEY_PREFIX}__{name}__{identifier}"
 98
 99
100def snapshot_version_key(name: str, version: t.Optional[str] = None) -> str:
101    if not version:
102        raise SQLMeshError("Version cannot be empty")
103    return f"{SNAPSHOT_VERSION_KEY_PREFIX}__{name}__{version}"
104
105
106def name_from_snapshot_version_key(key: str) -> str:
107    return key[len(f"{SNAPSHOT_VERSION_KEY_PREFIX}__") : key.rindex("__")]
108
109
110def dag_id_for_snapshot_info(info: SnapshotInfoLike) -> str:
111    assert info.version
112    return dag_id_for_name_version(info.name, info.version)
113
114
115def dag_id_for_name_version(name: str, version: str) -> str:
116    return f"sqlmesh_snapshot_{name}_{version}_dag"
117
118
119def plan_application_dag_id(environment: str, request_id: str) -> str:
120    return f"sqlmesh_plan_application__{environment}__{request_id}"
121
122
123def environment_key(env: str) -> str:
124    return f"{ENV_KEY_PREFIX}__{env}"
125
126
127def plan_dag_spec_key(request_id: str) -> str:
128    return f"{PLAN_DAG_SPEC_KEY_PREFIX}__{request_id}"
129
130
131def plan_dag_spec_key_from_dag_id(dag_id: str) -> str:
132    request_id = dag_id[dag_id.rindex("__") + 2 :]
133    return plan_dag_spec_key(request_id)
class PlanApplicationRequest(sqlmesh.utils.pydantic.PydanticModel):
41class PlanApplicationRequest(PydanticModel):
42    request_id: str
43    new_snapshots: t.List[Snapshot]
44    environment: Environment
45    no_gaps: bool
46    skip_backfill: bool
47    restatements: t.Set[str]
48    notification_targets: t.List[NotificationTarget]
49    backfill_concurrent_tasks: int
50    ddl_concurrent_tasks: int
51    users: t.List[User]
52    is_dev: bool
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
class BackfillIntervalsPerSnapshot(sqlmesh.utils.pydantic.PydanticModel):
55class BackfillIntervalsPerSnapshot(PydanticModel):
56    snapshot_id: SnapshotId
57    intervals: t.List[Interval]
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
class PlanDagSpec(sqlmesh.utils.pydantic.PydanticModel):
60class PlanDagSpec(PydanticModel):
61    request_id: str
62    environment_name: str
63    new_snapshots: t.List[Snapshot]
64    backfill_intervals_per_snapshot: t.List[BackfillIntervalsPerSnapshot]
65    promoted_snapshots: t.List[SnapshotTableInfo]
66    demoted_snapshots: t.List[SnapshotTableInfo]
67    start: TimeLike
68    end: t.Optional[TimeLike]
69    unpaused_dt: t.Optional[TimeLike]
70    no_gaps: bool
71    plan_id: str
72    previous_plan_id: t.Optional[str]
73    notification_targets: t.List[NotificationTarget]
74    backfill_concurrent_tasks: int
75    ddl_concurrent_tasks: int
76    users: t.List[User]
77    is_dev: bool
78    environment_expiration_ts: t.Optional[int]
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
class EnvironmentsResponse(sqlmesh.utils.pydantic.PydanticModel):
81class EnvironmentsResponse(PydanticModel):
82    environments: t.List[Environment]
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
class SnapshotsResponse(sqlmesh.utils.pydantic.PydanticModel):
85class SnapshotsResponse(PydanticModel):
86    snapshots: t.List[Snapshot]
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
class SnapshotIdsResponse(sqlmesh.utils.pydantic.PydanticModel):
89class SnapshotIdsResponse(PydanticModel):
90    snapshot_ids: t.List[SnapshotId]
Inherited Members
pydantic.main.BaseModel
BaseModel
parse_obj
parse_raw
parse_file
from_orm
construct
copy
schema
schema_json
validate
update_forward_refs
sqlmesh.utils.pydantic.PydanticModel
Config
dict
json
missing_required_fields
extra_fields
all_fields
required_fields
93def snapshot_key(snapshot: SnapshotIdLike) -> str:
94    return snapshot_key_from_name_identifier(snapshot.name, snapshot.identifier)
def snapshot_key_from_name_identifier(name: str, identifier: str) -> str:
97def snapshot_key_from_name_identifier(name: str, identifier: str) -> str:
98    return f"{SNAPSHOT_PAYLOAD_KEY_PREFIX}__{name}__{identifier}"
def snapshot_version_key(name: str, version: Optional[str] = None) -> str:
101def snapshot_version_key(name: str, version: t.Optional[str] = None) -> str:
102    if not version:
103        raise SQLMeshError("Version cannot be empty")
104    return f"{SNAPSHOT_VERSION_KEY_PREFIX}__{name}__{version}"
def name_from_snapshot_version_key(key: str) -> str:
107def name_from_snapshot_version_key(key: str) -> str:
108    return key[len(f"{SNAPSHOT_VERSION_KEY_PREFIX}__") : key.rindex("__")]
def dag_id_for_snapshot_info( info: Union[sqlmesh.core.snapshot.definition.SnapshotTableInfo, sqlmesh.core.snapshot.definition.Snapshot]) -> str:
111def dag_id_for_snapshot_info(info: SnapshotInfoLike) -> str:
112    assert info.version
113    return dag_id_for_name_version(info.name, info.version)
def dag_id_for_name_version(name: str, version: str) -> str:
116def dag_id_for_name_version(name: str, version: str) -> str:
117    return f"sqlmesh_snapshot_{name}_{version}_dag"
def plan_application_dag_id(environment: str, request_id: str) -> str:
120def plan_application_dag_id(environment: str, request_id: str) -> str:
121    return f"sqlmesh_plan_application__{environment}__{request_id}"
def environment_key(env: str) -> str:
124def environment_key(env: str) -> str:
125    return f"{ENV_KEY_PREFIX}__{env}"
def plan_dag_spec_key(request_id: str) -> str:
128def plan_dag_spec_key(request_id: str) -> str:
129    return f"{PLAN_DAG_SPEC_KEY_PREFIX}__{request_id}"
def plan_dag_spec_key_from_dag_id(dag_id: str) -> str:
132def plan_dag_spec_key_from_dag_id(dag_id: str) -> str:
133    request_id = dag_id[dag_id.rindex("__") + 2 :]
134    return plan_dag_spec_key(request_id)