Revert "Revert BigQueryAsyncExtractor to release 1.3.0 (#332)"

This reverts commit 39102d0.

Author: pankajkoti

astronomer/providers/google/cloud/extractors/__init__.py

@@ -0,0 +1,81 @@
+from typing import Any, List, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.models.taskinstance import TaskInstance
+from airflow.utils.log.logging_mixin import LoggingMixin
+from google.cloud.bigquery import Client
+from openlineage.airflow.extractors.base import BaseExtractor, TaskMetadata
+from openlineage.airflow.utils import get_job_name
+from openlineage.common.provider.bigquery import BigQueryDatasetsProvider
+
+from astronomer.providers.google.cloud.operators.bigquery import (
+    BigQueryInsertJobOperatorAsync,
+)
+
+
+class BigQueryAsyncExtractor(BaseExtractor, LoggingMixin):
+    """
+    This extractor provides visibility on the metadata of a BigQuery Insert Job
+    including ``billedBytes``, ``rowCount``, ``size``, etc. submitted from a
+    ``BigQueryInsertJobOperatorAsync`` operator.
+    """
+
+    def __init__(self, operator: BigQueryInsertJobOperatorAsync):
+        super().__init__(operator)
+        self._big_query_client = self._get_big_query_client()
+
+    def _get_big_query_client(self) -> Client:
+        """
+        Gets the BigQuery client to fetch job metadata.
+        The method checks whether a connection hook is available with the Airflow configuration for the operator, and
+        if yes, returns the same connection. Otherwise, returns the Client instance of``google.cloud.bigquery``.
+        """
+        if hasattr(self.operator, "hook") and self.operator.hook:
+            hook = self.operator.hook
+            return hook.get_client(project_id=hook.project_id, location=hook.location)
+        return Client()
+
+    def _get_xcom_bigquery_job_id(self, task_instance: TaskInstance) -> Any:
+        """
+        Pulls the BigQuery Job ID from XCOM for the task instance whose metadata needs to be extracted.
+
+        :param task_instance: Instance of the Airflow task whose BigQuery ``job_id`` needs to be pulled from XCOM.
+        """
+        bigquery_job_id = task_instance.xcom_pull(task_ids=task_instance.task_id, key="job_id")
+        if not bigquery_job_id:
+            raise AirflowException("Could not pull relevant BigQuery job ID from XCOM")
+        self.log.debug("Big Query Job Id %s", bigquery_job_id)
+        return bigquery_job_id
+
+    @classmethod
+    def get_operator_classnames(cls) -> List[str]:
+        """Returns the list of operators this extractor works on."""
+        return ["BigQueryInsertJobOperatorAsync"]
+
+    def extract(self) -> Optional[TaskMetadata]:
+        """Empty extract implementation for the abstractmethod of the ``BaseExtractor`` class."""
+        return None
+
+    def extract_on_complete(self, task_instance: TaskInstance) -> Optional[TaskMetadata]:
+        """
+        Callback on task completion to fetch metadata extraction details that are to be pushed to the Lineage server.
+
+        :param task_instance: Instance of the Airflow task whose metadata needs to be extracted.
+        """
+        try:
+            bigquery_job_id = self._get_xcom_bigquery_job_id(task_instance)
+        except AirflowException as ae:
+            exception_message = str(ae)
+            self.log.exception("%s", exception_message)
+            return TaskMetadata(name=get_job_name(task=self.operator))
+        stats = BigQueryDatasetsProvider(client=self._big_query_client).get_facets(bigquery_job_id)
+        inputs = stats.inputs
+        output = stats.output
+        run_facets = stats.run_facets
+
+        return TaskMetadata(
+            name=get_job_name(task=self.operator),
+            inputs=[ds.to_openlineage_dataset() for ds in inputs],
+            outputs=[output.to_openlineage_dataset()] if output else [],
+            run_facets=run_facets,
+        )

astronomer/providers/google/cloud/extractors/bigquery_async_extractor.py

