Add Looker PDT operators (#20882)

apache · Mar 7, 2022 · 6db9b00 · 6db9b00
1 parent db76f0b
commit 6db9b00
Show file tree

Hide file tree

Showing 12 changed files with 1,037 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_looker.py b/airflow/providers/google/cloud/example_dags/example_looker.py
@@ -0,0 +1,63 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Example Airflow DAG that show how to use various Looker
+operators to submit PDT materialization job and manage it.
+"""
+
+from datetime import datetime
+
+from airflow import models
+from airflow.providers.google.cloud.operators.looker import LookerStartPdtBuildOperator
+from airflow.providers.google.cloud.sensors.looker import LookerCheckPdtBuildSensor
+
+with models.DAG(
+ dag_id='example_gcp_looker',
+ schedule_interval=None,
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ # [START cloud_looker_async_start_pdt_sensor]
+ start_pdt_task_async = LookerStartPdtBuildOperator(
+ task_id='start_pdt_task_async',
+ looker_conn_id='your_airflow_connection_for_looker',
+ model='your_lookml_model',
+ view='your_lookml_view',
+ asynchronous=True,
+ )
+
+ check_pdt_task_async_sensor = LookerCheckPdtBuildSensor(
+ task_id='check_pdt_task_async_sensor',
+ looker_conn_id='your_airflow_connection_for_looker',
+ materialization_id=start_pdt_task_async.output,
+ poke_interval=10,
+ )
+ # [END cloud_looker_async_start_pdt_sensor]
+
+ # [START how_to_cloud_looker_start_pdt_build_operator]
+ build_pdt_task = LookerStartPdtBuildOperator(
+ task_id='build_pdt_task',
+ looker_conn_id='your_airflow_connection_for_looker',
+ model='your_lookml_model',
+ view='your_lookml_view',
+ )
+ # [END how_to_cloud_looker_start_pdt_build_operator]
+
+ start_pdt_task_async >> check_pdt_task_async_sensor
+
+ build_pdt_task
diff --git a/airflow/providers/google/cloud/hooks/looker.py b/airflow/providers/google/cloud/hooks/looker.py
@@ -0,0 +1,258 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+"""This module contains a Google Cloud Looker hook."""
+
+import json
+import time
+from enum import Enum
+from typing import Dict, Optional
+
+from looker_sdk.rtl import api_settings, auth_session, requests_transport, serialize
+from looker_sdk.sdk.api40 import methods as methods40
+from packaging.version import parse as parse_version
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.models.connection import Connection
+from airflow.version import version
+
+
+class LookerHook(BaseHook):
+ """Hook for Looker APIs."""
+
+ def __init__(
+ self,
+ looker_conn_id: str,
+ ) -> None:
+ super().__init__()
+ self.looker_conn_id = looker_conn_id
+ # source is used to track origin of the requests
+ self.source = f'airflow:{version}'
+
+ def start_pdt_build(
+ self,
+ model: str,
+ view: str,
+ query_params: Optional[Dict] = None,
+ ):
+ """
+ Submits a PDT materialization job 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.
+ """
+ self.log.info("Submitting PDT materialization job. Model: '%s', view: '%s'.", model, view)
+ self.log.debug("PDT materialization job source: '%s'.", self.source)
+
+ sdk = self.get_looker_sdk()
+ looker_ver = sdk.versions().looker_release_version
+ if parse_version(looker_ver) < parse_version("22.2.0"):
+ raise AirflowException(f'This API requires Looker version 22.2+. Found: {looker_ver}.')
+
+ # unpack query_params dict into kwargs (if not None)
+ if query_params:
+ resp = sdk.start_pdt_build(model_name=model, view_name=view, source=self.source, **query_params)
+ else:
+ resp = sdk.start_pdt_build(model_name=model, view_name=view, source=self.source)
+
+ self.log.info("Start PDT build response: '%s'.", resp)
+
+ return resp
+
+ def check_pdt_build(
+ self,
+ materialization_id: str,
+ ):
+ """
+ Gets the PDT materialization job status from Looker.
+
+ :param materialization_id: Required. The materialization id to check status for.
+ """
+ self.log.info("Requesting PDT materialization job status. Job id: %s.", materialization_id)
+
+ sdk = self.get_looker_sdk()
+ resp = sdk.check_pdt_build(materialization_id=materialization_id)
+
+ self.log.info("Check PDT build response: '%s'.", resp)
+ return resp
+
+ def pdt_build_status(
+ self,
+ materialization_id: str,
+ ) -> Dict:
+ """
+ Gets the PDT materialization job status.
+
+ :param materialization_id: Required. The materialization id to check status for.
+ """
+ resp = self.check_pdt_build(materialization_id=materialization_id)
+
+ status_json = resp['resp_text']
+ status_dict = json.loads(status_json)
+
+ self.log.info(
+ "PDT materialization job id: %s. Status: '%s'.", materialization_id, status_dict['status']
+ )
+
+ return status_dict
+
+ def stop_pdt_build(
+ self,
+ materialization_id: str,
+ ):
+ """
+ Starts a PDT materialization job cancellation request.
+
+ :param materialization_id: Required. The materialization id to stop.
+ """
+ self.log.info("Stopping PDT materialization. Job id: %s.", materialization_id)
+ self.log.debug("PDT materialization job source: '%s'.", self.source)
+
+ sdk = self.get_looker_sdk()
+ resp = sdk.stop_pdt_build(materialization_id=materialization_id, source=self.source)
+
+ self.log.info("Stop PDT build response: '%s'.", resp)
+ return resp
+
+ def wait_for_job(
+ self,
+ materialization_id: str,
+ wait_time: int = 10,
+ timeout: Optional[int] = None,
+ ) -> None:
+ """
+ Helper method which polls a PDT materialization job to check if it finishes.
+
+ :param materialization_id: Required. The materialization id to wait for.
+ :param wait_time: Optional. Number of seconds between checks.
+ :param timeout: Optional. How many seconds wait for job to be ready.
+ Used only if ``asynchronous`` is False.
+ """
+ self.log.info('Waiting for PDT materialization job to complete. Job id: %s.', materialization_id)
+
+ status = None
+ start = time.monotonic()
+
+ while status not in (
+ JobStatus.DONE.value,
+ JobStatus.ERROR.value,
+ JobStatus.CANCELLED.value,
+ JobStatus.UNKNOWN.value,
+ ):
+
+ if timeout and start + timeout < time.monotonic():
+ self.stop_pdt_build(materialization_id=materialization_id)
+ raise AirflowException(
+ f"Timeout: PDT materialization job is not ready after {timeout}s. "
+ f"Job id: {materialization_id}."
+ )
+
+ time.sleep(wait_time)
+
+ status_dict = self.pdt_build_status(materialization_id=materialization_id)
+ status = status_dict['status']
+
+ if status == JobStatus.ERROR.value:
+ msg = status_dict['message']
+ raise AirflowException(
+ f'PDT materialization job failed. Job id: {materialization_id}. Message:\n"{msg}"'
+ )
+ if status == JobStatus.CANCELLED.value:
+ raise AirflowException(f'PDT materialization job was cancelled. Job id: {materialization_id}.')
+ if status == JobStatus.UNKNOWN.value:
+ raise AirflowException(
+ f'PDT materialization job has unknown status. Job id: {materialization_id}.'
+ )
+
+ self.log.info('PDT materialization job completed successfully. Job id: %s.', materialization_id)
+
+ def get_looker_sdk(self):
+ """Returns Looker SDK client for Looker API 4.0."""
+
+ conn = self.get_connection(self.looker_conn_id)
+ settings = LookerApiSettings(conn)
+
+ transport = requests_transport.RequestsTransport.configure(settings)
+ return methods40.Looker40SDK(
+ auth_session.AuthSession(settings, transport, serialize.deserialize40, "4.0"),
+ serialize.deserialize40,
+ serialize.serialize,
+ transport,
+ "4.0",
+ )
+
+
+class LookerApiSettings(api_settings.ApiSettings):
+ """Custom implementation of Looker SDK's `ApiSettings` class."""
+
+ def __init__(
+ self,
+ conn: Connection,
+ ) -> None:
+ self.conn = conn # need to init before `read_config` is called in super
+ super().__init__()
+
+ def read_config(self):
+ """
+ Overrides the default logic of getting connection settings. Fetches
+ the connection settings from Airflow's connection object.
+ """
+
+ config = {}
+
+ if self.conn.host is None:
+ raise AirflowException(f'No `host` was supplied in connection: {self.conn.id}.')
+
+ if self.conn.port:
+ config["base_url"] = f"{self.conn.host}:{self.conn.port}" # port is optional
+ else:
+ config["base_url"] = self.conn.host
+
+ if self.conn.login:
+ config["client_id"] = self.conn.login
+ else:
+ raise AirflowException(f'No `login` was supplied in connection: {self.conn.id}.')
+
+ if self.conn.password:
+ config["client_secret"] = self.conn.password
+ else:
+ raise AirflowException(f'No `password` was supplied in connection: {self.conn.id}.')
+
+ extras = self.conn.extra_dejson # type: Dict
+
+ if 'verify_ssl' in extras:
+ config["verify_ssl"] = extras["verify_ssl"] # optional
+
+ if 'timeout' in extras:
+ config["timeout"] = extras["timeout"] # optional
+
+ return config
+
+
+class JobStatus(Enum):
+ """The job status string."""
+
+ QUEUED = "added"
+ PENDING = 'pending'
+ RUNNING = 'running'
+ CANCELLED = 'killed'
+ DONE = 'complete'
+ ERROR = 'error'
+ UNKNOWN = 'unknown'