dataproc-spark-connect 0.1.0__py2.py3-none-any.whl

dataproc_spark_connect-0.1.0.dist-info/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

dataproc_spark_connect-0.1.0.dist-info/METADATA

@@ -0,0 +1,11 @@
+Metadata-Version: 2.1
+Name: dataproc-spark-connect
+Version: 0.1.0
+Summary: Dataproc client library for Spark Connect
+License-File: LICENSE
+Requires-Dist: wheel
+Requires-Dist: websockets
+Requires-Dist: pyspark >=3.5
+Requires-Dist: pandas
+Requires-Dist: pyarrow
+

dataproc_spark_connect-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+google/cloud/dataproc_spark_connect/__init__.py,sha256=pybAofW6rmWI-4C8VYm1q0NOZD_sBvFQz43jUBSQW30,616
+google/cloud/dataproc_spark_connect/session.py,sha256=XCOym1Llp_vXCgMQ0EMsORoImY1BSKfJrZpooGCurNc,23013
+google/cloud/dataproc_spark_connect/client/__init__.py,sha256=6hCNSsgYlie6GuVpc5gjFsPnyeMTScTpXSPYqp1fplY,615
+google/cloud/dataproc_spark_connect/client/core.py,sha256=7Wy6QwkcWxlHBdo4NsktJEknggPpGkx9F5CS5IpQ7iM,3630
+google/cloud/dataproc_spark_connect/client/proxy.py,sha256=o5ppDkvUUQLy4uDMjW6roTMgRFKlVfQm5faC5b60rRk,7263
+dataproc_spark_connect-0.1.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-0.1.0.dist-info/METADATA,sha256=Ke0zfRLuybTquQJgm-Ve3t8tVMT26l6q5i3R8wUEwq4,261
+dataproc_spark_connect-0.1.0.dist-info/WHEEL,sha256=ED2S3aolPA63OeAfpmbtiQd0NXmZX4SrTHXvXF6JNGc,109
+dataproc_spark_connect-0.1.0.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-0.1.0.dist-info/RECORD,,

dataproc_spark_connect-0.1.0.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: setuptools (72.0.0)
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any
+

dataproc_spark_connect-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+google

google/cloud/dataproc_spark_connect/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .session import DataprocSparkSession

google/cloud/dataproc_spark_connect/client/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .core import DataprocChannelBuilder

google/cloud/dataproc_spark_connect/client/core.py

@@ -0,0 +1,113 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import google
+import grpc
+from pyspark.sql.connect.client import ChannelBuilder
+
+from . import proxy
+
+
+class DataprocChannelBuilder(ChannelBuilder):
+    """
+    This is a helper class that is used to create a GRPC channel based on the given
+    connection string per the documentation of Spark Connect.
+
+    This implementation of ChannelBuilder uses `secure_authorized_channel` from the
+    `google.auth.transport.grpc` package for authenticating secure channel.
+
+    Examples
+    --------
+    >>> cb =  ChannelBuilder("sc://localhost")
+    ... cb.endpoint
+
+    >>> cb = ChannelBuilder("sc://localhost/;use_ssl=true;token=aaa")
+    ... cb.secure
+    True
+    """
+
+    def toChannel(self) -> grpc.Channel:
+        """
+        Applies the parameters of the connection string and creates a new
+        GRPC channel according to the configuration. Passes optional channel options to
+        construct the channel.
+
+        Returns
+        -------
+        GRPC Channel instance.
+        """
+        # TODO: Replace with a direct channel once all compatibility issues with
+        # grpc have been resolved.
+        return self._proxied_channel()
+
+    def _proxied_channel(self) -> grpc.Channel:
+        return ProxiedChannel(self.host)
+
+    def _direct_channel(self) -> grpc.Channel:
+        destination = f"{self.host}:{self.port}"
+
+        credentials, project = google.auth.default(
+            scopes=["https://www.googleapis.com/auth/cloud-platform"]
+        )
+        # Get an HTTP request function to refresh credentials.
+        request = google.auth.transport.requests.Request()
+        # Create a channel.
+
+        return google.auth.transport.grpc.secure_authorized_channel(
+            credentials,
+            request,
+            destination,
+            None,
+            None,
+            options=self._channel_options,
+        )
+
+
+class ProxiedChannel(grpc.Channel):
+
+    def __init__(self, target_host):
+        self._proxy = proxy.DataprocSessionProxy(0, target_host)
+        self._proxy.start()
+        self._proxied_connect_url = f"sc://localhost:{self._proxy.port}"
+        self._wrapped = ChannelBuilder(self._proxied_connect_url).toChannel()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        ret = self._wrapped.__exit__(*args)
+        self._proxy.stop()
+        return ret
+
+    def close(self):
+        ret = self._wrapped.close()
+        self._proxy.stop()
+        return ret
+
+    def stream_stream(self, *args, **kwargs):
+        return self._wrapped.stream_stream(*args, **kwargs)
+
+    def stream_unary(self, *args, **kwargs):
+        return self._wrapped.stream_unary(*args, **kwargs)
+
+    def subscribe(self, *args, **kwargs):
+        return self._wrapped.subscribe(*args, **kwargs)
+
+    def unary_stream(self, *args, **kwargs):
+        return self._wrapped.unary_stream(*args, **kwargs)
+
+    def unary_unary(self, *args, **kwargs):
+        return self._wrapped.unary_unary(*args, **kwargs)
+
+    def unsubscribe(self, *args, **kwargs):
+        return self._wrapped.unsubscribe(*args, **kwargs)

