dataproc-spark-connect 0.8.1__py2.py3-none-any.whl → 0.8.3__py2.py3-none-any.whl

{dataproc_spark_connect-0.8.1.dist-info → dataproc_spark_connect-0.8.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataproc-spark-connect
-Version: 0.8.1
+Version: 0.8.3
 Summary: Dataproc client library for Spark Connect
 Home-page: https://github.com/GoogleCloudDataproc/dataproc-spark-connect-python
 Author: Google LLC

{dataproc_spark_connect-0.8.1.dist-info → dataproc_spark_connect-0.8.3.dist-info}/RECORD

@@ -1,12 +1,12 @@
-dataproc_spark_connect-0.8.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+dataproc_spark_connect-0.8.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 google/cloud/dataproc_spark_connect/__init__.py,sha256=dIqHNWVWWrSuRf26x11kX5e9yMKSHCtmI_GBj1-FDdE,1101
 google/cloud/dataproc_spark_connect/exceptions.py,sha256=WF-qdzgdofRwILCriIkjjsmjObZfF0P3Ecg4lv-Hmec,968
 google/cloud/dataproc_spark_connect/pypi_artifacts.py,sha256=gd-VMwiVP-EJuPp9Vf9Shx8pqps3oSKp0hBcSSZQS-A,1575
-google/cloud/dataproc_spark_connect/session.py,sha256=o7N68KddasHkANrNPxsq4Uij711eAaNz4oKa7jGdsaY,32017
+google/cloud/dataproc_spark_connect/session.py,sha256=ZWoW9-otaCJnttPt7h9W3pmhHpdbQsAOl8ypOX3fVbo,33556
 google/cloud/dataproc_spark_connect/client/__init__.py,sha256=6hCNSsgYlie6GuVpc5gjFsPnyeMTScTpXSPYqp1fplY,615
 google/cloud/dataproc_spark_connect/client/core.py,sha256=m3oXTKBm3sBy6jhDu9GRecrxLb5CdEM53SgMlnJb6ag,4616
 google/cloud/dataproc_spark_connect/client/proxy.py,sha256=qUZXvVY1yn934vE6nlO495XUZ53AUx9O74a9ozkGI9U,8976
-dataproc_spark_connect-0.8.1.dist-info/METADATA,sha256=vrvJk30s46CeD-1AZC05zaUtX4ENurK8Ee90Bsb4XEQ,3465
-dataproc_spark_connect-0.8.1.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-dataproc_spark_connect-0.8.1.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
-dataproc_spark_connect-0.8.1.dist-info/RECORD,,
+dataproc_spark_connect-0.8.3.dist-info/METADATA,sha256=croGipnWGtSrd2NLyMCHrcVagYCk9yJ6cEOqCEAm-Qc,3465
+dataproc_spark_connect-0.8.3.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+dataproc_spark_connect-0.8.3.dist-info/top_level.txt,sha256=_1QvSJIhFAGfxb79D6DhB7SUw2X6T4rwnz_LLrbcD3c,7
+dataproc_spark_connect-0.8.3.dist-info/RECORD,,

google/cloud/dataproc_spark_connect/session.py

@@ -22,9 +22,10 @@ import re
 import string
 import threading
 import time
-from typing import Any, cast, ClassVar, Dict, Optional, Union
 import uuid
 import tqdm
+from types import MethodType
+from typing import Any, cast, ClassVar, Dict, Optional, Union
 
 from google.api_core import retry
 from google.api_core.client_options import ClientOptions
@@ -234,6 +235,9 @@ class DataprocSparkSession(SparkSession):
                     target=create_session_pbar
                 )
 
+                # Activate Spark Connect mode for Spark client
+                os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
+
                 try:
                     if (
                         os.getenv(
@@ -538,19 +542,57 @@ class DataprocSparkSession(SparkSession):
 
         super().__init__(connection, user_id)
 
-        base_method = self.client._execute_plan_request_with_metadata
+        execute_plan_request_base_method = (
+            self.client._execute_plan_request_with_metadata
+        )
+        execute_base_method = self.client._execute
+        execute_and_fetch_as_iterator_base_method = (
+            self.client._execute_and_fetch_as_iterator
+        )
 
-        def wrapped_method(*args, **kwargs):
-            req = base_method(*args, **kwargs)
+        def execute_plan_request_wrapped_method(*args, **kwargs):
+            req = execute_plan_request_base_method(*args, **kwargs)
             if not req.operation_id:
                 req.operation_id = str(uuid.uuid4())
                 logger.debug(
                     f"No operation_id found. Setting operation_id: {req.operation_id}"
                 )
-            self._display_operation_link(req.operation_id)
             return req
 
-        self.client._execute_plan_request_with_metadata = wrapped_method
+        self.client._execute_plan_request_with_metadata = (
+            execute_plan_request_wrapped_method
+        )
+
+        def execute_wrapped_method(client_self, req, *args, **kwargs):
+            if not self._sql_lazy_transformation(req):
+                self._display_operation_link(req.operation_id)
+            execute_base_method(req, *args, **kwargs)
+
+        self.client._execute = MethodType(execute_wrapped_method, self.client)
+
+        def execute_and_fetch_as_iterator_wrapped_method(
+            client_self, req, *args, **kwargs
+        ):
+            if not self._sql_lazy_transformation(req):
+                self._display_operation_link(req.operation_id)
+            return execute_and_fetch_as_iterator_base_method(
+                req, *args, **kwargs
+            )
+
+        self.client._execute_and_fetch_as_iterator = MethodType(
+            execute_and_fetch_as_iterator_wrapped_method, self.client
+        )
+
+    @staticmethod
+    def _sql_lazy_transformation(req):
+        # Select SQL command
+        if req.plan and req.plan.command and req.plan.command.sql_command:
+            return (
+                "select"
+                in req.plan.command.sql_command.sql.strip().lower().split()
+            )
+
+        return False
 
     def _repr_html_(self) -> str:
         if not self._active_s8s_session_id:
@@ -593,7 +635,7 @@ class DataprocSparkSession(SparkSession):
                 return
             html_element = f"""
               <div>
-                  <p><a href="{url}">Spark UI</a> (Operation: {operation_id})</p>
+                  <p><a href="{url}">Spark Query</a> (Operation: {operation_id})</p>
               </div>
               """
             display(HTML(html_element))