def execute(self, context: "Context") -> str:
self.hook = LookerHook(looker_conn_id=self.looker_conn_id)
resp = self.hook.start_pdt_build(
model=self.model,
view=self.view,
query_params=self.query_params,
)
self.materialization_id = resp.materialization_id
if not self.materialization_id:
raise AirflowException(
f'No `materialization_id` was returned for model: {self.model}, view: {self.view}.'
)
self.log.info(
"PDT materialization job submitted successfully. Job id: %s.",
self.materialization_id)
if not self.asynchronous:
self.hook.wait_for_job(
materialization_id=self.materialization_id,
wait_time=self.wait_time,
timeout=self.wait_timeout,
)
return self.materialization_id
def poke(self, context: "Context") -> bool:
self.hook = LookerHook(looker_conn_id=self.looker_conn_id)
if not self.materialization_id:
raise AirflowException('Invalid `materialization_id`.')
# materialization_id is templated var pulling output from start task
status_dict = self.hook.pdt_build_status(materialization_id=self.materialization_id)
status = status_dict['status']
if status == JobStatus.ERROR.value:
msg = status_dict["message"]
raise AirflowException(
f'PDT materialization job failed. Job id: {self.materialization_id}. Message:\n"{msg}"'
)
elif status == JobStatus.CANCELLED.value:
raise AirflowException(
f'PDT materialization job was cancelled. Job id: {self.materialization_id}.'
)
elif status == JobStatus.UNKNOWN.value:
raise AirflowException(
f'PDT materialization job has unknown status. Job id: {self.materialization_id}.'
)
elif status == JobStatus.DONE.value:
self.log.debug(
"PDT materialization job completed successfully. Job id: %s.", self.materialization_id
)
return True
self.log.info("Waiting for PDT materialization job to complete. Job id: %s.", self.materialization_id)
return False
class LookerCheckPdtBuildSensor(BaseSensorOperator):
"""
Check for the state of a previously submitted PDT materialization job.
:param materialization_id: Required. The materialization job ID to poll. (templated)
:param looker_conn_id: Required. The connection ID to use connecting to Looker.
:param cancel_on_kill: Optional. Flag which indicates whether cancel the hook's job or not,
when on_kill is called.
"""
template_fields = ["materialization_id"]
def __init__(self,
materialization_id: str,
looker_conn_id: str,
cancel_on_kill: bool = True,
**kwargs) -> None:
super().__init__(**kwargs)
self.materialization_id = materialization_id
self.looker_conn_id = looker_conn_id
self.cancel_on_kill = cancel_on_kill
self.hook: Optional[LookerHook] = None
def poke(self, context: "Context") -> bool:
self.hook = LookerHook(looker_conn_id=self.looker_conn_id)
# materialization_id is templated var pulling output from start task
status_dict = self.hook.pdt_build_status(
materialization_id=self.materialization_id)
status = status_dict['status']
if status == JobStatus.ERROR.value:
msg = status_dict["message"]
raise AirflowException(
f'PDT materialization job failed. Job id: {self.materialization_id}. Message:\n"{msg}"'
)
elif status == JobStatus.CANCELLED.value:
raise AirflowException(
f'PDT materialization job was cancelled. Job id: {self.materialization_id}.'
)
elif status == JobStatus.UNKNOWN.value:
raise AirflowException(
f'PDT materialization job has unknown status. Job id: {self.materialization_id}.'
)
elif status == JobStatus.DONE.value:
self.log.debug(
"PDT materialization job completed successfully. Job id: %s.",
self.materialization_id)
return True
self.log.info(
"Waiting for PDT materialization job to complete. Job id: %s.",
self.materialization_id)
return False
def on_kill(self):
if self.materialization_id and self.cancel_on_kill:
self.hook.stop_pdt_build(
materialization_id=self.materialization_id)
class LookerStartPdtBuildOperator(BaseOperator):
"""
Submits a PDT materialization job to Looker.
:param looker_conn_id: Required. The connection ID to use connecting to Looker.
:param model: Required. The model of the PDT to start building.
:param view: Required. The view of the PDT to start building.
:param query_params: Optional. Additional materialization parameters.
:param asynchronous: Optional. Flag indicating whether to wait for the job
to finish or return immediately.
This is useful for submitting long running jobs and
waiting on them asynchronously using the LookerCheckPdtBuildSensor
:param cancel_on_kill: Optional. Flag which indicates whether cancel the
hook's job or not, when on_kill is called.
:param wait_time: Optional. Number of seconds between checks for job to be
ready. Used only if ``asynchronous`` is False.
:param wait_timeout: Optional. How many seconds wait for job to be ready.
Used only if ``asynchronous`` is False.
"""
def __init__(
self,
looker_conn_id: str,
model: str,
view: str,
query_params: Optional[Dict] = None,
asynchronous: bool = False,
cancel_on_kill: bool = True,
wait_time: int = 10,
wait_timeout: Optional[int] = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.model = model
self.view = view
self.query_params = query_params
self.looker_conn_id = looker_conn_id
self.asynchronous = asynchronous
self.cancel_on_kill = cancel_on_kill
self.wait_time = wait_time
self.wait_timeout = wait_timeout
self.hook: Optional[LookerHook] = None
self.materialization_id: Optional[str] = None
def execute(self, context: "Context") -> str:
self.hook = LookerHook(looker_conn_id=self.looker_conn_id)
resp = self.hook.start_pdt_build(
model=self.model,
view=self.view,
query_params=self.query_params,
)
self.materialization_id = resp.materialization_id
if not self.materialization_id:
raise AirflowException(
f'No `materialization_id` was returned for model: {self.model}, view: {self.view}.'
)
self.log.info(
"PDT materialization job submitted successfully. Job id: %s.",
self.materialization_id)
if not self.asynchronous:
self.hook.wait_for_job(
materialization_id=self.materialization_id,
wait_time=self.wait_time,
timeout=self.wait_timeout,
)
return self.materialization_id
def on_kill(self):
if self.materialization_id and self.cancel_on_kill:
self.hook.stop_pdt_build(
materialization_id=self.materialization_id)