google/cloud/dataproc_spark_connect/client/proxy.py

@@ -0,0 +1,227 @@
+#!/bin/env python
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import contextlib
+import logging
+import socket
+import threading
+import time
+
+import websockets.sync.client as websocketclient
+
+from google import auth as googleauth
+from google.auth.transport import requests as googleauthrequests
+
+parser = argparse.ArgumentParser()
+parser.add_argument("port")
+parser.add_argument("target_host")
+
+
+logger = logging.getLogger(__name__)
+
+
+class bridged_socket(object):
+    """Socket-like object that uses a websocket-over-TCP Bridge transport.
+
+    See: https://github.com/google/inverting-proxy/tree/master/utils/tcpbridge
+    """
+
+    def __init__(self, websocket_conn):
+        self._conn = websocket_conn
+
+    def recv(self, buff_size):
+        msg = self._conn.recv()
+        return bytes.fromhex(msg)
+
+    def send(self, msg_bytes):
+        msg = bytes.hex(msg_bytes)
+        self._conn.send(msg)
+
+
+def connect_tcp_bridge(hostname):
+    """Create a socket-like connection to the given hostname using websocket.
+
+    The backend server connected to over the websocket connection must be
+    running the TCP-bridge backend corresponding to this frontend.
+
+    Args:
+        hostname: The hostname of the server running the TCP-bridge backend.
+
+    Returns:
+        A socket-like object with `recv` and `send` methods.
+    """
+    path = "tcp-over-websocket-bridge/35218cb7-1201-4940-89e8-48d8f03fed96"
+    creds, _ = googleauth.default(
+        scopes=["https://www.googleapis.com/auth/cloud-platform"]
+    )
+    creds.refresh(googleauthrequests.Request())
+
+    return websocketclient.connect(
+        f"wss://{hostname}/{path}",
+        additional_headers={"Authorization": f"Bearer {creds.token}"},
+    )
+
+
+def forward_bytes(name, from_sock, to_sock):
+    """Continuously stream bytes from the `from_sock` to the `to_sock`.
+
+    This method terminates when either the `from_sock` is closed (causing
+    it to return a Falsy value from its `recv` method), or the first time
+    it hits an exception.
+
+    This method is intended to be run in a separate thread of execution.
+
+    Args:
+        from_sock: A socket-like object to stream bytes from.
+        to_sock: A socket-like object to stream bytes to.
+    """
+    while True:
+        try:
+            bs = from_sock.recv(1024)
+            if not bs:
+                return
+            to_sock.send(bs)
+        except Exception as ex:
+            logger.debug(f"[{name}] Exception forwarding bytes: {ex}")
+            return
+
+
+def forward_connection(conn_number, conn, addr, target_host):
+    """Create a connection to the target and forward `conn` to it.
+
+    This method creates a socket-like object holding a connection to the given
+    target host, and then continuously streams bytes in both directions between
+    `conn` and that newly created connection.
+
+    Both the supplied incoming connection (`conn`) and the created outgoing
+    connection are automatically closed when this method terminates.
+
+    This method should be run inside of a daemon thread so that it will not
+    block program termination.
+    """
+    with conn:
+        with connect_tcp_bridge(target_host) as websocket_conn:
+            backend_socket = bridged_socket(websocket_conn)
+            forward_name = f"{conn_number}-forward"
+            t1 = threading.Thread(
+                name=forward_name,
+                target=forward_bytes,
+                args=[forward_name, conn, backend_socket],
+                daemon=True,
+            )
+            t1.start()
+            backward_name = f"{conn_number}-backward"
+            t2 = threading.Thread(
+                name=backward_name,
+                target=forward_bytes,
+                args=[backward_name, backend_socket, conn],
+                daemon=True,
+            )
+            t2.start()
+            t1.join()
+            t2.join()
+
+
+class DataprocSessionProxy(object):
+    """A TCP proxy for forwarding requests to Dataproc Serverless Sessions.
+
+    Spark Connect clients connect to this proxy using the h2c (without-SSL)
+    protocol, and this proxy adds SSL by tunneling those connections over
+    an HTTPS/WebSocket connection to the backend server.
+
+    The tunneled requests are authenticated using the Google Application
+    Default Credentials.
+    """
+
+    def __init__(self, port, target_host):
+        self._port = port
+        self._target_host = target_host
+        self._started = False
+        self._killed = False
+        self._conn_number = 0
+
+    @property
+    def port(self):
+        """The local port the proxy is listening on"""
+        return self._port
+
+    def start(self, daemon=True):
+        """Start the proxy.
+
+        By the time this method returns the proxy has already started listening
+        on its local port will accept incoming connections.
+        """
+        if self._started:
+            raise Exception("Dataproc session proxy already started")
+        self._started = True
+        s = threading.Semaphore(value=0)
+        t = threading.Thread(target=self._run, args=[s], daemon=daemon)
+        t.start()
+        s.acquire()
+
+    def _run(self, s):
+        with socket.create_server(("127.0.0.1", self._port)) as frontend_socket:
+            if self._port == 0:
+                self._port = frontend_socket.getsockname()[1]
+            s.release()
+            while not self._killed:
+                conn, addr = frontend_socket.accept()
+                logger.debug(f"Accepted a connection from {addr}...")
+                self._conn_number += 1
+                threading.Thread(
+                    target=forward_connection,
+                    args=[self._conn_number, conn, addr, self._target_host],
+                    daemon=True,
+                ).start()
+
+    def stop(self):
+        """Stop the proxy."""
+        self._killed = True
+
+
+@contextlib.contextmanager
+def dataproc_session_proxy(port, target_host):
+    """Context manager for creating a Dataproc Session proxy.
+
+    Usage:
+        with dataproc_session_proxy(0, backend_hostname) as p:
+           local_port = p.port
+           ...
+
+    Args:
+        port: The local port to listen on. Use `0` to pick a free port.
+        target_host: The backend to proxy connections to.
+
+    Returns:
+        A context manager wrapping a DataprocSessionProxy instance.
+    """
+    proxy = DataprocSessionProxy(port, target_host)
+    try:
+        proxy.start(daemon=False)
+        yield proxy
+    finally:
+        proxy.stop()
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    with dataproc_session_proxy(int(args.port), args.target_host) as p:
+        print(f"Proxy listening on port {p.port}")
+        try:
+            while True:
+                pass
+        except KeyboardInterrupt:
+            pass

