import base64
import logging
import os
import time
from enum import Enum
from importlib.metadata import version
from typing import IO, Any, List, Mapping, Optional, Tuple, Union, cast

import dagster
import dagster._check as check
import dagster_pyspark
import databricks_api
import databricks_cli.sdk
import requests.exceptions
from dagster._annotations import deprecated, public
from databricks.sdk import WorkspaceClient
from databricks.sdk.core import (
    Config,
    DefaultCredentials,
    azure_service_principal,
    oauth_service_principal,
    pat_auth,
)
from databricks.sdk.service import compute, jobs
from typing_extensions import Final

import dagster_databricks

from .types import DatabricksRunState
from .version import __version__

# wait at most 24 hours by default for run execution
DEFAULT_RUN_MAX_WAIT_TIME_SEC: Final = 24 * 60 * 60


class DatabricksError(Exception):
    pass


class AuthTypeEnum(Enum):
    OAUTH_M2M = "oauth-m2m"
    PAT = "pat"
    AZURE_CLIENT_SECRET = "azure-client-secret"
    DEFAULT = "default"


class WorkspaceClientFactory:
    def __init__(
        self,
        host: Optional[str],
        token: Optional[str],
        oauth_client_id: Optional[str],
        oauth_client_secret: Optional[str],
        azure_client_id: Optional[str],
        azure_client_secret: Optional[str],
        azure_tenant_id: Optional[str],
    ):
        """Initialize the Databricks Workspace client. Users may provide explicit credentials for a PAT, databricks
        service principal oauth credentials, or azure service principal credentials. If no credentials are provided,
        the underlying WorkspaceClient from `databricks.sdk` will attempt to read credentials from the environment or
        from the `~/.databrickscfg` file. For more information, see the Databricks SDK docs on various ways you can
        authenticate with the WorkspaceClient, through which most interactions with the Databricks API occur.
        <https://docs.databricks.com/en/dev-tools/auth.html#authentication-for-databricks-automation>`_.
        """
        self._raise_if_multiple_auth_types(
            token=token,
            oauth_client_id=oauth_client_id,
            oauth_client_secret=oauth_client_secret,
            azure_client_id=azure_client_id,
            azure_client_secret=azure_client_secret,
            azure_tenant_id=azure_tenant_id,
        )
        self._assert_valid_credentials_combos(
            oauth_client_id=oauth_client_id,
            oauth_client_secret=oauth_client_secret,
            azure_client_id=azure_client_id,
            azure_client_secret=azure_client_secret,
            azure_tenant_id=azure_tenant_id,
        )
        auth_type = self._get_auth_type(
            token,
            oauth_client_id,
            oauth_client_secret,
            azure_client_id,
            azure_client_secret,
            azure_tenant_id,
        )
        product_info = {"product": "dagster-databricks", "product_version": __version__}

        # Figure out what credentials provider to use based on any explicitly-provided credentials. If none were
        # provided, then fallback to the default credentials provider, which will attempt to read credentials from
        # the environment or from a `~/.databrickscfg` file, if it exists.

        if auth_type == AuthTypeEnum.OAUTH_M2M:
            host = self._resolve_host(host)
            c = Config(
                host=host,
                client_id=oauth_client_id,
                client_secret=oauth_client_secret,
                credentials_provider=oauth_service_principal,
                **product_info,
            )
        elif auth_type == AuthTypeEnum.PAT:
            host = self._resolve_host(host)
            c = Config(host=host, token=token, credentials_provider=pat_auth, **product_info)
        elif auth_type == AuthTypeEnum.AZURE_CLIENT_SECRET:
            host = self._resolve_host(host)
            c = Config(
                host=host,
                azure_client_id=azure_client_id,
                azure_client_secret=azure_client_secret,
                azure_tenant_id=azure_tenant_id,
                credentials_provider=azure_service_principal,
                **product_info,
            )
        elif auth_type == AuthTypeEnum.DEFAULT:
            # Can be used to automatically read credentials from environment or ~/.databrickscfg file. This is common
            # when launching Databricks jobs from a laptop development setting through Dagster
            if host is not None:
                # This allows for explicit override of the host, while letting other credentials be read from the
                # environment or ~/.databrickscfg file
                c = Config(host=host, credentials_provider=DefaultCredentials(), **product_info)
            else:
                # The initialization machinery in the Config object will look for the host and other auth info in the
                # environment, as long as no values are provided for those attributes (including None)
                c = Config(credentials_provider=DefaultCredentials(), **product_info)
        else:
            raise ValueError(f"Unexpected auth type {auth_type}")
        self.config = c

    def _raise_if_multiple_auth_types(
        self,
        token: Optional[str] = None,
        oauth_client_id: Optional[str] = None,
        oauth_client_secret: Optional[str] = None,
        azure_client_id: Optional[str] = None,
        azure_client_secret: Optional[str] = None,
        azure_tenant_id: Optional[str] = None,
    ):
        more_than_one_auth_type_provided = (
            sum(
                [
                    True
                    for _ in [
                        token,
                        (oauth_client_id and oauth_client_secret),
                        (azure_client_id and azure_client_secret and azure_tenant_id),
                    ]
                    if _
                ]
            )
            > 1
        )
        if more_than_one_auth_type_provided:
            raise ValueError(
                "Can only provide one of token, oauth credentials, or azure credentials"
            )

    @staticmethod
    def _get_auth_type(
        token: Optional[str],
        oauth_client_id: Optional[str],
        oauth_client_secret: Optional[str],
        azure_client_id: Optional[str],
        azure_client_secret: Optional[str],
        azure_tenant_id: Optional[str],
    ) -> AuthTypeEnum:
        """Get the type of authentication used to initialize the WorkspaceClient."""
        if oauth_client_id and oauth_client_secret:
            auth_type = AuthTypeEnum.OAUTH_M2M
        elif token:
            auth_type = AuthTypeEnum.PAT
        elif azure_client_id and azure_client_secret and azure_tenant_id:
            auth_type = AuthTypeEnum.AZURE_CLIENT_SECRET
        else:
            auth_type = AuthTypeEnum.DEFAULT
        return auth_type

    @staticmethod
    def _assert_valid_credentials_combos(
        oauth_client_id: Optional[str] = None,
        oauth_client_secret: Optional[str] = None,
        azure_client_id: Optional[str] = None,
        azure_client_secret: Optional[str] = None,
        azure_tenant_id: Optional[str] = None,
    ):
        """Ensure that all required credentials are provided for the given auth type."""
        if (
            oauth_client_id
            and not oauth_client_secret
            or oauth_client_secret
            and not oauth_client_id
        ):
            raise ValueError(
                "If using databricks service principal oauth credentials, both oauth_client_id and"
                " oauth_client_secret must be provided"
            )
        if (
            (azure_client_id and not azure_client_secret and not azure_tenant_id)
            or (azure_client_secret and not azure_client_id and not azure_tenant_id)
            or (azure_tenant_id and not azure_client_id and not azure_client_secret)
        ):
            raise ValueError(
                "If using azure service principal auth, azure_client_id, azure_client_secret, and"
                " azure_tenant_id must be provided"
            )

    def get_workspace_client(self) -> WorkspaceClient:
        return WorkspaceClient(config=self.config)

    @staticmethod
    def _resolve_host(host: Optional[str]) -> str:
        host = host if host else os.getenv("DATABRICKS_HOST")
        if host is None:
            raise ValueError(
                "Must provide host explicitly or in DATABRICKS_HOST env var when providing"
                " credentials explicitly"
            )
        return host