@@ -76,6 +76,10 @@ apache.hive =
     impyla
 microsoft.azure =
     apache-airflow-providers-microsoft-azure
+
+# If in future we move Openlineage extractors out of the repo, this dependency should be removed
+openlineage =
+    openlineage-airflow==0.6.2
 docs =
     sphinx
     sphinx-autoapi
@@ -128,6 +132,7 @@ all =
     kubernetes_asyncio
     paramiko
     impyla
+    openlineage-airflow==0.6.2
 
 [options.packages.find]
 include =

setup.cfg

@@ -0,0 +1,171 @@
+import json
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+from airflow.exceptions import TaskDeferred
+from airflow.models.dagrun import DagRun
+from airflow.models.taskinstance import TaskInstance
+from airflow.utils.timezone import datetime
+from airflow.utils.types import DagRunType
+from openlineage.client.facet import OutputStatisticsOutputDatasetFacet
+from openlineage.common.dataset import Dataset, Source
+from openlineage.common.provider.bigquery import (
+    BigQueryFacets,
+    BigQueryJobRunFacet,
+    BigQueryStatisticsDatasetFacet,
+)
+
+from astronomer.providers.google.cloud.extractors.bigquery_async_extractor import (
+    BigQueryAsyncExtractor,
+)
+from astronomer.providers.google.cloud.operators.bigquery import (
+    BigQueryInsertJobOperatorAsync,
+)
+
+TEST_DATASET_LOCATION = "EU"
+TEST_GCP_PROJECT_ID = "test-project"
+TEST_DATASET = "test-dataset"
+TEST_TABLE = "test-table"
+EXECUTION_DATE = datetime(2022, 1, 1, 0, 0, 0)
+INSERT_DATE = EXECUTION_DATE.strftime("%Y-%m-%d")
+INSERT_ROWS_QUERY = (
+    f"INSERT {TEST_DATASET}.{TEST_TABLE} VALUES "
+    f"(42, 'monthy python', '{INSERT_DATE}'), "
+    f"(42, 'fishy fish', '{INSERT_DATE}');"
+)
+
+INPUT_STATS = [
+    Dataset(
+        source=Source(scheme="bigquery"),
+        name=f"astronomer-airflow-providers.{TEST_DATASET}.{TEST_TABLE}",
+        fields=[],
+        custom_facets={},
+        input_facets={},
+        output_facets={},
+    )
+]
+
+OUTPUT_STATS = Dataset(
+    source=Source(scheme="bigquery"),
+    name=f"astronomer-airflow-providers.{TEST_DATASET}.{TEST_TABLE}",
+    fields=[],
+    custom_facets={"stats": BigQueryStatisticsDatasetFacet(rowCount=2, size=0)},
+    input_facets={},
+    output_facets={"outputStatistics": OutputStatisticsOutputDatasetFacet(rowCount=2, size=0)},
+)
+
+with open("tests/google/cloud/extractors/job_details.json") as jd_json:
+    JOB_PROPERTIES = json.load(jd_json)
+
+RUN_FACETS = {
+    "bigQuery_job": BigQueryJobRunFacet(billedBytes=0, cached=False, properties=json.dumps(JOB_PROPERTIES))
+}
+
+
+@pytest.fixture
+def context():
+    """
+    Creates an empty context.
+    """
+    context = {}
+    yield context
+
+
+@mock.patch("astronomer.providers.google.cloud.operators.bigquery._BigQueryHook")
+@mock.patch("airflow.models.TaskInstance.xcom_pull")
+@mock.patch("openlineage.common.provider.bigquery.BigQueryDatasetsProvider.get_facets")
+def test_extract_on_complete(mock_bg_dataset_provider, mock_xcom_pull, mock_hook):
+    """
+    Tests that  the custom extractor's implementation for the BigQueryInsertJobOperatorAsync is able to process the
+    operator's metadata that needs to be extracted as per OpenLineage.
+    """
+    configuration = {
+        "query": {
+            "query": INSERT_ROWS_QUERY,
+            "useLegacySql": False,
+        }
+    }
+    job_id = "123456"
+    mock_hook.return_value.insert_job.return_value = MagicMock(job_id=job_id, error_result=False)
+    mock_bg_dataset_provider.return_value = BigQueryFacets(
+        run_facets=RUN_FACETS, inputs=INPUT_STATS, output=OUTPUT_STATS
+    )
+
+    task_id = "insert_query_job"
+    operator = BigQueryInsertJobOperatorAsync(
+        task_id=task_id,
+        configuration=configuration,
+        location=TEST_DATASET_LOCATION,
+        job_id=job_id,
+        project_id=TEST_GCP_PROJECT_ID,
+    )
+
+    task_instance = TaskInstance(task=operator)
+    with pytest.raises(TaskDeferred):
+        operator.execute(context)
+
+    bq_extractor = BigQueryAsyncExtractor(operator)
+    task_meta_extract = bq_extractor.extract()
+    assert task_meta_extract is None
+
+    task_meta = bq_extractor.extract_on_complete(task_instance)
+
+    mock_xcom_pull.assert_called_once_with(task_ids=task_instance.task_id, key="job_id")
+
+    assert task_meta.name == f"adhoc_airflow.{task_id}"
+
+    assert task_meta.inputs[0].facets["dataSource"].name == INPUT_STATS[0].source.scheme
+    assert task_meta.inputs[0].name == INPUT_STATS[0].name
+
+    assert task_meta.outputs[0].name == OUTPUT_STATS.name
+    assert task_meta.outputs[0].facets["stats"].rowCount == 2
+    assert task_meta.outputs[0].facets["stats"].size == 0
+
+    assert task_meta.run_facets["bigQuery_job"].billedBytes == 0
+    run_facet_properties = json.loads(task_meta.run_facets["bigQuery_job"].properties)
+    assert run_facet_properties == JOB_PROPERTIES
+
+
+def test_extractor_works_on_operator():
+    """Tests that the custom extractor implementation is available for the BigQueryInsertJobOperatorAsync Operator."""
+    task_id = "insert_query_job"
+    operator = BigQueryInsertJobOperatorAsync(task_id=task_id, configuration={})
+    assert type(operator).__name__ in BigQueryAsyncExtractor.get_operator_classnames()
+
+
+@mock.patch("astronomer.providers.google.cloud.operators.bigquery._BigQueryHook")
+def test_unavailable_xcom_raises_exception(mock_hook):
+    """
+    Tests that an exception is raised when the custom extractor is not available to retrieve required XCOM for the
+    BigQueryInsertJobOperatorAsync Operator.
+    """
+    configuration = {
+        "query": {
+            "query": INSERT_ROWS_QUERY,
+            "useLegacySql": False,
+        }
+    }
+    job_id = "123456"
+    mock_hook.return_value.insert_job.return_value = MagicMock(job_id=job_id, error_result=False)
+    task_id = "insert_query_job"
+    operator = BigQueryInsertJobOperatorAsync(
+        task_id=task_id,
+        configuration=configuration,
+        location=TEST_DATASET_LOCATION,
+        job_id=job_id,
+        project_id=TEST_GCP_PROJECT_ID,
+    )
+
+    task_instance = TaskInstance(task=operator)
+    execution_date = datetime(2022, 1, 1, 0, 0, 0)
+    task_instance.run_id = DagRun.generate_run_id(DagRunType.MANUAL, execution_date)
+
+    with pytest.raises(TaskDeferred):
+        operator.execute(context)
+    bq_extractor = BigQueryAsyncExtractor(operator)
+    with mock.patch.object(bq_extractor.log, "exception") as mock_log_exception:
+        task_meta = bq_extractor.extract_on_complete(task_instance)
+
+    mock_log_exception.assert_called_with("%s", "Could not pull relevant BigQuery job ID from XCOM")
+    assert task_meta.name == f"adhoc_airflow.{task_id}"