acryl-datahub 0.15.0.6rc3__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/ingestion/run/pipeline.py

@@ -9,7 +9,7 @@ import sys
 import threading
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, Iterator, List, Optional, cast
+from typing import Any, Dict, Iterable, Iterator, List, Optional
 
 import click
 import humanfriendly
@@ -26,7 +26,7 @@ from datahub.ingestion.api.common import EndOfStream, PipelineContext, RecordEnv
 from datahub.ingestion.api.global_context import set_graph_context
 from datahub.ingestion.api.pipeline_run_listener import PipelineRunListener
 from datahub.ingestion.api.report import Report
-from datahub.ingestion.api.sink import Sink, SinkReport, WriteCallback
+from datahub.ingestion.api.sink import Sink, SinkReport
 from datahub.ingestion.api.source import Extractor, Source
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.extractor.extractor_registry import extractor_registry
@@ -35,15 +35,15 @@ from datahub.ingestion.reporting.reporting_provider_registry import (
     reporting_provider_registry,
 )
 from datahub.ingestion.run.pipeline_config import PipelineConfig, ReporterConfig
+from datahub.ingestion.run.sink_callback import DeadLetterQueueCallback, LoggingCallback
 from datahub.ingestion.sink.datahub_rest import DatahubRestSink
-from datahub.ingestion.sink.file import FileSink, FileSinkConfig
 from datahub.ingestion.sink.sink_registry import sink_registry
 from datahub.ingestion.source.source_registry import source_registry
 from datahub.ingestion.transformer.system_metadata_transformer import (
     SystemMetadataTransformer,
 )
 from datahub.ingestion.transformer.transform_registry import transform_registry
-from datahub.metadata.schema_classes import MetadataChangeProposalClass
+from datahub.sdk._attribution import KnownAttribution, change_default_attribution
 from datahub.telemetry import stats
 from datahub.telemetry.telemetry import telemetry_instance
 from datahub.utilities._custom_package_loader import model_version_name
@@ -57,68 +57,6 @@ logger = logging.getLogger(__name__)
 _REPORT_PRINT_INTERVAL_SECONDS = 60
 
 
-class LoggingCallback(WriteCallback):
-    def __init__(self, name: str = "") -> None:
-        super().__init__()
-        self.name = name
-
-    def on_success(
-        self, record_envelope: RecordEnvelope, success_metadata: dict
-    ) -> None:
-        logger.debug(
-            f"{self.name} sink wrote workunit {record_envelope.metadata['workunit_id']}"
-        )
-
-    def on_failure(
-        self,
-        record_envelope: RecordEnvelope,
-        failure_exception: Exception,
-        failure_metadata: dict,
-    ) -> None:
-        logger.error(
-            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}",
-            extra={"failure_metadata": failure_metadata},
-            exc_info=failure_exception,
-        )
-
-
-class DeadLetterQueueCallback(WriteCallback):
-    def __init__(self, ctx: PipelineContext, config: Optional[FileSinkConfig]) -> None:
-        if not config:
-            config = FileSinkConfig.parse_obj({"filename": "failed_events.json"})
-        self.file_sink: FileSink = FileSink(ctx, config)
-        self.logging_callback = LoggingCallback(name="failure-queue")
-        logger.info(f"Failure logging enabled. Will log to {config.filename}.")
-
-    def on_success(
-        self, record_envelope: RecordEnvelope, success_metadata: dict
-    ) -> None:
-        pass
-
-    def on_failure(
-        self,
-        record_envelope: RecordEnvelope,
-        failure_exception: Exception,
-        failure_metadata: dict,
-    ) -> None:
-        if "workunit_id" in record_envelope.metadata:
-            if isinstance(record_envelope.record, MetadataChangeProposalClass):
-                mcp = cast(MetadataChangeProposalClass, record_envelope.record)
-                if mcp.systemMetadata:
-                    if not mcp.systemMetadata.properties:
-                        mcp.systemMetadata.properties = {}
-                    if "workunit_id" not in mcp.systemMetadata.properties:
-                        # update the workunit id
-                        mcp.systemMetadata.properties["workunit_id"] = (
-                            record_envelope.metadata["workunit_id"]
-                        )
-                record_envelope.record = mcp
-        self.file_sink.write_record_async(record_envelope, self.logging_callback)
-
-    def close(self) -> None:
-        self.file_sink.close()
-
-
 class PipelineInitError(Exception):
     pass
 
