[GitHub] [airflow] mik-laj commented on a change in pull request #20882: Add Looker PDT operators

[GitBox <gi...@apache.org>]

mik-laj commented on a change in pull request #20882:
URL: https://github.com/apache/airflow/pull/20882#discussion_r817106938



##########
File path: airflow/providers/google/cloud/hooks/looker.py
##########
@@ -0,0 +1,268 @@
+#
+# 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
+        self.source = self.get_api_source()
+
+    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():
+                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_api_source(self):
+        """Returns origin of the API call."""
+
+        # Infer API source based on Airflow version, e.g.:
+        # "2.1.4+composer" -> Cloud Composer
+        # "2.1.4" -> Airflow
+
+        self.log.debug(
+            "Got the following version for connection id: '%s'. Version: '%s'.", self.looker_conn_id, version
+        )
+
+        return 'composer' if 'composer' in version else 'airflow'

Review comment:
       >  Does it mean that a user need to specify it? 
   
   Yes. The Cloud Composer team needs to add this variable in their Docker image. 
   ```
   ENV AIRFLOW_GOOGLE_PARTNER=composer
   ```
   
   > If so, is there a ways to identify the type of Airflow (standalone/managed) based in a build-in data?
   
   No, because we are not involved in the development of a managed versions of Airflow and we want to be vendor-vendor neutral. This means that we cannot change the behavior of the code depending on which vendor is running the code. If a vendor wants to make changes to the behavior of the software, this change should only be made in the vendor's environment.
   




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org