star

venv/lib/python3.9/site-packages/databricks/sdk/runtime/__init__.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import logging
+from typing import Dict, Optional, Union, cast
+
+logger = logging.getLogger("databricks.sdk")
+is_local_implementation = True
+
+# All objects that are injected into the Notebook's user namespace should also be made
+# available to be imported from databricks.sdk.runtime.globals. This import can be used
+# in Python modules so users can access these objects from Files more easily.
+dbruntime_objects = [
+    "display",
+    "displayHTML",
+    "dbutils",
+    "table",
+    "sql",
+    "udf",
+    "getArgument",
+    "sc",
+    "sqlContext",
+    "spark",
+]
+
+# DO NOT MOVE THE TRY-CATCH BLOCK BELOW AND DO NOT ADD THINGS BEFORE IT! WILL MAKE TEST FAIL.
+try:
+    # We don't want to expose additional entity to user namespace, so
+    # a workaround here for exposing required information in notebook environment
+    from dbruntime.sdk_credential_provider import init_runtime_native_auth
+
+    logger.debug("runtime SDK credential provider available")
+    dbruntime_objects.append("init_runtime_native_auth")
+except ImportError:
+    init_runtime_native_auth = None
+
+globals()["init_runtime_native_auth"] = init_runtime_native_auth
+
+
+def init_runtime_repl_auth():
+    try:
+        from dbruntime.databricks_repl_context import get_context
+
+        ctx = get_context()
+        if ctx is None:
+            logger.debug("Empty REPL context returned, skipping runtime auth")
+            return None, None
+        if ctx.workspaceUrl is None:
+            logger.debug("Workspace URL is not available, skipping runtime auth")
+            return None, None
+        host = f"https://{ctx.workspaceUrl}"
+
+        def inner() -> Dict[str, str]:
+            ctx = get_context()
+            return {"Authorization": f"Bearer {ctx.apiToken}"}
+
+        return host, inner
+    except ImportError:
+        return None, None
+
+
+def init_runtime_legacy_auth():
+    try:
+        import IPython
+
+        ip_shell = IPython.get_ipython()
+        if ip_shell is None:
+            return None, None
+        global_ns = ip_shell.ns_table["user_global"]
+        if "dbutils" not in global_ns:
+            return None, None
+        dbutils = global_ns["dbutils"].notebook.entry_point.getDbutils()
+        if dbutils is None:
+            return None, None
+        ctx = dbutils.notebook().getContext()
+        if ctx is None:
+            return None, None
+        host = getattr(ctx, "apiUrl")().get()
+
+        def inner() -> Dict[str, str]:
+            ctx = dbutils.notebook().getContext()
+            return {"Authorization": f'Bearer {getattr(ctx, "apiToken")().get()}'}
+
+        return host, inner
+    except ImportError:
+        return None, None
+
+
+try:
+    # Internal implementation
+    # Separated from above for backward compatibility
+    from dbruntime import UserNamespaceInitializer
+
+    userNamespaceGlobals = UserNamespaceInitializer.getOrCreate().get_namespace_globals()
+    _globals = globals()
+    for var in dbruntime_objects:
+        if var not in userNamespaceGlobals:
+            continue
+        _globals[var] = userNamespaceGlobals[var]
+    is_local_implementation = False
+except ImportError:
+    # OSS implementation
+    is_local_implementation = True
+
+    for var in dbruntime_objects:
+        globals()[var] = None
+
+    # The next few try-except blocks are for initialising globals in a best effort
+    # mannaer. We separate them to try to get as many of them working as possible
+    try:
+        # We expect this to fail and only do this for providing types
+        from pyspark.sql.context import SQLContext
+
+        sqlContext: SQLContext = None  # type: ignore
+        table = sqlContext.table
+    except Exception as e:
+        logging.debug(f"Failed to initialize globals 'sqlContext' and 'table', continuing. Cause: {e}")
+
+    try:
+        from pyspark.sql.functions import udf  # type: ignore
+    except ImportError as e:
+        logging.debug(f"Failed to initialise udf global: {e}")
+
+    try:
+        from databricks.connect import DatabricksSession  # type: ignore
+
+        spark = DatabricksSession.builder.getOrCreate()
+        sql = spark.sql  # type: ignore
+    except Exception as e:
+        # We are ignoring all failures here because user might want to initialize
+        # spark session themselves and we don't want to interfere with that
+        logging.debug(f"Failed to initialize globals 'spark' and 'sql', continuing. Cause: {e}")
+
+    try:
+        # We expect this to fail locally since dbconnect does not support sparkcontext. This is just for typing
+        sc = spark.sparkContext  # type: ignore
+    except Exception as e:
+        logging.debug(f"Failed to initialize global 'sc', continuing. Cause: {e}")
+
+    def display(input=None, *args, **kwargs) -> None:  # type: ignore
+        """
+        Display plots or data.
+        Display plot:
+                        - display() # no-op
+                        - display(matplotlib.figure.Figure)
+        Display dataset:
+                        - display(spark.DataFrame)
+                        - display(list) # if list can be converted to DataFrame, e.g., list of named tuples
+                        - display(pandas.DataFrame)
+                        - display(koalas.DataFrame)
+                        - display(pyspark.pandas.DataFrame)
+        Display any other value that has a _repr_html_() method
+        For Spark 2.0 and 2.1:
+                        - display(DataFrame, streamName='optional', trigger=optional pyspark.sql.streaming.Trigger,
+                                                        checkpointLocation='optional')
+        For Spark 2.2+:
+                        - display(DataFrame, streamName='optional', trigger=optional interval like '1 second',
+                                                        checkpointLocation='optional')
+        """
+        # Import inside the function so that imports are only triggered on usage.
+        from IPython import display as IPDisplay
+
+        return IPDisplay.display(input, *args, **kwargs)  # type: ignore
+
+    def displayHTML(html) -> None:  # type: ignore
+        """
+        Display HTML data.
+        Parameters
+        ----------
+        data : URL or HTML string
+                        If data is a URL, display the resource at that URL, the resource is loaded dynamically by the browser.
+                        Otherwise data should be the HTML to be displayed.
+        See also:
+        IPython.display.HTML
+        IPython.display.display_html
+        """
+        # Import inside the function so that imports are only triggered on usage.
+        from IPython import display as IPDisplay
+
+        return IPDisplay.display_html(html, raw=True)  # type: ignore
+
+    # We want to propagate the error in initialising dbutils because this is a core
+    # functionality of the sdk
+    from databricks.sdk.dbutils import RemoteDbUtils
+
+    from . import dbutils_stub
+
+    dbutils_type = Union[dbutils_stub.dbutils, RemoteDbUtils]
+
+    dbutils = RemoteDbUtils()
+    dbutils = cast(dbutils_type, dbutils)
+
+    # We do this to prevent importing widgets implementation prematurely
+    # The widget import should prompt users to use the implementation
+    # which has ipywidget support.
+    def getArgument(name: str, defaultValue: Optional[str] = None):
+        return dbutils.widgets.getArgument(name, defaultValue)
+
+
+__all__ = dbruntime_objects