class DatabricksClient:
    """A thin wrapper over the Databricks REST API."""

    def __init__(
        self,
        host: Optional[str] = None,
        token: Optional[str] = None,
        oauth_client_id: Optional[str] = None,
        oauth_client_secret: Optional[str] = None,
        azure_client_id: Optional[str] = None,
        azure_client_secret: Optional[str] = None,
        azure_tenant_id: Optional[str] = None,
        workspace_id: Optional[str] = None,
    ):
        self.host = host
        self.workspace_id = workspace_id

        workspace_client_factory = WorkspaceClientFactory(
            oauth_client_id=oauth_client_id,
            oauth_client_secret=oauth_client_secret,
            azure_client_id=azure_client_id,
            azure_client_secret=azure_client_secret,
            azure_tenant_id=azure_tenant_id,
            token=token,
            host=host,
        )
        self._workspace_client = workspace_client_factory.get_workspace_client()

        # TODO: This is the old shim client that we were previously using. Arguably this is
        # confusing for users to use since this is an unofficial wrapper around the documented
        # Databricks REST API. We should consider removing this in the next minor release.
        if token:
            self._client = databricks_api.DatabricksAPI(host=host, token=token)
            self.__setup_user_agent(self._client.client)
            # TODO: This is the old `databricks_cli` client that was previously recommended by Databricks.
            # It is no longer supported and should be removed in favour of `databricks-sdk` in the next
            # minor release.
            self._api_client = databricks_cli.sdk.ApiClient(host=host, token=token)
            self.__setup_user_agent(self._api_client)
        else:
            self._client = None
            self._api_client = None

    def __setup_user_agent(
        self,
        client: Union[WorkspaceClient, databricks_api.DatabricksAPI, databricks_cli.sdk.ApiClient],
    ) -> None:
        """Overrides the user agent for the Databricks API client."""
        client.default_headers["user-agent"] = f"dagster-databricks/{__version__}"

    @deprecated(
        breaking_version="0.21.0", additional_warn_text="Use `workspace_client` property instead."
    )
    @public
    @property
    def client(self) -> databricks_api.DatabricksAPI:
        """Retrieve the legacy Databricks API client. Note: accessing this property will throw an exception if oauth
        credentials are used to initialize the DatabricksClient, because oauth credentials are not supported by the
        legacy Databricks API client.
        """
        if self._client is None:
            raise ValueError(
                "Legacy Databricks API client from `databricks-api` was not initialized because"
                " oauth credentials were used instead of an access token. This legacy Databricks"
                " API client is not supported when using oauth credentials. Use the"
                " `workspace_client` property instead."
            )
        return self._client

    @client.setter
    def client(self, value: Optional[databricks_api.DatabricksAPI]) -> None:
        self._client = value

    @deprecated(
        breaking_version="0.21.0", additional_warn_text="Use `workspace_client` property instead."
    )
    @public
    @property
    def api_client(self) -> databricks_cli.sdk.ApiClient:
        """Retrieve a reference to the underlying Databricks API client. For more information,
        see the `Databricks Python API <https://docs.databricks.com/dev-tools/python-api.html>`_.
        Noe: accessing this property will throw an exception if oauth credentials are used to initialize the
        DatabricksClient, because oauth credentials are not supported by the legacy Databricks API client.
        **Examples:**.

        .. code-block:: python

            from dagster import op
            from databricks_cli.jobs.api import JobsApi
            from databricks_cli.runs.api import RunsApi
            from databricks.sdk import WorkspaceClient

            @op(required_resource_keys={"databricks_client"})
            def op1(context):
                # Initialize the Databricks Jobs API
                jobs_client = JobsApi(context.resources.databricks_client.api_client)
                runs_client = RunsApi(context.resources.databricks_client.api_client)
                client = context.resources.databricks_client.api_client

                # Example 1: Run a Databricks job with some parameters.
                jobs_client.run_now(...)
                client.jobs.run_now(...)

                # Example 2: Trigger a one-time run of a Databricks workload.
                runs_client.submit_run(...)
                client.jobs.submit(...)

                # Example 3: Get an existing run.
                runs_client.get_run(...)
                client.jobs.get_run(...)

                # Example 4: Cancel a run.
                runs_client.cancel_run(...)
                client.jobs.cancel_run(...)

        Returns:
            ApiClient: The authenticated Databricks API client.
        """
        if self._api_client is None:
            raise ValueError(
                "Legacy Databricks API client from `databricks-cli` was not initialized because"
                " oauth credentials were used instead of an access token. This legacy Databricks"
                " API client is not supported when using oauth credentials. Use the"
                " `workspace_client` property instead."
            )
        return self._api_client

    @public
    @property
    def workspace_client(self) -> WorkspaceClient:
        """Retrieve a reference to the underlying Databricks Workspace client. For more information,
        see the `Databricks SDK for Python <https://docs.databricks.com/dev-tools/sdk-python.html>`_.

        **Examples:**

        .. code-block:: python

            from dagster import op
            from databricks.sdk import WorkspaceClient

            @op(required_resource_keys={"databricks_client"})
            def op1(context):
                # Initialize the Databricks Jobs API
                client = context.resources.databricks_client.api_client

                # Example 1: Run a Databricks job with some parameters.
                client.jobs.run_now(...)

                # Example 2: Trigger a one-time run of a Databricks workload.
                client.jobs.submit(...)

                # Example 3: Get an existing run.
                client.jobs.get_run(...)

                # Example 4: Cancel a run.
                client.jobs.cancel_run(...)

        Returns:
            WorkspaceClient: The authenticated Databricks SDK Workspace Client.
        """
        return self._workspace_client

    def read_file(self, dbfs_path: str, block_size: int = 1024**2) -> bytes:
        """Read a file from DBFS to a **byte string**."""
        if dbfs_path.startswith("dbfs://"):
            dbfs_path = dbfs_path[7:]

        data = b""
        bytes_read = 0
        dbfs_service = self.workspace_client.dbfs

        jdoc = dbfs_service.read(path=dbfs_path, length=block_size)
        data += base64.b64decode(jdoc.data)
        while jdoc.bytes_read == block_size:
            bytes_read += jdoc.bytes_read
            jdoc = dbfs_service.read(path=dbfs_path, offset=bytes_read, length=block_size)
            data += base64.b64decode(jdoc.data)

        return data

    def put_file(
        self, file_obj: IO, dbfs_path: str, overwrite: bool = False, block_size: int = 1024**2
    ) -> None:
        """Upload an arbitrary large file to DBFS.

        This doesn't use the DBFS `Put` API because that endpoint is limited to 1MB.
        """
        if dbfs_path.startswith("dbfs://"):
            dbfs_path = dbfs_path[7:]

        dbfs_service = self.workspace_client.dbfs

        create_response = dbfs_service.create(path=dbfs_path, overwrite=overwrite)
        handle = create_response.handle

        block = file_obj.read(block_size)
        while block:
            data = base64.b64encode(block).decode("utf-8")
            dbfs_service.add_block(data=data, handle=handle)
            block = file_obj.read(block_size)

        dbfs_service.close(handle=handle)

    def get_run_state(self, databricks_run_id: int) -> "DatabricksRunState":
        """Get the state of a run by Databricks run ID.

        Return a `DatabricksRunState` object. Note that the `result_state`
        attribute may be `None` if the run hasn't yet terminated.
        """
        run = self.workspace_client.jobs.get_run(databricks_run_id)
        return DatabricksRunState.from_databricks(run.state)

    def poll_run_state(
        self,
        logger: logging.Logger,
        start_poll_time: float,
        databricks_run_id: int,
        max_wait_time_sec: float,
        verbose_logs: bool = True,
    ) -> bool:
        run_state = self.get_run_state(databricks_run_id)

        if run_state.has_terminated():
            if run_state.is_successful():
                logger.info(f"Run `{databricks_run_id}` completed successfully.")
                return True
            if run_state.is_skipped():
                logger.info(f"Run `{databricks_run_id}` was skipped.")
                return True
            else:
                error_message = (
                    f"Run `{databricks_run_id}` failed with result state:"
                    f" `{run_state.result_state}`. Message: {run_state.state_message}."
                )
                logger.error(error_message)
                raise DatabricksError(error_message)
        else:
            if verbose_logs:
                logger.debug(f"Run `{databricks_run_id}` in state {run_state}.")
        if time.time() - start_poll_time > max_wait_time_sec:
            raise DatabricksError(
                f"Run `{databricks_run_id}` took more than {max_wait_time_sec}s to complete."
                " Failing the run."
            )
        return False

    def wait_for_run_to_complete(
        self,
        logger: logging.Logger,
        databricks_run_id: int,
        poll_interval_sec: float,
        max_wait_time_sec: int,
        verbose_logs: bool = True,
    ) -> None:
        logger.info(f"Waiting for Databricks run `{databricks_run_id}` to complete...")

        start_poll_time = time.time()
        while True:
            if self.poll_run_state(
                logger=logger,
                start_poll_time=start_poll_time,
                databricks_run_id=databricks_run_id,
                max_wait_time_sec=max_wait_time_sec,
                verbose_logs=verbose_logs,
            ):
                return

            time.sleep(poll_interval_sec)


