def get_link( self, operator: "AbstractOperator", dttm: Optional[datetime] = None, *, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: """ Get link to qubole command result page. :param operator: operator :param dttm: datetime :return: url link """ conn = BaseHook.get_connection( getattr(operator, "qubole_conn_id", None) or operator.kwargs['qubole_conn_id'] # type: ignore[attr-defined] ) if conn and conn.host: host = re.sub(r'api$', 'v2/analyze?command_id=', conn.host) else: host = 'https://api.qubole.com/v2/analyze?command_id=' if ti_key: qds_command_id = XCom.get_one(key='qbol_cmd_id', ti_key=ti_key) else: assert dttm qds_command_id = XCom.get_one(key='qbol_cmd_id', dag_id=operator.dag_id, task_id=operator.task_id, execution_date=dttm) url = host + str(qds_command_id) if qds_command_id else '' return url
def get_link( self, operator, dttm=None, *, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key: run_id = XCom.get_one(key="run_id", ti_key=ti_key) else: assert dttm run_id = XCom.get_one( key="run_id", dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) conn = BaseHook.get_connection(operator.azure_data_factory_conn_id) subscription_id = conn.extra_dejson[ "extra__azure_data_factory__subscriptionId"] # Both Resource Group Name and Factory Name can either be declared in the Azure Data Factory # connection or passed directly to the operator. resource_group_name = operator.resource_group_name or conn.extra_dejson.get( "extra__azure_data_factory__resource_group_name") factory_name = operator.factory_name or conn.extra_dejson.get( "extra__azure_data_factory__factory_name") url = ( f"https://adf.azure.com/en-us/monitoring/pipelineruns/{run_id}" f"?factory=/subscriptions/{subscription_id}/" f"resourceGroups/{resource_group_name}/providers/Microsoft.DataFactory/" f"factories/{factory_name}") return url
def get_link(self, operator, dttm=None, *, ti_key=None): if ti_key: job_run_url = XCom.get_one(key="job_run_url", ti_key=ti_key) else: assert dttm job_run_url = XCom.get_one(dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key="job_run_url") return job_run_url
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: """ Get link to EMR cluster. :param operator: operator :param dttm: datetime :return: url link """ if ti_key is not None: flow_id = XCom.get_value(key="return_value", ti_key=ti_key) else: assert dttm flow_id = XCom.get_one( key="return_value", dag_id=operator.dag_id, task_id=operator.task_id, execution_date=dttm ) return ( f'https://console.aws.amazon.com/elasticmapreduce/home#cluster-details:{flow_id}' if flow_id else '' )
def test_xcom_enable_pickle_type(self): json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test2" dag_id = "test_dag2" task_id = "test_task2" configuration.set("core", "enable_xcom_pickling", "True") XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, execution_date=execution_date) ret_value = XCom.get_one(key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date) self.assertEqual(ret_value, json_obj) session = settings.Session() ret_value = session.query(XCom).filter( XCom.key == key, XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.execution_date == execution_date).first().value self.assertEqual(ret_value, json_obj)
def _get_result(dag, task_id, execution_date=None): execution_date = execution_date or days_ago(0) result = XCom.get_one(execution_date=execution_date, task_id=task_id, dag_id=dag.dag_id) assert result return result
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: """ Link to Amazon Web Services Console. :param operator: airflow operator :param ti_key: TaskInstance ID to return link for :param dttm: execution date. Uses for compatibility with Airflow 2.2 :return: link to external system """ if ti_key is not None: conf = XCom.get_value(key=self.key, ti_key=ti_key) elif not dttm: conf = {} else: conf = XCom.get_one( key=self.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return self.format_link(**conf) if conf else ""
def get_link(self, operator, dttm): job_run_url = XCom.get_one(dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, key="job_run_url") return job_run_url
def test_xcom_enable_pickle_type(self): json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test2" dag_id = "test_dag2" task_id = "test_task2" configuration.set("core", "enable_xcom_pickling", "True") XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, execution_date=execution_date) ret_value = XCom.get_one(key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date) self.assertEqual(ret_value, json_obj) session = settings.Session() ret_value = session.query(XCom).filter(XCom.key == key, XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.execution_date == execution_date ).first().value self.assertEqual(ret_value, json_obj)
def test_xcom_get_one_disable_pickle_type(self): json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test1" dag_id = "test_dag1" task_id = "test_task1" XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, execution_date=execution_date) ret_value = XCom.get_one(key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date) self.assertEqual(ret_value, json_obj) session = settings.Session() ret_value = session.query(XCom).filter( XCom.key == key, XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.execution_date == execution_date).first().value self.assertEqual(ret_value, json_obj)
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key: conf = XCom.get_one(key=self.key, ti_key=ti_key) else: assert dttm conf = XCom.get_one(key=self.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm) return (conf["url"].format(region=conf["region"], project_id=conf["project_id"], resource=conf["resource"]) if conf else "")
def get_link(self, operator: BaseOperator, dttm: datetime) -> str: conf = XCom.get_one( key=CloudComposerEnvironmentsLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return CLOUD_COMPOSER_ENVIRONMENTS_LINK.format( project_id=conf["project_id"]) if conf else ""
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key: list_conf = XCom.get_one(key=self.key, ti_key=ti_key) else: assert dttm list_conf = XCom.get_one( key=self.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return (list_conf["url"].format(project_id=list_conf["project_id"], ) if list_conf else "")
def get_link( self, operator, dttm: Optional[datetime] = None, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key: conf = XCom.get_one(key=self.key, ti_key=ti_key) else: assert dttm conf = XCom.get_one( key=self.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return self.format_str.format(**conf) if conf else ""
def get_link(self, operator: BaseOperator, dttm: datetime) -> str: conf = XCom.get_one( key=CloudComposerEnvironmentLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return (CLOUD_COMPOSER_DETAILS_LINK.format( project_id=conf["project_id"], region=conf['region'], environment_id=conf['environment_id']) if conf else "")
def get_link(self, operator: BaseOperator, dttm: datetime): conf = XCom.get_one( key=DataprocLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm ) return ( conf["url"].format( region=conf["region"], project_id=conf["project_id"], resource=conf["resource"] ) if conf else "" )
def get_link(self, operator: BaseOperator, dttm: datetime): list_conf = XCom.get_one( key=DataprocListLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return ( list_conf["url"].format( project_id=list_conf["project_id"], ) if list_conf else "" )
def get_link(self, operator: BaseOperator, dttm: datetime) -> str: conf = XCom.get_one( key=DataflowJobLink.key, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return ( DATAFLOW_JOB_LINK.format( project_id=conf["project_id"], region=conf['region'], job_id=conf['job_id'] ) if conf else "" )
def test_write_xcom(self): url_template = ('/api/experimental/dags/{}/tasks/{}' '/instances/{}/xcom/{}/{}') dag_id = 'example_bash_operator' task_id = 'xcom_task' key = 'xcom_key' value = 'xcom_value' execution_date = datetime.now().replace(microsecond=0) datetime_string = execution_date.isoformat() # Test Correct execution response = self.app.post( quote(url_template.format(dag_id, task_id, datetime_string, key, value)) ) self.assertEqual(200, response.status_code) # Check xcom value in database xcom_value = XCom.get_one(execution_date=execution_date, key=key, task_id=task_id, dag_id=dag_id) self.assertEqual(value, xcom_value, 'XCom value: expected {}, found {}'.format(value, xcom_value)) # Test error for nonexistent dag response = self.app.post( quote(url_template.format('does_not_exist_dag', task_id, datetime_string, key, value)) ) self.assertEqual(404, response.status_code) # Test error for bad datetime format response = self.app.post( quote(url_template.format(dag_id, task_id, 'not_a_datetime', key, value)) ) self.assertEqual(400, response.status_code)
def get_link(self, operator: BaseOperator, dttm: datetime) -> str: """ Get link to EMR cluster. :param operator: operator :param dttm: datetime :return: url link """ flow_id = XCom.get_one(key="return_value", dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm) return ( f'https://console.aws.amazon.com/elasticmapreduce/home#cluster-details:{flow_id}' if flow_id else '')
def get_link( self, operator, dttm=None, *, ti_key: Optional["TaskInstanceKey"] = None, ) -> str: if ti_key is not None: run_page_url = XCom.get_value(key=XCOM_RUN_PAGE_URL_KEY, ti_key=ti_key) else: assert dttm run_page_url = XCom.get_one( key=XCOM_RUN_PAGE_URL_KEY, dag_id=operator.dag.dag_id, task_id=operator.task_id, execution_date=dttm, ) return run_page_url
def test_write_xcom(self): url_template = ('/api/experimental/dags/{}/tasks/{}' '/instances/{}/xcom/{}/{}') dag_id = 'example_bash_operator' task_id = 'xcom_task' key = 'xcom_key' value = 'xcom_value' execution_date = datetime.now().replace(microsecond=0) datetime_string = execution_date.isoformat() # Test Correct execution response = self.app.post( quote( url_template.format(dag_id, task_id, datetime_string, key, value))) self.assertEqual(200, response.status_code) # Check xcom value in database xcom_value = XCom.get_one(execution_date=execution_date, key=key, task_id=task_id, dag_id=dag_id) self.assertEqual( value, xcom_value, 'XCom value: expected {}, found {}'.format(value, xcom_value)) # Test error for nonexistent dag response = self.app.post( quote( url_template.format('does_not_exist_dag', task_id, datetime_string, key, value))) self.assertEqual(404, response.status_code) # Test error for bad datetime format response = self.app.post( quote( url_template.format(dag_id, task_id, 'not_a_datetime', key, value))) self.assertEqual(400, response.status_code)