venv/lib/python3.9/site-packages/databricks/sdk/runtime/dbutils_stub.py

@@ -0,0 +1,373 @@
+import typing
+from collections import namedtuple
+
+
+class FileInfo(namedtuple("FileInfo", ["path", "name", "size", "modificationTime"])):
+    pass
+
+
+class MountInfo(namedtuple("MountInfo", ["mountPoint", "source", "encryptionType"])):
+    pass
+
+
+class SecretScope(namedtuple("SecretScope", ["name"])):
+
+    def getName(self):
+        return self.name
+
+
+class SecretMetadata(namedtuple("SecretMetadata", ["key"])):
+    pass
+
+
+class dbutils:
+
+    class credentials:
+        """
+        Utilities for interacting with credentials within notebooks
+        """
+
+        @staticmethod
+        def assumeRole(role: str) -> bool:
+            """
+            Sets the role ARN to assume when looking for credentials to authenticate with S3
+            """
+            ...
+
+        @staticmethod
+        def showCurrentRole() -> typing.List[str]:
+            """
+            Shows the currently set role
+            """
+            ...
+
+        @staticmethod
+        def showRoles() -> typing.List[str]:
+            """
+            Shows the set of possibly assumed roles
+            """
+            ...
+
+        @staticmethod
+        def getCurrentCredentials() -> typing.Mapping[str, str]: ...
+
+    class data:
+        """
+        Utilities for understanding and interacting with datasets (EXPERIMENTAL)
+        """
+
+        @staticmethod
+        def summarize(df: any, precise: bool = False) -> None:
+            """Summarize a Spark/pandas/Koalas DataFrame and visualize the statistics to get quick insights.
+
+            Example: dbutils.data.summarize(df)
+
+            :param df: A pyspark.sql.DataFrame, pyspark.pandas.DataFrame, databricks.koalas.DataFrame
+            or pandas.DataFrame object to summarize. Streaming dataframes are not supported.
+            :param precise: If false, percentiles, distinct item counts, and frequent item counts
+            will be computed approximately to reduce the run time.
+            If true, distinct item counts and frequent item counts will be computed exactly,
+            and percentiles will be computed with high precision.
+
+            :return: visualization of the computed summmary statistics.
+            """
+            ...
+
+    class fs:
+        """
+        Manipulates the Databricks filesystem (DBFS) from the console
+        """
+
+        @staticmethod
+        def cp(source: str, dest: str, recurse: bool = False) -> bool:
+            """
+            Copies a file or directory, possibly across FileSystems
+            """
+            ...
+
+        @staticmethod
+        def head(file: str, max_bytes: int = 65536) -> str:
+            """
+            Returns up to the first 'maxBytes' bytes of the given file as a String encoded in UTF-8
+            """
+            ...
+
+        @staticmethod
+        def ls(path: str) -> typing.List[FileInfo]:
+            """
+            Lists the contents of a directory
+            """
+            ...
+
+        @staticmethod
+        def mkdirs(dir: str) -> bool:
+            """
+            Creates the given directory if it does not exist, also creating any necessary parent directories
+            """
+            ...
+
+        @staticmethod
+        def mv(source: str, dest: str, recurse: bool = False) -> bool:
+            """
+            Moves a file or directory, possibly across FileSystems
+            """
+            ...
+
+        @staticmethod
+        def put(file: str, contents: str, overwrite: bool = False) -> bool:
+            """
+            Writes the given String out to a file, encoded in UTF-8
+            """
+            ...
+
+        @staticmethod
+        def rm(dir: str, recurse: bool = False) -> bool:
+            """
+            Removes a file or directory
+            """
+            ...
+
+        @staticmethod
+        def cacheFiles(*files): ...
+
+        @staticmethod
+        def cacheTable(name: str): ...
+
+        @staticmethod
+        def uncacheFiles(*files): ...
+
+        @staticmethod
+        def uncacheTable(name: str): ...
+
+        @staticmethod
+        def mount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
+            """
+            Mounts the given source directory into DBFS at the given mount point
+            """
+            ...
+
+        @staticmethod
+        def updateMount(
+            source: str,
+            mount_point: str,
+            encryption_type: str = "",
+            owner: typing.Optional[str] = None,
+            extra_configs: typing.Mapping[str, str] = {},
+        ) -> bool:
+            """
+            Similar to mount(), but updates an existing mount point (if present) instead of creating a new one
+            """
+            ...
+
+        @staticmethod
+        def mounts() -> typing.List[MountInfo]:
+            """
+            Displays information about what is mounted within DBFS
+            """
+            ...
+
+        @staticmethod
+        def refreshMounts() -> bool:
+            """
+            Forces all machines in this cluster to refresh their mount cache, ensuring they receive the most recent information
+            """
+            ...
+
+        @staticmethod
+        def unmount(mount_point: str) -> bool:
+            """
+            Deletes a DBFS mount point
+            """
+            ...
+
+    class jobs:
+        """
+        Utilities for leveraging jobs features
+        """
+
+        class taskValues:
+            """
+            Provides utilities for leveraging job task values
+            """
+
+            @staticmethod
+            def get(
+                taskKey: str,
+                key: str,
+                default: any = None,
+                debugValue: any = None,
+            ) -> None:
+                """
+                Returns the latest task value that belongs to the current job run
+                """
+                ...
+
+            @staticmethod
+            def set(key: str, value: any) -> None:
+                """
+                Sets a task value on the current task run
+                """
+                ...
+
+    class library:
+        """
+        Utilities for session isolated libraries
+        """
+
+        @staticmethod
+        def restartPython() -> None:
+            """
+            Restart python process for the current notebook session
+            """
+            ...
+
+    class notebook:
+        """
+        Utilities for the control flow of a notebook (EXPERIMENTAL)
+        """
+
+        @staticmethod
+        def exit(value: str) -> None:
+            """
+            This method lets you exit a notebook with a value
+            """
+            ...
+
+        @staticmethod
+        def run(
+            path: str,
+            timeout_seconds: int,
+            arguments: typing.Mapping[str, str],
+        ) -> str:
+            """
+            This method runs a notebook and returns its exit value
+            """
+            ...
+
+    class secrets:
+        """
+        Provides utilities for leveraging secrets within notebooks
+        """
+
+        @staticmethod
+        def get(scope: str, key: str) -> str:
+            """
+            Gets the string representation of a secret value with scope and key
+            """
+            ...
+
+        @staticmethod
+        def getBytes(self, scope: str, key: str) -> bytes:
+            """Gets the bytes representation of a secret value for the specified scope and key."""
+
+        @staticmethod
+        def list(scope: str) -> typing.List[SecretMetadata]:
+            """
+            Lists secret metadata for secrets within a scope
+            """
+            ...
+
+        @staticmethod
+        def listScopes() -> typing.List[SecretScope]:
+            """
+            Lists secret scopes
+            """
+            ...
+
+    class widgets:
+        """
+        provides utilities for working with notebook widgets. You can create different types of widgets and get their bound value
+        """
+
+        @staticmethod
+        def get(name: str) -> str:
+            """Returns the current value of a widget with give name.
+            :param name: Name of the argument to be accessed
+            :return: Current value of the widget or default value
+            """
+            ...
+
+        @staticmethod
+        def getArgument(name: str, defaultValue: typing.Optional[str] = None) -> typing.Optional[str]:
+            """Returns the current value of a widget with give name.
+            :param name: Name of the argument to be accessed
+            :param defaultValue: (Deprecated) default value
+            :return: Current value of the widget or default value
+            """
+            ...
+
+        @staticmethod
+        def text(name: str, defaultValue: str, label: str = None):
+            """Creates a text input widget with given name, default value and optional label for
+            display
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def dropdown(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: str = None,
+        ):
+            """Creates a dropdown input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget (must be one of choices)
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def combobox(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
+            """Creates a combobox input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def multiselect(
+            name: str,
+            defaultValue: str,
+            choices: typing.List[str],
+            label: typing.Optional[str] = None,
+        ):
+            """Creates a multiselect input widget with given specification.
+            :param name: Name of argument associated with the new input widget
+            :param defaultValue: Default value of the input widget (must be one of choices)
+            :param choices: List of choices for the dropdown input widget
+            :param label: Optional label string for display in notebook and dashboard
+            """
+            ...
+
+        @staticmethod
+        def remove(name: str):
+            """Removes given input widget. If widget does not exist it will throw an error.
+            :param name: Name of argument associated with input widget to be removed
+            """
+            ...
+
+        @staticmethod
+        def removeAll():
+            """Removes all input widgets in the notebook."""
+            ...
+
+
+getArgument = dbutils.widgets.getArgument