Edit on GitHub

sqlmesh.schedulers.airflow.plugin

 1from __future__ import annotations
 2
 3import logging
 4import time
 5import typing as t
 6
 7from airflow.plugins_manager import AirflowPlugin
 8
 9from sqlmesh.core import constants as c
10from sqlmesh.schedulers.airflow import util
11from sqlmesh.schedulers.airflow.api import sqlmesh_api_v1
12
13logger = logging.getLogger(__name__)
14
15
16class SqlmeshAirflowPlugin(AirflowPlugin):
17    name = c.SQLMESH
18    flask_blueprints = [sqlmesh_api_v1]
19
20    @classmethod
21    def on_load(cls, *args: t.Any, **kwargs: t.Any) -> None:
22        with util.scoped_state_sync() as state_sync:
23            try:
24                state_sync.init_schema()
25            except Exception as ex:
26                # This method is called once for each Gunicorn worker spawned by the Airflow Webserver,
27                # which leads to SQLMesh schema being initialized concurrently from multiple processes.
28                # There is a known issue in Postgres (https://stackoverflow.com/a/29908840) which occurs
29                # due to a race condition when a new schema is being created concurrently. Here we retry
30                # the schema initialization once as a workaround.
31                logger.warning("Failed to initialize the SQLMesh State Sync: %s. Retrying...", ex)
32                time.sleep(1)
33                state_sync.init_schema()
class SqlmeshAirflowPlugin(airflow.plugins_manager.AirflowPlugin):
17class SqlmeshAirflowPlugin(AirflowPlugin):
18    name = c.SQLMESH
19    flask_blueprints = [sqlmesh_api_v1]
20
21    @classmethod
22    def on_load(cls, *args: t.Any, **kwargs: t.Any) -> None:
23        with util.scoped_state_sync() as state_sync:
24            try:
25                state_sync.init_schema()
26            except Exception as ex:
27                # This method is called once for each Gunicorn worker spawned by the Airflow Webserver,
28                # which leads to SQLMesh schema being initialized concurrently from multiple processes.
29                # There is a known issue in Postgres (https://stackoverflow.com/a/29908840) which occurs
30                # due to a race condition when a new schema is being created concurrently. Here we retry
31                # the schema initialization once as a workaround.
32                logger.warning("Failed to initialize the SQLMesh State Sync: %s. Retrying...", ex)
33                time.sleep(1)
34                state_sync.init_schema()

Class used to define AirflowPlugin.

SqlmeshAirflowPlugin()
@classmethod
def on_load(cls, *args: Any, **kwargs: Any) -> None:
21    @classmethod
22    def on_load(cls, *args: t.Any, **kwargs: t.Any) -> None:
23        with util.scoped_state_sync() as state_sync:
24            try:
25                state_sync.init_schema()
26            except Exception as ex:
27                # This method is called once for each Gunicorn worker spawned by the Airflow Webserver,
28                # which leads to SQLMesh schema being initialized concurrently from multiple processes.
29                # There is a known issue in Postgres (https://stackoverflow.com/a/29908840) which occurs
30                # due to a race condition when a new schema is being created concurrently. Here we retry
31                # the schema initialization once as a workaround.
32                logger.warning("Failed to initialize the SQLMesh State Sync: %s. Retrying...", ex)
33                time.sleep(1)
34                state_sync.init_schema()

Executed when the plugin is loaded. This method is only called once during runtime.

Parameters
  • args: If future arguments are passed in on call.
  • kwargs: If future arguments are passed in on call.
Inherited Members
airflow.plugins_manager.AirflowPlugin
validate