Schorle patch (#17)

* test hard coded path
* ensure that content type is text/javascript for strict content type execution
* reorder as I changed it to a type of mutable dict
* decompose regardless when content is being sent
* remove anything hardcoded
* fix websocket for bundle.js
* use multidict instead of converting headers to dictionaries
* pop due to multidict
* pop all instances of a key
* use CIMultiDict because headers are case insensitive
* removed schorle due to it requiring cookies to work
* added inject_sql_warehouse example

README.md

@@ -94,8 +94,6 @@ support the access to request object to access headers, etc.
 * [ ] cloudflared
 * [x] dbtunnel custom relay (private only)
 
-
-
 ## Setup
 
 **Please do not use this in production!!**
@@ -105,6 +103,60 @@ support the access to request object to access headers, etc.
 3. Enjoy your proxy experience :-) 
 4. If you want to share the link ensure that the other user has permission to attach to your cluster.
 
+## Passing databricks auth to your app via `inject_auth`
+
+You can pass databricks user auth from your notebook session to any of the frameworks by doing the following:
+
+```python
+from dbtunnel import dbtunnel
+dbtunnel.<framework>(<script_path>).inject_auth().run()
+```
+
+For example: 
+
+```python
+from dbtunnel import dbtunnel
+dbtunnel.gradio(demo).inject_auth().run()
+```
+
+This exposes the user information via environment variable DATABRICKS_HOST and DATABRICKS_TOKEN.
+
+
+## Passing a warehouse to your app via `inject_sql_warehouse`
+
+You can pass databricks warehouse auth from your notebook session to any of the frameworks by doing the following:
+
+```python
+from dbtunnel import dbtunnel
+dbtunnel.<framework>(<script_path>).inject_sql_warehouse().run()
+```
+
+This exposes the warehouse information via environment variable DATABRICKS_HOST, DATABRICKS_TOKEN and DATABRICKS_HTTP_PATH.
+
+
+## Passing custom environment variables via `inject_env`
+
+You can pass custom environment variables from your notebook to any of the frameworks by doing the following:
+
+```python
+from dbtunnel import dbtunnel
+dbtunnel.<framework>(<script_path>).inject_env({
+   "MY_CUSTOM_ENV": "my_custom_env_value"
+}).run()
+```
+
+For example: 
+
+```python
+from dbtunnel import dbtunnel
+dbtunnel.gradio(demo).inject_env({
+    "MY_CUSTOM_ENV": "my_custom_env_value"
+}).run()
+```
+
+Keep in mind environment variables need to be strings!
+
+
 ## Exposing to internet using ngrok
 
 **WARNING: IT WILL BE PUBLICLY AVAILABLE TO ANYONE WITH THE LINK SO DO NOT EXPOSE ANYTHING SENSITIVE**

dbtunnel/__init__.py

@@ -85,10 +85,5 @@ def uvicorn(app, port: int = 8080):
         from dbtunnel.uvicorn import UvicornAppTunnel
         return UvicornAppTunnel(app, port)
 
-    @staticmethod
-    def schorle(app, port: int = 8080):
-        from dbtunnel.schorle import SchorleAppTunnel
-        return SchorleAppTunnel(app, port)
-
 
 dbtunnel = AppTunnels()

dbtunnel/tunnels.py

@@ -92,7 +92,7 @@ def get_cloud_proxy_settings(cloud: str, org_id: str, cluster_id: str, port: int
 
 Flavor = Literal[
     "gradio", "fastapi", "nicegui", "streamlit", "stable-diffusion-ui", "bokeh", "flask", "dash", "solara",
-    "code-server", "chainlit", "shiny-python", "uvicorn", "schorle"]
+    "code-server", "chainlit", "shiny-python", "uvicorn"]
 
 
 def get_current_username() -> str:
@@ -158,6 +158,13 @@ def shared(self):
         return self._share
 
     def inject_auth(self, host: str = None, token: str = None, write_cfg: bool = False):
+        """
+        Inject databricks host and token into the environment
+
+        :param host: any databricks url you may want to use, otherwise defaults to your notebook session host
+        :param token: any databricks token you may want to use, otherwise defaults to your notebook session token
+        :param write_cfg: this will create a .databrickscfg file in the root directory, only works on single user clusters.
+        """
         if os.getenv("DATABRICKS_HOST") is None:
             self._log.info("Setting databricks host from context")
             os.environ["DATABRICKS_HOST"] = host or ensure_scheme(ctx.host)
