Python models.DagBag方法代码示例

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def collect_process_output(process):
        output = REST_API.get_empty_process_output()
        if process.stderr is not None:
            output["stderr"] = ""
            for line in process.stderr.readlines():
                output["stderr"] += str(line.decode('utf-8'))
        if process.stdin is not None:
            output["stdin"] = ""
            for line in process.stdin.readlines():
                output["stdin"] += str(line.decode('utf-8'))
        if process.stdout is not None:
            output["stdout"] = ""
            for line in process.stdout.readlines():
                output["stdout"] += str(line.decode('utf-8'))
        logging.info("RestAPI Output: " + str(output))
        return output

    # Filtering out logging statements from the standard output
    # Content like:
    #
    # [2017-04-19 10:04:34,927] {__init__.py:36} INFO - Using executor CeleryExecutor
    # [2017-04-19 10:04:35,926] {models.py:154} INFO - Filling up the DagBag from /Users/... 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def prepare_dagruns(self):
        dagbag = models.DagBag(include_examples=True)
        self.bash_dag = dagbag.dags['example_bash_operator']
        self.sub_dag = dagbag.dags['example_subdag_operator']

        self.bash_dagrun = self.bash_dag.create_dagrun(
            run_type=DagRunType.SCHEDULED,
            execution_date=self.default_date,
            start_date=timezone.utcnow(),
            state=State.RUNNING)

        self.sub_dagrun = self.sub_dag.create_dagrun(
            run_type=DagRunType.SCHEDULED,
            execution_date=self.default_date,
            start_date=timezone.utcnow(),
            state=State.RUNNING) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_safe_mode_heuristic_match(self):
        """With safe mode enabled, a file matching the discovery heuristics
        should be discovered.
        """
        with NamedTemporaryFile(dir=self.empty_dir, suffix=".py") as f:
            f.write(b"# airflow")
            f.write(b"# DAG")
            f.flush()

            with conf_vars({('core', 'dags_folder'): self.empty_dir}):
                dagbag = models.DagBag(include_examples=False, safe_mode=True)

            self.assertEqual(len(dagbag.dagbag_stats), 1)
            self.assertEqual(
                dagbag.dagbag_stats[0].file,
                "/{}".format(os.path.basename(f.name))) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_get_dag_fileloc(self):
        """
        Test that fileloc is correctly set when we load example DAGs,
        specifically SubDAGs and packaged DAGs.
        """
        dagbag = models.DagBag(dag_folder=self.empty_dir, include_examples=True)
        dagbag.process_file(os.path.join(TEST_DAGS_FOLDER, "test_zip.zip"))

        expected = {
            'example_bash_operator': 'airflow/example_dags/example_bash_operator.py',
            'example_subdag_operator': 'airflow/example_dags/example_subdag_operator.py',
            'example_subdag_operator.section-1': 'airflow/example_dags/subdags/subdag.py',
            'test_zip_dag': 'dags/test_zip.zip/test_zip.py'
        }

        for dag_id, path in expected.items():
            dag = dagbag.get_dag(dag_id)
            self.assertTrue(dag.fileloc.endswith(path)) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_deactivate_unknown_dags(self):
        """
        Test that dag_ids not passed into deactivate_unknown_dags
        are deactivated when function is invoked
        """
        dagbag = DagBag(include_examples=True)
        dag_id = "test_deactivate_unknown_dags"
        expected_active_dags = dagbag.dags.keys()

        model_before = DagModel(dag_id=dag_id, is_active=True)
        with create_session() as session:
            session.merge(model_before)

        models.DAG.deactivate_unknown_dags(expected_active_dags)

        after_model = DagModel.get_dagmodel(dag_id)
        self.assertTrue(model_before.is_active)
        self.assertFalse(after_model.is_active)

        # clean up
        with create_session() as session:
            session.query(DagModel).filter(DagModel.dag_id == 'test_deactivate_unknown_dags').delete() 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def dag_bag_multiple():
    """
    Create a DagBag containing two DAGs, linked by multiple ExternalTaskMarker.
    """
    dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False)
    daily_dag = DAG("daily_dag", start_date=DEFAULT_DATE, schedule_interval="@daily")
    agg_dag = DAG("agg_dag", start_date=DEFAULT_DATE, schedule_interval="@daily")
    dag_bag.bag_dag(daily_dag, None, daily_dag)
    dag_bag.bag_dag(agg_dag, None, agg_dag)

    daily_task = DummyOperator(task_id="daily_tas", dag=daily_dag)

    start = DummyOperator(task_id="start", dag=agg_dag)
    for i in range(25):
        task = ExternalTaskMarker(task_id=f"{daily_task.task_id}_{i}",
                                  external_dag_id=daily_dag.dag_id,
                                  external_task_id=daily_task.task_id,
                                  execution_date="{{ macros.ds_add(ds, -1 * %s) }}" % i,
                                  dag=agg_dag)
        start >> task

    yield dag_bag 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_should_parse_only_unpaused_dags(self):
        dag_file = os.path.join(
            os.path.dirname(os.path.realpath(__file__)), '../dags/test_multiple_dags.py'
        )
        dag_file_processor = DagFileProcessor(dag_ids=[], log=mock.MagicMock())
        dagbag = DagBag(dag_folder=dag_file, include_examples=False)
        dagbag.sync_to_db()
        with create_session() as session:
            session.query(TaskInstance).delete()
            (
                session.query(DagModel)
                .filter(DagModel.dag_id == "test_multiple_dags__dag_1")
                .update({DagModel.is_paused: True}, synchronize_session=False)
            )

        simple_dags, import_errors_count = dag_file_processor.process_file(
            file_path=dag_file, failure_callback_requests=[]
        )
        with create_session() as session:
            tis = session.query(TaskInstance).all()

        self.assertEqual(0, import_errors_count)
        self.assertEqual(['test_multiple_dags__dag_2'], [dag.dag_id for dag in simple_dags])
        self.assertEqual({'test_multiple_dags__dag_2'}, {ti.dag_id for ti in tis}) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_process_dags_queries_count(
        self, expected_query_counts, dag_count, task_count, start_ago, schedule_interval, shape
    ):
        with mock.patch.dict("os.environ", {
            "PERF_DAGS_COUNT": str(dag_count),
            "PERF_TASKS_COUNT": str(task_count),
            "PERF_START_AGO": start_ago,
            "PERF_SCHEDULE_INTERVAL": schedule_interval,
            "PERF_SHAPE": shape,
        }), conf_vars({
            ('scheduler', 'use_job_schedule'): 'True',
        }):
            dagbag = DagBag(dag_folder=ELASTIC_DAG_FILE, include_examples=False)
            processor = DagFileProcessor([], mock.MagicMock())
            for expected_query_count in expected_query_counts:
                with assert_queries_count(expected_query_count):
                    processor._process_dags(dagbag.dags.values()) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def check_and_get_dag(dag_id: str, task_id: Optional[str] = None) -> DagModel:
    """Checks that DAG exists and in case it is specified that Task exist"""
    dag_model = DagModel.get_current(dag_id)
    if dag_model is None:
        raise DagNotFound("Dag id {} not found in DagModel".format(dag_id))

    dagbag = DagBag(
        dag_folder=dag_model.fileloc,
        store_serialized_dags=conf.getboolean('core', 'store_serialized_dags')
    )
    dag = dagbag.get_dag(dag_id)  # prefetch dag if it is stored serialized
    if dag_id not in dagbag.dags:
        error_message = "Dag id {} not found".format(dag_id)
        raise DagNotFound(error_message)
    if task_id and not dag.has_task(task_id):
        error_message = 'Task {} not found in dag {}'.format(task_id, dag_id)
        raise TaskNotFound(error_message)
    return dag 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def is_arg_not_provided(arg):
        return arg is None or arg == ""

    # Get the DagBag which has a list of all the current Dags 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def get_dagbag():
        return DagBag() 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def make_example_dags(module_path):
    """Loads DAGs from a module for test."""
    dagbag = DagBag(module_path)
    return dagbag.dags 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def setUpClass(cls):
        super().setUpClass()
        cls.dagbag = models.DagBag(include_examples=True)
        cls.app.dag_bag = cls.dagbag
        DAG.bulk_sync_to_db(cls.dagbag.dags.values()) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_view_uses_existing_dagbag(self, endpoint):
        """
        Test that Graph, Tree & Dag Details View uses the DagBag already created in views.py
        instead of creating a new one.
        """
        url = f'{endpoint}?dag_id=example_bash_operator'
        resp = self.client.get(url, follow_redirects=True)
        self.check_content_in_response('example_bash_operator', resp) 

# 需要导入模块: from airflow import models [as 别名]
# 或者: from airflow.models import DagBag [as 别名]
def test_code_from_db(self):
        from airflow.models.dagcode import DagCode
        dag = models.DagBag(include_examples=True).get_dag("example_bash_operator")
        DagCode(dag.fileloc, DagCode._get_code_from_file(dag.fileloc)).sync_to_db()
        url = 'code?dag_id=example_bash_operator'
        resp = self.client.get(url)
        self.check_content_not_in_response('Failed to load file', resp)
        self.check_content_in_response('example_bash_operator', resp)