google/cloud/dataproc_spark_connect/session.py

@@ -0,0 +1,540 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import logging
+import os
+import random
+import string
+import time
+import datetime
+from time import sleep
+from typing import Any, cast, ClassVar, Dict, Optional
+
+from google.api_core import retry
+from google.api_core.future.polling import POLLING_PREDICATE
+from google.api_core.client_options import ClientOptions
+from google.api_core.exceptions import FailedPrecondition, InvalidArgument, NotFound
+from google.cloud.dataproc_v1.types import sessions
+
+from google.cloud.dataproc_spark_connect.client import DataprocChannelBuilder
+from google.cloud.dataproc_v1 import (
+    CreateSessionRequest,
+    GetSessionRequest,
+    Session,
+    SessionControllerClient,
+    SessionTemplate,
+    TerminateSessionRequest,
+)
+from google.protobuf import text_format
+from google.protobuf.text_format import ParseError
+from pyspark.sql.connect.session import SparkSession
+from pyspark.sql.utils import to_str
+
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DataprocSparkSession(SparkSession):
+    """The entry point to programming Spark with the Dataset and DataFrame API.
+
+    A DataprocRemoteSparkSession can be used to create :class:`DataFrame`, register :class:`DataFrame` as
+    tables, execute SQL over tables, cache tables, and read parquet files.
+
+    Examples
+    --------
+
+    Create a Spark session with Dataproc Spark Connect.
+
+    >>> spark = (
+    ...     DataprocSparkSession.builder
+    ...         .appName("Word Count")
+    ...         .dataprocConfig(Session())
+    ...         .getOrCreate()
+    ... ) # doctest: +SKIP
+    """
+
+    _active_s8s_session_uuid: ClassVar[Optional[str]] = None
+    _project_id = None
+    _region = None
+    _client_options = None
+    _active_s8s_session_id: ClassVar[Optional[str]] = None
+
+    class Builder(SparkSession.Builder):
+
+        _dataproc_runtime_spark_version = {"3.0": "3.5.1", "2.2": "3.5.0"}
+
+        _session_static_configs = [
+            "spark.executor.cores",
+            "spark.executor.memoryOverhead",
+            "spark.executor.memory",
+            "spark.driver.memory",
+            "spark.driver.cores",
+            "spark.eventLog.dir",
+            "spark.history.fs.logDirectory",
+        ]
+
+        def __init__(self):
+            self._options: Dict[str, Any] = {}
+            self._channel_builder: Optional[DataprocChannelBuilder] = None
+            self._dataproc_config: Optional[Session] = None
+            self._project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
+            self._region = os.environ.get("GOOGLE_CLOUD_REGION")
+            self._client_options = ClientOptions(
+                api_endpoint=os.environ.get(
+                    "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT",
+                    f"{self._region}-dataproc.googleapis.com",
+                )
+            )
+
+        def __apply_options(self, session: "SparkSession") -> None:
+            with self._lock:
+                self._options = {
+                    key: value
+                    for key, value in self._options.items()
+                    if key not in self._session_static_configs
+                }
+                self._apply_options(session)
+
+        def projectId(self, project_id):
+            self._project_id = project_id
+            return self
+
+        def region(self, region):
+            self._region = region
+            self._client_options.api_endpoint = os.environ.get(
+                "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT",
+                f"{self._region}-dataproc.googleapis.com",
+            )
+            return self
+
+        def dataprocConfig(self, dataproc_config: Session):
+            with self._lock:
+                self._dataproc_config = dataproc_config
+                for k, v in dataproc_config.runtime_config.properties.items():
+                    self._options[cast(str, k)] = to_str(v)
+                return self
+
+        def remote(self, url: Optional[str] = None) -> "SparkSession.Builder":
+            if url:
+                raise NotImplemented(
+                    "DataprocSparkSession does not support connecting to an existing remote server"
+                )
+            else:
+                return self
+
+        def create(self) -> "SparkSession":
+            raise NotImplemented(
+                "DataprocSparkSession allows session creation only through getOrCreate"
+            )
+
+        def __create_spark_connect_session_from_s8s(
+            self, session_response
+        ) -> "SparkSession":
+            DataprocSparkSession._active_s8s_session_uuid = (
+                session_response.uuid
+            )
+            DataprocSparkSession._project_id = self._project_id
+            DataprocSparkSession._region = self._region
+            DataprocSparkSession._client_options = self._client_options
+            spark_connect_url = session_response.runtime_info.endpoints.get(
+                "Spark Connect Server"
+            )
+            spark_connect_url = spark_connect_url.replace("https", "sc")
+            url = f"{spark_connect_url.replace('.com/', '.com:443/')};session_id={session_response.uuid};use_ssl=true"
+            logger.debug(f"Spark Connect URL: {url}")
+            self._channel_builder = DataprocChannelBuilder(url)
+
+            assert self._channel_builder is not None
+            session = DataprocSparkSession(connection=self._channel_builder)
+
+            DataprocSparkSession._set_default_and_active_session(session)
+            self.__apply_options(session)
+            return session
+
+        def __create(self) -> "SparkSession":
+            with self._lock:
+
+                if self._options.get("spark.remote", False):
+                    raise NotImplemented(
+                        "DataprocSparkSession does not support connecting to an existing remote server"
+                    )
+
+                from google.cloud.dataproc_v1 import SessionControllerClient
+
+                dataproc_config: Session = self._get_dataproc_config()
+                session_template: SessionTemplate = self._get_session_template()
+
+                self._get_and_validate_version(
+                    dataproc_config, session_template
+                )
+
+                spark_connect_session = self._get_spark_connect_session(
+                    dataproc_config, session_template
+                )
+
+                spark = self._get_spark(dataproc_config, session_template)
+
+                if not spark_connect_session:
+                    dataproc_config.spark_connect_session = {}
+                if not spark:
+                    dataproc_config.spark = {}
+                os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
+                session_request = CreateSessionRequest()
+                session_id = self.generate_dataproc_session_id()
+
+                session_request.session_id = session_id
+                dataproc_config.name = f"projects/{self._project_id}/regions/{self._region}/sessions/{session_id}"
+                logger.debug(
+                    f"Configurations used to create serverless session:\n {dataproc_config}"
+                )
+                session_request.session = dataproc_config
+                session_request.parent = (
+                    f"projects/{self._project_id}/regions/{self._region}"
+                )
+
+                logger.debug("Creating serverless session")
+                DataprocSparkSession._active_s8s_session_id = session_id
+                s8s_creation_start_time = time.time()
+                try:
+                    session_polling = retry.Retry(
+                        predicate=POLLING_PREDICATE,
+                        initial=5.0,  # seconds
+                        maximum=5.0,  # seconds
+                        multiplier=1.0,
+                        timeout=600,  # seconds
+                    )
+                    logger.info(
+                        "Creating Spark session. It may take few minutes."
+                    )
+                    operation = SessionControllerClient(
+                        client_options=self._client_options
+                    ).create_session(session_request)
+                    print(
+                        f"Interactive Session Detail View:  https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id}"
+                    )
+                    session_response: Session = operation.result(
+                        polling=session_polling
+                    )
+                    if (
+                        "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH"
+                        in os.environ
+                    ):
+                        file_path = os.environ[
+                            "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH"
+                        ]
+                        try:
+                            session_data = {
+                                "session_name": session_response.name,
+                                "session_uuid": session_response.uuid,
+                            }
+                            os.makedirs(
+                                os.path.dirname(file_path), exist_ok=True
+                            )
+                            with open(file_path, "w") as json_file:
+                                json.dump(session_data, json_file, indent=4)
+                        except Exception as e:
+                            logger.error(
+                                f"Exception while writing active session to file {file_path} , {e}"
+                            )
+                except InvalidArgument as e:
+                    DataprocSparkSession._active_s8s_session_id = None
+                    raise RuntimeError(
+                        f"Error while creating serverless session: {e}"
+                    ) from None
+                except Exception as e:
+                    DataprocSparkSession._active_s8s_session_id = None
+                    raise RuntimeError(
+                        f"Error while creating serverless session https://console.cloud.google.com/dataproc/interactive/{self._region}/{session_id} : {e}"
+                    ) from None
+
+                logger.debug(
+                    f"Serverless session created: {session_id}, creation time taken: {int(time.time() - s8s_creation_start_time)} seconds"
+                )
+                return self.__create_spark_connect_session_from_s8s(
+                    session_response
+                )
+
+        def _is_s8s_session_active(
+            self, s8s_session_id: str
+        ) -> Optional[sessions.Session]:
+            session_name = f"projects/{self._project_id}/locations/{self._region}/sessions/{s8s_session_id}"
+            get_session_request = GetSessionRequest()
+            get_session_request.name = session_name
+            state = None
+            try:
+                get_session_response = SessionControllerClient(
+                    client_options=self._client_options
+                ).get_session(get_session_request)
+                state = get_session_response.state
+            except Exception as e:
+                logger.debug(f"{s8s_session_id} deleted: {e}")
+                return None
+
+            if state is not None and (
+                state == Session.State.ACTIVE or state == Session.State.CREATING
+            ):
+                return get_session_response
+            return None
+
+        def _get_exiting_active_session(self) -> Optional["SparkSession"]:
+            s8s_session_id = DataprocSparkSession._active_s8s_session_id
+            session_response = self._is_s8s_session_active(s8s_session_id)
+
+            session = DataprocSparkSession.getActiveSession()
+            if session is None:
+                session = DataprocSparkSession._default_session
+
+            if session_response is not None:
+                print(
+                    f"Using existing session: https://console.cloud.google.com/dataproc/interactive/{self._region}/{s8s_session_id}, configuration changes may not be applied."
+                )
+                if session is None:
+                    session = self.__create_spark_connect_session_from_s8s(
+                        session_response
+                    )
+                return session
+            else:
+                logger.info(
+                    f"Session: {s8s_session_id} not active, stopping previous spark session and creating new"
+                )
+                if session is not None:
+                    session.stop()
+
+                return None
+
+        def getOrCreate(self) -> "SparkSession":
+            with DataprocSparkSession._lock:
+                session = self._get_exiting_active_session()
+                if session is None:
+                    session = self.__create()
+                if session:
+                    self.__apply_options(session)
+                return session
+
+        def _get_dataproc_config(self):
+            dataproc_config = Session()
+            if self._dataproc_config:
+                dataproc_config = self._dataproc_config
+                for k, v in self._options.items():
+                    dataproc_config.runtime_config.properties[k] = v
+            elif "DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG" in os.environ:
+                filepath = os.environ[
+                    "DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG"
+                ]
+                try:
+                    with open(filepath, "r") as f:
+                        dataproc_config = Session.wrap(
+                            text_format.Parse(
+                                f.read(), Session.pb(dataproc_config)
+                            )
+                        )
+                except FileNotFoundError:
+                    raise FileNotFoundError(f"File '{filepath}' not found")
+                except ParseError as e:
+                    raise ParseError(f"Error parsing file '{filepath}': {e}")
+            if "COLAB_NOTEBOOK_RUNTIME_ID" in os.environ:
+                dataproc_config.labels["colab-notebook-runtime-id"] = (
+                    os.environ["COLAB_NOTEBOOK_RUNTIME_ID"]
+                )
+            if "COLAB_NOTEBOOK_KERNEL_ID" in os.environ:
+                dataproc_config.labels["colab-notebook-kernel-id"] = os.environ[
+                    "COLAB_NOTEBOOK_KERNEL_ID"
+                ]
+            return dataproc_config
+
+        def _get_session_template(self):
+            from google.cloud.dataproc_v1 import (
+                GetSessionTemplateRequest,
+                SessionTemplateControllerClient,
+            )
+
+            session_template = None
+            if self._dataproc_config and self._dataproc_config.session_template:
+                session_template = self._dataproc_config.session_template
+                get_session_template_request = GetSessionTemplateRequest()
+                get_session_template_request.name = session_template
+                client = SessionTemplateControllerClient(
+                    client_options=self._client_options
+                )
+                try:
+                    session_template = client.get_session_template(
+                        get_session_template_request
+                    )
+                except Exception as e:
+                    logger.error(
+                        f"Failed to get session template {session_template}: {e}"
+                    )
+                    raise
+            return session_template
+
+        def _get_and_validate_version(self, dataproc_config, session_template):
+            trimmed_version = lambda v: ".".join(v.split(".")[:2])
+            version = None
+            if (
+                dataproc_config
+                and dataproc_config.runtime_config
+                and dataproc_config.runtime_config.version
+            ):
+                version = dataproc_config.runtime_config.version
+            elif (
+                session_template
+                and session_template.runtime_config
+                and session_template.runtime_config.version
+            ):
+                version = session_template.runtime_config.version
+
+            if not version:
+                version = "3.0"
+                dataproc_config.runtime_config.version = version
+            elif (
+                trimmed_version(version)
+                not in self._dataproc_runtime_spark_version
+            ):
+                raise ValueError(
+                    f"runtime_config.version {version} is not supported. "
+                    f"Supported versions: {self._dataproc_runtime_spark_version.keys()}"
+                )
+
+            server_version = self._dataproc_runtime_spark_version[
+                trimmed_version(version)
+            ]
+            import importlib.metadata
+
+            dataproc_connect_version = importlib.metadata.version(
+                "dataproc-spark-connect"
+            )
+            client_version = importlib.metadata.version("pyspark")
+            version_message = f"Dataproc Spark Connect: {dataproc_connect_version} (PySpark: {client_version}) Dataproc Session Runtime: {version} (Spark: {server_version})"
+            logger.info(version_message)
+            if trimmed_version(client_version) != trimmed_version(
+                server_version
+            ):
+                logger.warning(
+                    f"client and server on different versions: {version_message}"
+                )
+            return version
+
+        def _get_spark_connect_session(self, dataproc_config, session_template):
+            spark_connect_session = None
+            if dataproc_config and dataproc_config.spark_connect_session:
+                spark_connect_session = dataproc_config.spark_connect_session
+            elif session_template and session_template.spark_connect_session:
+                spark_connect_session = session_template.spark_connect_session
+            return spark_connect_session
+
+        def _get_spark(self, dataproc_config, session_template):
+            spark = None
+            if dataproc_config and dataproc_config.spark:
+                spark = dataproc_config.spark
+            elif session_template and session_template.spark:
+                spark = session_template.spark
+            return spark
+
+        def generate_dataproc_session_id(self):
+            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+            suffix_length = 6
+            random_suffix = "".join(
+                random.choices(
+                    string.ascii_lowercase + string.digits, k=suffix_length
+                )
+            )
+            return f"sc-{timestamp}-{random_suffix}"
+
+    def _repr_html_(self) -> str:
+        if not self._active_s8s_session_id:
+            return """
+            <div>No Active Dataproc Spark Session</div>
+            """
+
+        s8s_session = f"https://console.cloud.google.com/dataproc/interactive/{self._region}/{self._active_s8s_session_id}"
+        ui = f"{s8s_session}/sparkApplications/applications"
+        version = ""
+        return f"""
+        <div>
+            <p><b>Spark Connect</b></p>
+
+            <p><a href="{s8s_session}">Dataproc Session</a></p>
+            <p><a href="{ui}">Spark UI</a></p>
+        </div>
+        """
+
+    def _remove_stoped_session_from_file(self):
+        if "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH" in os.environ:
+            file_path = os.environ[
+                "DATAPROC_SPARK_CONNECT_ACTIVE_SESSION_FILE_PATH"
+            ]
+            try:
+                with open(file_path, "w"):
+                    pass
+            except Exception as e:
+                logger.error(
+                    f"Exception while removing active session in file {file_path} , {e}"
+                )
+
+    def stop(self) -> None:
+        with DataprocSparkSession._lock:
+            if DataprocSparkSession._active_s8s_session_id is not None:
+                from google.cloud.dataproc_v1 import SessionControllerClient
+
+                logger.debug(
+                    f"Terminating serverless session: {DataprocSparkSession._active_s8s_session_id}"
+                )
+                terminate_session_request = TerminateSessionRequest()
+                session_name = f"projects/{DataprocSparkSession._project_id}/locations/{DataprocSparkSession._region}/sessions/{DataprocSparkSession._active_s8s_session_id}"
+                terminate_session_request.name = session_name
+                state = None
+                try:
+                    SessionControllerClient(
+                        client_options=self._client_options
+                    ).terminate_session(terminate_session_request)
+                    get_session_request = GetSessionRequest()
+                    get_session_request.name = session_name
+                    state = Session.State.ACTIVE
+                    while (
+                        state != Session.State.TERMINATING
+                        and state != Session.State.TERMINATED
+                        and state != Session.State.FAILED
+                    ):
+                        session = SessionControllerClient(
+                            client_options=self._client_options
+                        ).get_session(get_session_request)
+                        state = session.state
+                        sleep(1)
+                except NotFound:
+                    logger.debug(
+                        f"Session {DataprocSparkSession._active_s8s_session_id} already deleted"
+                    )
+                except FailedPrecondition:
+                    logger.debug(
+                        f"Session {DataprocSparkSession._active_s8s_session_id} already terminated manually or terminated automatically through session ttl limits"
+                    )
+                if state is not None and state == Session.State.FAILED:
+                    raise RuntimeError("Serverless session termination failed")
+
+                self._remove_stoped_session_from_file()
+                DataprocSparkSession._active_s8s_session_uuid = None
+                DataprocSparkSession._active_s8s_session_id = None
+                DataprocSparkSession._project_id = None
+                DataprocSparkSession._region = None
+                DataprocSparkSession._client_options = None
+
+            self.client.close()
+            if self is DataprocSparkSession._default_session:
+                DataprocSparkSession._default_session = None
+            if self is getattr(
+                DataprocSparkSession._active_session, "session", None
+            ):
+                DataprocSparkSession._active_session.session = None