class DatabricksJobRunner:
    """Submits jobs created using Dagster config to Databricks, and monitors their progress.

    Attributes:
        host (str): Databricks host, e.g. https://uksouth.azuredatabricks.net.
        token (str): Databricks authentication token.
        poll_interval_sec (float): How often to poll Databricks for run status.
        max_wait_time_sec (int): How long to wait for a run to complete before failing.
    """

    def __init__(
        self,
        host: Optional[str] = None,
        token: Optional[str] = None,
        oauth_client_id: Optional[str] = None,
        oauth_client_secret: Optional[str] = None,
        azure_client_id: Optional[str] = None,
        azure_client_secret: Optional[str] = None,
        azure_tenant_id: Optional[str] = None,
        poll_interval_sec: float = 5,
        max_wait_time_sec: int = DEFAULT_RUN_MAX_WAIT_TIME_SEC,
    ):
        self.host = check.opt_str_param(host, "host")
        self.token = check.opt_str_param(token, "token")
        self.poll_interval_sec = check.numeric_param(poll_interval_sec, "poll_interval_sec")
        self.max_wait_time_sec = check.int_param(max_wait_time_sec, "max_wait_time_sec")

        oauth_client_id = check.opt_str_param(oauth_client_id, "oauth_client_id")
        oauth_client_secret = check.opt_str_param(oauth_client_secret, "oauth_client_secret")
        azure_client_id = check.opt_str_param(azure_client_id, "azure_client_id")
        azure_client_secret = check.opt_str_param(azure_client_secret, "azure_client_secret")
        azure_tenant_id = check.opt_str_param(azure_tenant_id, "azure_tenant_id")

        self._client: DatabricksClient = DatabricksClient(
            host=self.host,
            token=self.token,
            oauth_client_id=oauth_client_id,
            oauth_client_secret=oauth_client_secret,
            azure_client_id=azure_client_id,
            azure_client_secret=azure_client_secret,
            azure_tenant_id=azure_tenant_id,
        )

    @property
    def client(self) -> DatabricksClient:
        """Return the underlying `DatabricksClient` object."""
        return self._client

    def submit_run(self, run_config: Mapping[str, Any], task: Mapping[str, Any]) -> int:
        """Submit a new run using the 'Runs submit' API."""
        existing_cluster_id = run_config["cluster"].get("existing")

        new_cluster = run_config["cluster"].get("new")

        # The Databricks API needs different keys to be present in API calls depending
        # on new/existing cluster, so we need to process the new_cluster
        # config first.
        if new_cluster:
            new_cluster = new_cluster.copy()

            nodes = new_cluster.pop("nodes")

            if "driver_instance_pool_id" in nodes:
                check.invariant(
                    "instance_pool_id" in nodes,
                    "Usage of `driver_instance_pool_id` requires that `instance_pool_id` be specified"
                    " for worker nodes",
                )

            if "instance_pool_id" in nodes:
                new_cluster["instance_pool_id"] = nodes["instance_pool_id"]

                if "driver_instance_pool_id" in nodes:
                    new_cluster["driver_instance_pool_id"] = nodes["driver_instance_pool_id"]
            else:
                node_types = nodes["node_types"]
                new_cluster["node_type_id"] = node_types["node_type_id"]
                if "driver_node_type_id" in node_types:
                    new_cluster["driver_node_type_id"] = node_types["driver_node_type_id"]

            cluster_size = new_cluster.pop("size")
            if "num_workers" in cluster_size:
                new_cluster["num_workers"] = cluster_size["num_workers"]
            else:
                new_cluster["autoscale"] = cluster_size["autoscale"]

            tags = new_cluster.get("custom_tags", {})
            if isinstance(tags, list):
                tags = {x["key"]: x["value"] for x in tags}
            tags["__dagster_version"] = dagster.__version__
            new_cluster["custom_tags"] = tags

        check.invariant(
            existing_cluster_id is not None or new_cluster is not None,
            "Invalid value for run_config.cluster",
        )

        # We'll always need some libraries, namely dagster/dagster_databricks/dagster_pyspark,
        # since they're imported by our scripts.
        # Add them if they're not already added by users in config.
        libraries = list(run_config.get("libraries", []))
        install_default_libraries = run_config.get("install_default_libraries", True)
        if install_default_libraries:
            python_libraries = {
                x["pypi"]["package"].split("==")[0].replace("_", "-")
                for x in libraries
                if "pypi" in x
            }

            for library_name, library in [
                ("dagster", dagster),
                ("dagster-databricks", dagster_databricks),
                ("dagster-pyspark", dagster_pyspark),
            ]:
                if library_name not in python_libraries:
                    libraries.append(
                        {"pypi": {"package": f"{library_name}=={library.__version__}"}}
                    )

            if "databricks-sdk" not in python_libraries:
                libraries.append(
                    {"pypi": {"package": f"databricks-sdk=={version('databricks-sdk')}"}}
                )

        # Only one task should be able to be chosen really; make sure of that here.
        check.invariant(
            sum(
                task.get(key) is not None
                for key in [
                    "notebook_task",
                    "spark_python_task",
                    "spark_jar_task",
                    "spark_submit_task",
                ]
            )
            == 1,
            "Multiple tasks specified in Databricks run",
        )

        notification_settings = self._get_notification_settings(run_config)
        job_health_settings = self._get_job_health_settings(run_config)

        return self.client.workspace_client.jobs.submit(
            run_name=run_config.get("run_name"),
            tasks=[
                jobs.SubmitTask.from_dict(
                    {
                        "new_cluster": new_cluster,
                        "existing_cluster_id": existing_cluster_id,
                        "libraries": libraries,
                        **task,
                        "task_key": "dagster-task",
                    },
                )
            ],
            idempotency_token=run_config.get("idempotency_token"),
            timeout_seconds=run_config.get("timeout_seconds"),
            health=job_health_settings,
            **notification_settings,
        ).bind()["run_id"]

    def _get_job_health_settings(
        self, run_config: Mapping[str, Any]
    ) -> Optional[List[jobs.JobsHealthRule]]:
        if "job_health_settings" in run_config:
            job_health_settings = [
                jobs.JobsHealthRule.from_dict(h) for h in run_config["job_health_settings"]
            ]
        else:
            job_health_settings = None
        return job_health_settings

    def _get_notification_settings(
        self, run_config: Mapping[str, Any]
    ) -> Mapping[
        str,
        Union[jobs.JobEmailNotifications, jobs.JobNotificationSettings, jobs.WebhookNotifications],
    ]:
        notification_configs = {}
        if "email_notifications" in run_config:
            email_notifications = jobs.JobEmailNotifications.from_dict(
                run_config["email_notifications"]
            )
            notification_configs.update({"email_notifications": email_notifications})
        if "notification_settings" in run_config:
            notification_settings = jobs.JobNotificationSettings.from_dict(
                run_config["notification_settings"]
            )
            notification_configs.update({"notification_settings": notification_settings})
        if "webhook_notifications" in run_config:
            webhook_notifications = jobs.WebhookNotifications.from_dict(
                run_config["webhook_notifications"]
            )
            notification_configs.update({"webhook_notifications": webhook_notifications})
        return notification_configs

    def retrieve_logs_for_run_id(
        self, log: logging.Logger, databricks_run_id: int
    ) -> Optional[Tuple[Optional[str], Optional[str]]]:
        """Retrieve the stdout and stderr logs for a run."""
        run = self.client.workspace_client.jobs.get_run(databricks_run_id)

        # Run.cluster_instance can be None. In that case, fall back to cluster instance on first
        # task. Currently pyspark step launcher runs jobs with singleton tasks.
        cluster_instance = run.cluster_instance or run.tasks[0].cluster_instance
        cluster_id = check.inst(
            cluster_instance.cluster_id,
            str,
            "cluster_id should be string like `1234-123456-abcdefgh` got:"
            f" `{cluster_instance.cluster_id}`",
        )
        cluster = self.client.workspace_client.clusters.get(cluster_id)
        log_config = cluster.cluster_log_conf
        if log_config is None:
            log.warn(
                f"Logs not configured for cluster {cluster_id} used for run {databricks_run_id}"
            )
            return None
        if cast(Optional[compute.S3StorageInfo], log_config.s3) is not None:
            logs_prefix = log_config.s3.destination
            log.warn("Retrieving S3 logs not yet implemented")
            return None
        elif cast(Optional[compute.DbfsStorageInfo], log_config.dbfs) is not None:
            logs_prefix = log_config.dbfs.destination
            stdout = self.wait_for_dbfs_logs(log, logs_prefix, cluster_id, "stdout")
            stderr = self.wait_for_dbfs_logs(log, logs_prefix, cluster_id, "stderr")
            return stdout, stderr

    def wait_for_dbfs_logs(
        self,
        log: logging.Logger,
        prefix: str,
        cluster_id: str,
        filename: str,
        waiter_delay: int = 10,
        waiter_max_attempts: int = 10,
    ) -> Optional[str]:
        """Attempt up to `waiter_max_attempts` attempts to get logs from DBFS."""
        path = "/".join([prefix, cluster_id, "driver", filename])
        log.info(f"Retrieving logs from {path}")
        num_attempts = 0
        while num_attempts <= waiter_max_attempts:
            try:
                logs = self.client.read_file(path)
                return logs.decode("utf-8")
            except requests.exceptions.HTTPError:
                num_attempts += 1
                time.sleep(waiter_delay)
        log.warn("Could not retrieve cluster logs!")
