dataproc-spark-connect 0.1.0__tar.gz → 0.2.0__tar.gz

dataproc_spark_connect-0.2.0/PKG-INFO

@@ -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.1.0 → dataproc_spark_connect-0.2.0}/README.md

@@ -6,15 +6,16 @@ Spark cluster using the Spark Connect protocol without requiring additional step
 
 ## Install
 
-```
-pip install dataproc_spark_connect
-```
+.. code-block:: console
+
+      pip install dataproc_spark_connect
 
 ## Uninstall
 
-```
-pip uninstall dataproc_spark_connect
-```
+.. 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).
@@ -28,33 +29,37 @@ If you are running the client outside of Google Cloud, you must set following en
 ## Usage
 
 1. Install the latest version of Dataproc Python client and Dataproc Spark Connect modules:
-```
-pip install google_cloud_dataproc --force-reinstall
-pip install dataproc_spark_connect --force-reinstall
-```
+
+      .. 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:
-```python
-from google.cloud.dataproc_spark_connect import DataprocSparkSession
 
-```
+      .. 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`:
-      ```python
-      spark = DataprocSparkSession.builder.getOrCreate()
-      ```
+
+      .. code-block:: python
+
+            spark = DataprocSparkSession.builder.getOrCreate()
 
    2. Start a Spark session with the following code instead of using a config file:
-      ```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()
-      ```
+
+      .. 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).
@@ -62,29 +67,32 @@ 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.
 
-```
-pip install -r requirements.txt
-```
+      .. code-block:: console
+
+            pip install -r requirements.txt
+
 2. Build the code.
 
-```
-python setup.py sdist bdist_wheel
-```
+      .. 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
 
-2. Copy the generated `.whl` file to Cloud Storage. Use the version specified in the `setup.py` file.
+            VERSION=<version> gsutil cp dist/dataproc_spark_connect-${VERSION}-py2.py3-none-any.whl gs://<your_bucket_name>
 
-```
-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.
 
-3. 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
-```
+            %%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/dataproc_spark_connect.egg-info/PKG-INFO

@@ -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/dataproc_spark_connect.egg-info/requires.txt

@@ -0,0 +1,7 @@
+google-api-core>=2.19.1
+google-cloud-dataproc>=5.15.1
+wheel
+websockets
+pyspark>=3.5
+pandas
+pyarrow

{dataproc_spark_connect-0.1.0 → dataproc_spark_connect-0.2.0}/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(

{dataproc_spark_connect-0.1.0 → dataproc_spark_connect-0.2.0}/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 → dataproc_spark_connect-0.2.0}/setup.py

@@ -12,13 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from setuptools import find_namespace_packages, setup
+from pathlib import Path
+
+this_directory = Path(__file__).parent
+long_description = (this_directory / "README.md").read_text()
+
 
 setup(
     name="dataproc-spark-connect",
-    version="0.1.0",
+    version="0.2.0",
     description="Dataproc client library for Spark Connect",
+    long_description=long_description,
+    author="Google LLC",
+    url="https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python",
+    license="Apache 2.0",
     packages=find_namespace_packages(include=["google.*"]),
     install_requires=[
+        "google-api-core>=2.19.1",
+        "google-cloud-dataproc>=5.15.1",
         "wheel",
         "websockets",
         "pyspark>=3.5",

dataproc_spark_connect-0.1.0/PKG-INFO

@@ -1,10 +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/dataproc_spark_connect.egg-info/PKG-INFO

@@ -1,10 +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/dataproc_spark_connect.egg-info/requires.txt

@@ -1,5 +0,0 @@
-wheel
-websockets
-pyspark>=3.5
-pandas
-pyarrow