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)
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
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
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
Inherited Members
- pydantic.main.BaseModel
- BaseModel
- parse_obj
- parse_raw
- parse_file
- from_orm
- construct
- copy
- schema
- schema_json
- validate
- update_forward_refs
Inherited Members
- pydantic.main.BaseModel
- BaseModel
- parse_obj
- parse_raw
- parse_file
- from_orm
- construct
- copy
- schema
- schema_json
- validate
- update_forward_refs
Inherited Members
- pydantic.main.BaseModel
- BaseModel
- parse_obj
- parse_raw
- parse_file
- from_orm
- construct
- copy
- schema
- schema_json
- validate
- update_forward_refs
def
snapshot_key( snapshot: Union[sqlmesh.core.snapshot.definition.SnapshotId, sqlmesh.core.snapshot.definition.SnapshotTableInfo, sqlmesh.core.snapshot.definition.Snapshot]) -> str:
def
snapshot_key_from_name_identifier(name: str, identifier: str) -> str:
def
snapshot_version_key(name: str, version: Optional[str] = None) -> str:
def
name_from_snapshot_version_key(key: str) -> str:
def
dag_id_for_snapshot_info( info: Union[sqlmesh.core.snapshot.definition.SnapshotTableInfo, sqlmesh.core.snapshot.definition.Snapshot]) -> str:
def
dag_id_for_name_version(name: str, version: str) -> str:
def
plan_application_dag_id(environment: str, request_id: str) -> str:
def
environment_key(env: str) -> str:
def
plan_dag_spec_key(request_id: str) -> str:
def
plan_dag_spec_key_from_dag_id(dag_id: str) -> str: