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

dataproc_spark_connect-0.2.0.dist-info/METADATA

@@ -0,0 +1,114 @@
+Metadata-Version: 2.1
+Name: dataproc-spark-connect
+Version: 0.2.0
+Summary: Dataproc client library for Spark Connect
+Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
+Author: Google LLC
+License: Apache 2.0
+License-File: LICENSE
+Requires-Dist: google-api-core>=2.19.1
+Requires-Dist: google-cloud-dataproc>=5.15.1
+Requires-Dist: wheel
+Requires-Dist: websockets
+Requires-Dist: pyspark>=3.5
+Requires-Dist: pandas
+Requires-Dist: pyarrow
+
+# Dataproc Spark Connect Client
+
+A wrapper of the Apache [Spark Connect](https://spark.apache.org/spark-connect/) client with
+additional functionalities that allow applications to communicate with a remote Dataproc
+Spark cluster using the Spark Connect protocol without requiring additional steps.
+
+## Install
+
+.. code-block:: console
+
+      pip install dataproc_spark_connect
+
+## Uninstall
+
+.. code-block:: console
+
+      pip uninstall dataproc_spark_connect
+
+
+## Setup
+This client requires permissions to manage [Dataproc sessions and session templates](https://cloud.google.com/dataproc-serverless/docs/concepts/iam).
+If you are running the client outside of Google Cloud, you must set following environment variables:
+
+* GOOGLE_CLOUD_PROJECT - The Google Cloud project you use to run Spark workloads
+* GOOGLE_CLOUD_REGION - The Compute Engine [region](https://cloud.google.com/compute/docs/regions-zones#available) where you run the Spark workload.
+* GOOGLE_APPLICATION_CREDENTIALS - Your [Application Credentials](https://cloud.google.com/docs/authentication/provide-credentials-adc)
+* DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG (Optional) - The config location, such as `tests/integration/resources/session.textproto`
+
+## Usage
+
+1. Install the latest version of Dataproc Python client and Dataproc Spark Connect modules:
+
+      .. code-block:: console
+
+            pip install google_cloud_dataproc --force-reinstall
+            pip install dataproc_spark_connect --force-reinstall
+
+2. Add the required import into your PySpark application or notebook:
+
+      .. code-block:: python
+
+            from google.cloud.dataproc_spark_connect import DataprocSparkSession
+
+3. There are two ways to create a spark session,
+
+   1. Start a Spark session using properties defined in `DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG`:
+
+      .. code-block:: python
+
+            spark = DataprocSparkSession.builder.getOrCreate()
+
+   2. Start a Spark session with the following code instead of using a config file:
+
+      .. code-block:: python
+
+            from google.cloud.dataproc_v1 import SparkConnectConfig
+            from google.cloud.dataproc_v1 import Session
+            dataproc_config = Session()
+            dataproc_config.spark_connect_session = SparkConnectConfig()
+            dataproc_config.environment_config.execution_config.subnetwork_uri = "<subnet>"
+            dataproc_config.runtime_config.version = '3.0'
+            spark = DataprocSparkSession.builder.dataprocConfig(dataproc_config).getOrCreate()
+
+## Billing
+As this client runs the spark workload on Dataproc, your project will be billed as per [Dataproc Serverless Pricing](https://cloud.google.com/dataproc-serverless/pricing).
+This will happen even if you are running the client from a non-GCE instance.
+
+## Contributing
+### Building and Deploying SDK
+
+1. Install the requirements in virtual environment.
+
+      .. code-block:: console
+
+            pip install -r requirements.txt
+
+2. Build the code.
+
+      .. code-block:: console
+
+            python setup.py sdist bdist_wheel
+
+
+3. Copy the generated `.whl` file to Cloud Storage. Use the version specified in the `setup.py` file.
+
+      .. code-block:: console
+
+            VERSION=<version> gsutil cp dist/dataproc_spark_connect-${VERSION}-py2.py3-none-any.whl gs://<your_bucket_name>
+
+4. Download the new SDK on Vertex, then uninstall the old version and install the new one.
+
+      .. code-block:: console
+
+            %%bash
+            export VERSION=<version>
+            gsutil cp gs://<your_bucket_name>/dataproc_spark_connect-${VERSION}-py2.py3-none-any.whl .
+            yes | pip uninstall dataproc_spark_connect
+            pip install dataproc_spark_connect-${VERSION}-py2.py3-none-any.whl

dataproc_spark_connect-0.2.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=A42Wo87VSunG0D3sB-biWyNvU33WhI92mmrJbXI1oNo,23017
+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=ScrbaGsEvqi8wp4ngfD-T9K9mFHXBkVMZkTSr7mdNBs,8926
+dataproc_spark_connect-0.2.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-0.2.0.dist-info/METADATA,sha256=UivMTIfzkp6fzGHG4hiXPUAsRP9P7VBQMKJdEcjmowk,4200
+dataproc_spark_connect-0.2.0.dist-info/WHEEL,sha256=OpXWERl2xLPRHTvd2ZXo_iluPEQd8uSbYkJ53NAER_Y,109
+dataproc_spark_connect-0.2.0.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-0.2.0.dist-info/RECORD,,

{dataproc_spark_connect-0.1.0.dist-info → dataproc_spark_connect-0.2.0.dist-info}/WHEEL

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

google/cloud/dataproc_spark_connect/client/proxy.py

@@ -43,13 +43,22 @@ class bridged_socket(object):
         self._conn = websocket_conn
 
     def recv(self, buff_size):
-        msg = self._conn.recv()
+        # N.B. The websockets [recv method](https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv)
+        # does not support the buff_size parameter, but it does add a `timeout` keyword parameter not supported by normal
+        # socket objects.
+        #
+        # We set that timeout to 60 seconds to prevent any scenarios where we wind up stuck waiting for a message from a websocket connection
+        # that never comes.
+        msg = self._conn.recv(timeout=60)
         return bytes.fromhex(msg)
 
     def send(self, msg_bytes):
         msg = bytes.hex(msg_bytes)
         self._conn.send(msg)
 
+    def close(self):
+        return self._conn.close()
+
 
 def connect_tcp_bridge(hostname):
     """Create a socket-like connection to the given hostname using websocket.
@@ -93,12 +102,51 @@ def forward_bytes(name, from_sock, to_sock):
             bs = from_sock.recv(1024)
             if not bs:
                 return
-            to_sock.send(bs)
+            while bs:
+                try:
+                    to_sock.send(bs)
+                    bs = None
+                except TimeoutError:
+                    # On timeouts during a send, we retry just the send
+                    # to make sure we don't lose any bytes.
+                    pass
+        except TimeoutError:
+            # On timeouts during a receive, we retry the entire flow.
+            pass
         except Exception as ex:
             logger.debug(f"[{name}] Exception forwarding bytes: {ex}")
+            to_sock.close()
             return
 
 
+def connect_sockets(conn_number, from_sock, to_sock):
+    """Create a connection between the two given ports.
+
+    This method continuously streams bytes in both directions between the
+    given `from_sock` and `to_sock` socket-like objects.
+
+    The caller is responsible for creating and closing the supplied socekts.
+    """
+    forward_name = f"{conn_number}-forward"
+    t1 = threading.Thread(
+        name=forward_name,
+        target=forward_bytes,
+        args=[forward_name, from_sock, to_sock],
+        daemon=True,
+    )
+    t1.start()
+    backward_name = f"{conn_number}-backward"
+    t2 = threading.Thread(
+        name=backward_name,
+        target=forward_bytes,
+        args=[backward_name, to_sock, from_sock],
+        daemon=True,
+    )
+    t2.start()
+    t1.join()
+    t2.join()
+
+
 def forward_connection(conn_number, conn, addr, target_host):
     """Create a connection to the target and forward `conn` to it.
 
@@ -115,24 +163,7 @@ def forward_connection(conn_number, conn, addr, target_host):
     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()
+            connect_sockets(conn_number, conn, backend_socket)
 
 
 class DataprocSessionProxy(object):
@@ -179,6 +210,14 @@ class DataprocSessionProxy(object):
             s.release()
             while not self._killed:
                 conn, addr = frontend_socket.accept()
+                # Set a timeout on how long we will allow send/recv calls to block
+                #
+                # The code that reads and writes to this connection will retry
+                # on timeouts, so this is a safe change.
+                #
+                # The chosen timeout is a very short one because it allows us
+                # to more quickly detect when a connection has been closed.
+                conn.settimeout(1)
                 logger.debug(f"Accepted a connection from {addr}...")
                 self._conn_number += 1
                 threading.Thread(

google/cloud/dataproc_spark_connect/session.py

@@ -196,13 +196,13 @@ class DataprocSparkSession(SparkSession):
                 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}"
+                dataproc_config.name = f"projects/{self._project_id}/locations/{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}"
+                    f"projects/{self._project_id}/locations/{self._region}"
                 )
 
                 logger.debug("Creating serverless session")

dataproc_spark_connect-0.1.0.dist-info/METADATA

@@ -1,11 +0,0 @@
-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

@@ -1,10 +0,0 @@
-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,,