@@ -173,6 +180,13 @@ def inject_auth(self, host: str = None, token: str = None, write_cfg: bool = Fal
         return self
 
     def inject_sql_warehouse(self, http_path: str, server_hostname: str = None, token: str = None):
+        """
+        Inject databricks warehouse http path into the environment and auth
+        :param http_path: the http path to the warehouse
+        :param server_hostname: the hostname of the databricks workspace, defaults to the current workspace where the notebook is running
+        :param token: any databricks token you may want to use, otherwise defaults to your notebook session token
+        :return:
+        """
         if os.getenv("DATABRICKS_SERVER_HOSTNAME") is None:
             self._log.info("Setting databricks server hostname from context")
             os.environ["DATABRICKS_SERVER_HOSTNAME"] = server_hostname or extract_hostname(
@@ -189,6 +203,15 @@ def inject_sql_warehouse(self, http_path: str, server_hostname: str = None, toke
         return self
 
     def inject_env(self, **kwargs):
+        """
+        Inject environment variables into the environment. Keep in mind environment variables are case sensitive
+
+        Example usage:
+        dbtunnel.chainlit("path/to/script").inject_env(ENV_VAR1="value1", ENV_VAR2="value2").run()
+
+        :param kwargs: keyword arguments for the environment variables you want to set
+        :return:
+        """
         for k, v in kwargs.items():
             if type(v) != str:
                 raise ValueError(f"Value for environment variable {k} must be a string")
@@ -197,6 +220,11 @@ def inject_env(self, **kwargs):
         return self
 
     def with_token_auth(self):
+        """
+        Experimental feature do not use.
+        :return:
+        """
+
         self._basic_tunnel_auth["token_auth"] = True
         self._basic_tunnel_auth["token_auth_workspace_url"] = ctx.host
         return self

dbtunnel/vendor/asgiproxy/config.py

@@ -1,13 +1,14 @@
-from typing import Optional, Union, Iterable, Dict, Callable
+from typing import Optional, Iterable, Dict, Callable
 from urllib.parse import urljoin
 
 import aiohttp
+from multidict import MultiDict, CIMultiDict
 from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.types import Scope
 from starlette.websockets import WebSocket
 
-Headerlike = Union[dict, Headers]
+Headerlike = MultiDict
 
 
 class ProxyConfig:
@@ -48,7 +49,8 @@ def process_upstream_headers(
         """
         Process upstream HTTP headers before they're passed to the client.
         """
-        return proxy_response.headers  # type: ignore
+        headers = CIMultiDict(proxy_response.headers)
+        return headers  # type: ignore
 
     def get_upstream_http_options(
             self, *, scope: Scope, client_request: Request, data

dbtunnel/vendor/asgiproxy/frameworks.py

@@ -133,49 +133,14 @@ def _modify_js_bundle(content, root_path):
     )()
     return config
 
-
-def _make_schorle_local_proxy_config(
-        url_base_path,
-        service_host: str = "0.0.0.0",
-        service_port: int = 9989,
-        auth_config: dict = None
-):
-    auth_config = auth_config or {}
-
-    def _modify_root(content, root_path):
-        list_of_uris = [b"/_schorle",]
-        for uri in list_of_uris:
-            content = content.replace(uri, root_path.rstrip("/").encode("utf-8") + uri)
-        return content
-
-    modify_root = functools.partial(_modify_root, root_path=url_base_path)
-
-    config = type(
-        "Config",
-        (BaseURLProxyConfigMixin, ProxyConfig),
-        {
-            "upstream_base_url": f"http://{service_host}:{service_port}",
-            "rewrite_host_header": f"{service_host}:{service_port}",
-            "modify_content": {
-                "/": modify_root,
-                # some reason gradio also has caps index bundled calling out explicitly
-            },
-            **auth_config,
-        },
-    )()
-    return config
-
-
 class Frameworks:
     STREAMLIT: str = "streamlit"
     GRADIO: str = "gradio"
     CHAINLIT: str = "chainlit"
-    SCHORLE: str = "schorle"
 
 
 framework_specific_proxy_config = {
     Frameworks.STREAMLIT: _make_streamlit_local_proxy_config,
     Frameworks.GRADIO: _make_gradio_local_proxy_config,
     Frameworks.CHAINLIT: _make_chainlit_local_proxy_config,
-    Frameworks.SCHORLE: _make_schorle_local_proxy_config,
 }

dbtunnel/vendor/asgiproxy/proxies/http.py

@@ -86,9 +86,10 @@ async def convert_proxy_response_to_user_response(
     if is_from_databricks_proxy(scope) is True:
         if response_content is not None and len(response_content) > 0 and context.config.modify_content is not None:
             for path_pattern, modify_func in context.config.modify_content.items():
+                # if path is .js it should be of type text/javascript;charset=utf-8
                 if fnmatch.fnmatch(scope["path"], path_pattern):
                     response_content = modify_func(response_content)
-                    new_headers = {k: v for k, v in headers_to_client.items()}
+                    new_headers.popall("Content-Length", None)
                     new_headers["Content-Length"] = str(len(response_content))
 
     return Response(

setup.py

@@ -29,10 +29,6 @@
             # Specify dependencies for building documentation here
             "uvicorn",
         ],
-        "schorle": [
-            "schorle",
-            "uvicorn",
-        ],
         "streamlit": [
             # Specify dependencies for building documentation here
             "streamlit",