@@ -236,76 +174,99 @@ class Pipeline:
         self.last_time_printed = int(time.time())
         self.cli_report = CliReport()
 
-        self.graph = None
-        with _add_init_error_context("connect to DataHub"):
-            if self.config.datahub_api:
-                self.graph = DataHubGraph(self.config.datahub_api)
-                self.graph.test_connection()
-
-        with _add_init_error_context("set up framework context"):
-            self.ctx = PipelineContext(
-                run_id=self.config.run_id,
-                graph=self.graph,
-                pipeline_name=self.config.pipeline_name,
-                dry_run=dry_run,
-                preview_mode=preview_mode,
-                pipeline_config=self.config,
-            )
-
-        if self.config.sink is None:
-            logger.info(
-                "No sink configured, attempting to use the default datahub-rest sink."
-            )
-            with _add_init_error_context("configure the default rest sink"):
-                self.sink_type = "datahub-rest"
-                self.sink = _make_default_rest_sink(self.ctx)
-        else:
-            self.sink_type = self.config.sink.type
-            with _add_init_error_context(
-                f"find a registered sink for type {self.sink_type}"
-            ):
-                sink_class = sink_registry.get(self.sink_type)
-
-            with _add_init_error_context(f"configure the sink ({self.sink_type})"):
-                sink_config = self.config.sink.dict().get("config") or {}
-                self.sink = sink_class.create(sink_config, self.ctx)
-                logger.debug(f"Sink type {self.sink_type} ({sink_class}) configured")
-        logger.info(f"Sink configured successfully. {self.sink.configured()}")
-
-        if self.graph is None and isinstance(self.sink, DatahubRestSink):
-            with _add_init_error_context("setup default datahub client"):
-                self.graph = self.sink.emitter.to_graph()
-                self.graph.test_connection()
-        self.ctx.graph = self.graph
-        telemetry_instance.set_context(server=self.graph)
-
-        with set_graph_context(self.graph):
-            with _add_init_error_context("configure reporters"):
-                self._configure_reporting(report_to)
-
-            with _add_init_error_context(
-                f"find a registered source for type {self.source_type}"
-            ):
-                source_class = source_registry.get(self.source_type)
-
-            with _add_init_error_context(f"configure the source ({self.source_type})"):
-                self.source = source_class.create(
-                    self.config.source.dict().get("config", {}), self.ctx
-                )
-                logger.debug(
-                    f"Source type {self.source_type} ({source_class}) configured"
+        with contextlib.ExitStack() as exit_stack, contextlib.ExitStack() as inner_exit_stack:
+            self.graph: Optional[DataHubGraph] = None
+            with _add_init_error_context("connect to DataHub"):
+                if self.config.datahub_api:
+                    self.graph = exit_stack.enter_context(
+                        DataHubGraph(self.config.datahub_api)
+                    )
+                    self.graph.test_connection()
+
+            with _add_init_error_context("set up framework context"):
+                self.ctx = PipelineContext(
+                    run_id=self.config.run_id,
+                    graph=self.graph,
+                    pipeline_name=self.config.pipeline_name,
+                    dry_run=dry_run,
+                    preview_mode=preview_mode,
+                    pipeline_config=self.config,
                 )
-                logger.info("Source configured successfully.")
 
-            extractor_type = self.config.source.extractor
-            with _add_init_error_context(f"configure the extractor ({extractor_type})"):
-                extractor_class = extractor_registry.get(extractor_type)
-                self.extractor = extractor_class(
-                    self.config.source.extractor_config, self.ctx
+            if self.config.sink is None:
+                logger.info(
+                    "No sink configured, attempting to use the default datahub-rest sink."
                 )
+                with _add_init_error_context("configure the default rest sink"):
+                    self.sink_type = "datahub-rest"
+                    self.sink = exit_stack.enter_context(
+                        _make_default_rest_sink(self.ctx)
+                    )
+            else:
+                self.sink_type = self.config.sink.type
+                with _add_init_error_context(
+                    f"find a registered sink for type {self.sink_type}"
+                ):
+                    sink_class = sink_registry.get(self.sink_type)
+
+                with _add_init_error_context(f"configure the sink ({self.sink_type})"):
+                    sink_config = self.config.sink.dict().get("config") or {}
+                    self.sink = exit_stack.enter_context(
+                        sink_class.create(sink_config, self.ctx)
+                    )
+                    logger.debug(
+                        f"Sink type {self.sink_type} ({sink_class}) configured"
+                    )
+            logger.info(f"Sink configured successfully. {self.sink.configured()}")
+
+            if self.graph is None and isinstance(self.sink, DatahubRestSink):
+                with _add_init_error_context("setup default datahub client"):
+                    self.graph = self.sink.emitter.to_graph()
+                    self.graph.test_connection()
+            self.ctx.graph = self.graph
+            telemetry_instance.set_context(server=self.graph)
+
+            with set_graph_context(self.graph):
+                with _add_init_error_context("configure reporters"):
+                    self._configure_reporting(report_to)
 
-            with _add_init_error_context("configure transformers"):
-                self._configure_transforms()
+                with _add_init_error_context(
+                    f"find a registered source for type {self.source_type}"
+                ):
+                    source_class = source_registry.get(self.source_type)
+
+                with _add_init_error_context(
+                    f"configure the source ({self.source_type})"
+                ):
+                    self.source = inner_exit_stack.enter_context(
+                        source_class.create(
+                            self.config.source.dict().get("config", {}), self.ctx
+                        )
+                    )
+                    logger.debug(
+                        f"Source type {self.source_type} ({source_class}) configured"
+                    )
+                    logger.info("Source configured successfully.")
+
+                extractor_type = self.config.source.extractor
+                with _add_init_error_context(
+                    f"configure the extractor ({extractor_type})"
+                ):
+                    extractor_class = extractor_registry.get(extractor_type)
+                    self.extractor = inner_exit_stack.enter_context(
+                        extractor_class(self.config.source.extractor_config, self.ctx)
+                    )
+
+                with _add_init_error_context("configure transformers"):
+                    self._configure_transforms()
+
+            # If all of the initialization succeeds, we can preserve the exit stack until the pipeline run.
+            # We need to use an exit stack so that if we have an exception during initialization,
+            # things that were already initialized are still cleaned up.
+            # We need to separate the source/extractor from the rest because stateful
+            # ingestion requires the source to be closed before the state can be updated.
+            self.inner_exit_stack = inner_exit_stack.pop_all()
+            self.exit_stack = exit_stack.pop_all()
 
     @property
     def source_type(self) -> str:
@@ -439,18 +400,20 @@ class Pipeline:
             return True
         return False
 
-    def run(self) -> None:  # noqa: C901
-        with contextlib.ExitStack() as stack:
+    def run(self) -> None:
+        with self.exit_stack, self.inner_exit_stack:
             if self.config.flags.generate_memory_profiles:
                 import memray
 
-                stack.enter_context(
+                self.exit_stack.enter_context(
                     memray.Tracker(
                         f"{self.config.flags.generate_memory_profiles}/{self.config.run_id}.bin"
                     )
                 )
 
-            stack.enter_context(self.sink)
+            self.exit_stack.enter_context(
+                change_default_attribution(KnownAttribution.INGESTION)
+            )
 
             self.final_status = PipelineStatus.UNKNOWN
             self._notify_reporters_on_ingestion_start()
@@ -459,8 +422,10 @@ class Pipeline:
                 callback = (
                     LoggingCallback()
                     if not self.config.failure_log.enabled
-                    else DeadLetterQueueCallback(
-                        self.ctx, self.config.failure_log.log_config
+                    else self.exit_stack.enter_context(
+                        DeadLetterQueueCallback(
+                            self.ctx, self.config.failure_log.log_config
+                        )
                     )
                 )
                 for wu in itertools.islice(
@@ -506,12 +471,11 @@ class Pipeline:
                             "Failed to process some records. Continuing.",
                             exc_info=e,
                         )
-                        # TODO: Transformer errors should cause the pipeline to fail.
+                        # TODO: Transformer errors should be reported more loudly / as part of the pipeline report.
 
                     if not self.dry_run:
                         self.sink.handle_work_unit_end(wu)
-                self.extractor.close()
-                self.source.close()
+
                 # no more data is coming, we need to let the transformers produce any additional records if they are holding on to state
                 for record_envelope in self.transform(
                     [
@@ -527,6 +491,11 @@ class Pipeline:
                         # TODO: propagate EndOfStream and other control events to sinks, to allow them to flush etc.
                         self.sink.write_record_async(record_envelope, callback)
 
+                # Stateful ingestion generates the updated state objects as part of the
+                # source's close method. Because of that, we need to close the source
+                # before we call process_commits.
+                self.inner_exit_stack.close()
+
                 self.process_commits()
                 self.final_status = PipelineStatus.COMPLETED
             except (SystemExit, KeyboardInterrupt) as e:
@@ -539,9 +508,6 @@ class Pipeline:
             finally:
                 clear_global_warnings()
 
-                if callback and hasattr(callback, "close"):
-                    callback.close()  # type: ignore
-
                 self._notify_reporters_on_ingestion_completion()
 
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:
@@ -560,10 +526,8 @@ class Pipeline:
         Evaluates the commit_policy for each committable in the context and triggers the commit operation
         on the committable if its required commit policies are satisfied.
         """
-        has_errors: bool = (
-            True
-            if self.source.get_report().failures or self.sink.get_report().failures
-            else False
+        has_errors: bool = bool(
+            self.source.get_report().failures or self.sink.get_report().failures
         )
         has_warnings: bool = bool(
             self.source.get_report().warnings or self.sink.get_report().warnings

datahub/ingestion/run/sink_callback.py

@@ -0,0 +1,77 @@
+import logging
+import threading
+from typing import Optional
+
+from datahub.ingestion.api.closeable import Closeable
+from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
+from datahub.ingestion.api.sink import WriteCallback
+from datahub.ingestion.sink.file import FileSink, FileSinkConfig
+from datahub.metadata.schema_classes import MetadataChangeProposalClass
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingCallback(WriteCallback):
+    def __init__(self, name: str = "") -> None:
+        super().__init__()
+        self.name = name
+
+    def on_success(
+        self, record_envelope: RecordEnvelope, success_metadata: dict
+    ) -> None:
+        logger.debug(
+            f"{self.name} sink wrote workunit {record_envelope.metadata['workunit_id']}"
+        )
+
+    def on_failure(
+        self,
+        record_envelope: RecordEnvelope,
+        failure_exception: Exception,
+        failure_metadata: dict,
+    ) -> None:
+        logger.error(
+            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}",
+            extra={"failure_metadata": failure_metadata},
+            exc_info=failure_exception,
+        )
+
+
+class DeadLetterQueueCallback(WriteCallback, Closeable):
+    def __init__(self, ctx: PipelineContext, config: Optional[FileSinkConfig]) -> None:
+        if not config:
+            config = FileSinkConfig.parse_obj({"filename": "failed_events.json"})
+        self.file_sink: FileSink = FileSink(ctx, config)
+        self.file_sink_lock = threading.Lock()
+        self.logging_callback = LoggingCallback(name="failure-queue")
+        logger.info(f"Failure logging enabled. Will log to {config.filename}.")
+
+    def on_success(
+        self, record_envelope: RecordEnvelope, success_metadata: dict
+    ) -> None:
+        pass
+
+    def on_failure(
+        self,
+        record_envelope: RecordEnvelope,
+        failure_exception: Exception,
+        failure_metadata: dict,
+    ) -> None:
+        if "workunit_id" in record_envelope.metadata and isinstance(
+            record_envelope.record, MetadataChangeProposalClass
+        ):
+            mcp: MetadataChangeProposalClass = record_envelope.record
+            if mcp.systemMetadata:
+                if not mcp.systemMetadata.properties:
+                    mcp.systemMetadata.properties = {}
+                if "workunit_id" not in mcp.systemMetadata.properties:
+                    # update the workunit id
+                    mcp.systemMetadata.properties["workunit_id"] = (
+                        record_envelope.metadata["workunit_id"]
+                    )
+            record_envelope.record = mcp
+        with self.file_sink_lock:
+            self.file_sink.write_record_async(record_envelope, self.logging_callback)
+
+    def close(self) -> None:
+        with self.file_sink_lock:
+            self.file_sink.close()

datahub/ingestion/sink/datahub_rest.py

@@ -20,7 +20,11 @@ from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
+    DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_TRACE_MODE,
     DataHubRestEmitter,
+    RestSinkEndpoint,
+    RestTraceMode,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -66,6 +70,8 @@ _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
 
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
+    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
+    default_trace_mode: RestTraceMode = DEFAULT_REST_TRACE_MODE
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -172,6 +178,8 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             ca_certificate_path=config.ca_certificate_path,
             client_certificate_path=config.client_certificate_path,
             disable_ssl_verification=config.disable_ssl_verification,
+            openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
+            default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
         )
 
     @property

datahub/ingestion/source/abs/config.py

@@ -144,10 +144,8 @@ class DataLakeSourceConfig(
         return path_specs
 
     @pydantic.validator("platform", always=True)
-    def platform_not_empty(cls, platform: str, values: dict) -> str:
-        inferred_platform = values.get(
-            "platform", None
-        )  # we may have inferred it above
+    def platform_not_empty(cls, platform: Any, values: dict) -> str:
+        inferred_platform = values.get("platform")  # we may have inferred it above
         platform = platform or inferred_platform
         if not platform:
             raise ValueError("platform must not be empty")

datahub/ingestion/source/bigquery_v2/bigquery_audit.py

@@ -165,7 +165,7 @@ class BigQueryTableRef:
     @classmethod
     def from_spec_obj(cls, spec: dict) -> "BigQueryTableRef":
         for key in ["projectId", "datasetId", "tableId"]:
-            if key not in spec.keys():
+            if key not in spec:
                 raise ValueError(f"invalid BigQuery table reference dict: {spec}")
 
         return cls(

datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -1,8 +1,6 @@
-import json
 import logging
 import os
 import re
-import tempfile
 from datetime import timedelta
 from typing import Any, Dict, List, Optional, Union
 
@@ -17,10 +15,10 @@ from datahub.configuration.source_common import (
     PlatformInstanceConfigMixin,
 )
 from datahub.configuration.validate_field_removal import pydantic_removed_field
-from datahub.configuration.validate_multiline_string import pydantic_multiline_string
 from datahub.ingestion.glossary.classification_mixin import (
     ClassificationSourceConfigMixin,
 )
+from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
 from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterConfig
 from datahub.ingestion.source.state.stateful_ingestion_base import (
@@ -107,50 +105,8 @@ class BigQueryUsageConfig(BaseUsageConfig):
     )
 
 
-class BigQueryCredential(ConfigModel):
-    project_id: str = Field(description="Project id to set the credentials")
-    private_key_id: str = Field(description="Private key id")
-    private_key: str = Field(
-        description="Private key in a form of '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n'"
-    )
-    client_email: str = Field(description="Client email")
-    client_id: str = Field(description="Client Id")
-    auth_uri: str = Field(
-        default="https://accounts.google.com/o/oauth2/auth",
-        description="Authentication uri",
-    )
-    token_uri: str = Field(
-        default="https://oauth2.googleapis.com/token", description="Token uri"
-    )
-    auth_provider_x509_cert_url: str = Field(
-        default="https://www.googleapis.com/oauth2/v1/certs",
-        description="Auth provider x509 certificate url",
-    )
-    type: str = Field(default="service_account", description="Authentication type")
-    client_x509_cert_url: Optional[str] = Field(
-        default=None,
-        description="If not set it will be default to https://www.googleapis.com/robot/v1/metadata/x509/client_email",
-    )
-
-    _fix_private_key_newlines = pydantic_multiline_string("private_key")
-
-    @root_validator(skip_on_failure=True)
-    def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("client_x509_cert_url") is None:
-            values["client_x509_cert_url"] = (
-                f"https://www.googleapis.com/robot/v1/metadata/x509/{values['client_email']}"
-            )
-        return values
-
-    def create_credential_temp_file(self) -> str:
-        with tempfile.NamedTemporaryFile(delete=False) as fp:
-            cred_json = json.dumps(self.dict(), indent=4, separators=(",", ": "))
-            fp.write(cred_json.encode())
-            return fp.name
-
-
 class BigQueryConnectionConfig(ConfigModel):
-    credential: Optional[BigQueryCredential] = Field(
+    credential: Optional[GCPCredential] = Field(
         default=None, description="BigQuery credential informations"
     )
 

datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -292,6 +292,11 @@ class BigQuerySchemaApi:
                         if hasattr(d, "_properties") and isinstance(d._properties, dict)
                         else None
                     ),
+                    # TODO: Fetch dataset description individually impacts overall performance if the number of datasets is high (hundreds); instead we should fetch in batch for all datasets.
+                    # TODO: Given we are calling get_dataset for each dataset, we may consume and publish other fields too, such as created, modified, etc...
+                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_get_dataset
+                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dataset.Dataset
+                    comment=self.bq_client.get_dataset(d.reference).description,
                 )
                 for d in datasets
             ]
@@ -339,7 +344,7 @@ class BigQuerySchemaApi:
         with_partitions: bool = False,
     ) -> Iterator[BigqueryTable]:
         with PerfTimer() as current_timer:
-            filter_clause: str = ", ".join(f"'{table}'" for table in tables.keys())
+            filter_clause: str = ", ".join(f"'{table}'" for table in tables)
 
             if with_partitions:
                 query_template = BigqueryQuery.tables_for_dataset

datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -296,6 +296,7 @@ class BigQuerySchemaGenerator:
         self,
         dataset: str,
         project_id: str,
+        description: Optional[str] = None,
         tags: Optional[Dict[str, str]] = None,
         extra_properties: Optional[Dict[str, str]] = None,
     ) -> Iterable[MetadataWorkUnit]:
@@ -336,6 +337,7 @@ class BigQuerySchemaGenerator:
             domain_config=self.config.domain,
             schema_container_key=schema_container_key,
             database_container_key=database_container_key,
+            description=description,
             external_url=(
                 BQ_EXTERNAL_DATASET_URL_TEMPLATE.format(
                     project=project_id, dataset=dataset
@@ -471,14 +473,15 @@ class BigQuerySchemaGenerator:
 
         if self.config.include_schema_metadata:
             yield from self.gen_dataset_containers(
-                dataset_name,
-                project_id,
-                bigquery_dataset.labels,
-                (
+                dataset=dataset_name,
+                project_id=project_id,
+                tags=bigquery_dataset.labels,
+                extra_properties=(
                     {"location": bigquery_dataset.location}
                     if bigquery_dataset.location
                     else None
                 ),
+                description=bigquery_dataset.comment,
             )
 
         columns = None