sqlmesh.schedulers.airflow.common API documentation

Edit on GitHub

# sqlmesh.schedulers.airflow.common

View Source

      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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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

def snapshot_key( snapshot: Union[sqlmesh.core.snapshot.definition.SnapshotId, sqlmesh.core.snapshot.definition.SnapshotTableInfo, sqlmesh.core.snapshot.definition.Snapshot]) -> str: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    124def environment_key(env: str) -> str:
    125    return f"{ENV_KEY_PREFIX}__{env}"

def plan_dag_spec_key(request_id: str) -> str: View Source

    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: View Source

    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)