ygg 0.1.35__py3-none-any.whl → 0.1.38__py3-none-any.whl

{ygg-0.1.35.dist-info → ygg-0.1.38.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.35
+Version: 0.1.38
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.35.dist-info → ygg-0.1.38.dist-info}/RECORD

@@ -1,17 +1,17 @@
-ygg-0.1.35.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.38.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=PfH7Xwt6uue6oqe6S5V8NhDJcVQClkKrBE1KXhdelZc,117
-yggdrasil/version.py,sha256=dvUFqQgabIeithBNW7swqwxg6R59T0-i289dJYgomYQ,22
+yggdrasil/version.py,sha256=pPo3-5ffVm5NbQms9PN259Sjq7R-Q9UY4HYJ8M2iom4,22
 yggdrasil/databricks/__init__.py,sha256=skctY2c8W-hI81upx9F_PWRe5ishL3hrdiTuizgDjdw,152
 yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
-yggdrasil/databricks/compute/cluster.py,sha256=yx3xEJ5Vgg-IPeyYxxdf74_9DSBPHvg2-FYs8ptJrl0,40404
+yggdrasil/databricks/compute/cluster.py,sha256=mnNzjCx7X3iK22oZ7K3pqot0AXq9JTdg97kT61j2_UU,40729
 yggdrasil/databricks/compute/execution_context.py,sha256=nxrNXoarq_JAB-Cpj0udHhq2jx-DmMbRWJdAezLrPis,22347
 yggdrasil/databricks/compute/remote.py,sha256=nEN_Fr1Ouul_iKOf4B5QjEGscYAcl7nHjGsl2toRzrU,2874
 yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
 yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
 yggdrasil/databricks/sql/__init__.py,sha256=y1n5yg-drZ8QVZbEgznsRG24kdJSnFis9l2YfYCsaCM,234
-yggdrasil/databricks/sql/engine.py,sha256=kUFBddJJQC0AgDqH0l7GFs7d_Ony5rc8fOv4inLU6Vw,41051
+yggdrasil/databricks/sql/engine.py,sha256=Azx3gKtWOMy3D9I2FhkLmpthZPWAJZ9iZkaDivmt_0s,41002
 yggdrasil/databricks/sql/exceptions.py,sha256=Jqd_gT_VyPL8klJEHYEzpv5eHtmdY43WiQ7HZBaEqSk,53
-yggdrasil/databricks/sql/statement_result.py,sha256=VlHXhTcvTVya_2aJ-uUfUooZF_MqQuOZ8k7g6PBDhOM,17227
+yggdrasil/databricks/sql/statement_result.py,sha256=tik53oB7Eq6nA4Unv8qH6fe9RJJO9HcumI5BAkW3omA,16684
 yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCrMhdA,6412
 yggdrasil/databricks/workspaces/__init__.py,sha256=Ti1I99JTC3koYJaCy8WYvkAox4KdcuMRk8b2rHroWCY,133
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
@@ -55,8 +55,8 @@ yggdrasil/types/cast/registry.py,sha256=_zdFGmUBB7P-e_LIcJlOxMcxAkXoA-UXB6HqLMgT
 yggdrasil/types/cast/spark_cast.py,sha256=_KAsl1DqmKMSfWxqhVE7gosjYdgiL1C5bDQv6eP3HtA,24926
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=BuTiWrdCANZCdD_p2MAytqm74eq-rdRXd-LGojBRrfU,5023
 yggdrasil/types/cast/spark_polars_cast.py,sha256=btmZNHXn2NSt3fUuB4xg7coaE0RezIBdZD92H8NK0Jw,9073
-ygg-0.1.35.dist-info/METADATA,sha256=qSh-Cd0LtjU_CV9Je0E1ns1AnuRbIUPMIpvVJoJtIUA,19204
-ygg-0.1.35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.35.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.35.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.35.dist-info/RECORD,,
+ygg-0.1.38.dist-info/METADATA,sha256=2G6ZcCvdN5TLyP9Z7z3GHQio7xjzBCd-ZEOO2qnL2Yc,19204
+ygg-0.1.38.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ygg-0.1.38.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.38.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.38.dist-info/RECORD,,

yggdrasil/databricks/compute/cluster.py

@@ -310,6 +310,11 @@ class Cluster(WorkspaceService):
             self.details = self.clusters_client().get(cluster_id=self.cluster_id)
         return self._details
 
+    def refresh(self, max_delay: float | None = None):
+        self.details = self.fresh_details(max_delay=max_delay)
+
+        return self
+
     @details.setter
     def details(self, value: "ClusterDetails"):
         """Cache cluster details and update identifiers."""
@@ -322,10 +327,10 @@ class Cluster(WorkspaceService):
     @property
     def state(self):
         """Return the current cluster state."""
-        details = self.fresh_details(max_delay=10)
+        self.refresh()
 
-        if details is not None:
-            return details.state
+        if self._details is not None:
+            return self._details.state
         return State.UNKNOWN
 
     @property
@@ -356,7 +361,7 @@ class Cluster(WorkspaceService):
     def wait_for_status(
         self,
         tick: float = 0.5,
-        timeout: float = 600,
+        timeout: Union[float, dt.timedelta] = 600,
         backoff: int = 2,
         max_sleep_time: float = 15
     ):
@@ -374,6 +379,9 @@ class Cluster(WorkspaceService):
         start = time.time()
         sleep_time = tick
 
+        if isinstance(timeout, dt.timedelta):
+            timeout = timeout.total_seconds()
+
         while self.is_pending:
             time.sleep(sleep_time)
 
@@ -691,7 +699,8 @@ class Cluster(WorkspaceService):
             )
 
             self.wait_for_status()
-            self.details = self.clusters_client().edit_and_wait(**update_details)
+            self.details = self.clusters_client().edit(**update_details)
+            self.wait_for_status()
 
             logger.info(
                 "Updated %s",
@@ -754,7 +763,10 @@ class Cluster(WorkspaceService):
                 return None
 
             return Cluster(
-                workspace=self.workspace, cluster_id=details.cluster_id, _details=details
+                workspace=self.workspace,
+                cluster_id=details.cluster_id,
+                cluster_name=details.cluster_name,
+                _details=details
             )
 
         for cluster in self.list_clusters():
@@ -796,10 +808,8 @@ class Cluster(WorkspaceService):
             logger.info("Starting %s", self)
 
             if wait_timeout:
-                self.details = (
-                    self.clusters_client()
-                    .start_and_wait(cluster_id=self.cluster_id, timeout=wait_timeout)
-                )
+                self.clusters_client().start(cluster_id=self.cluster_id)
+                self.wait_for_status(timeout=wait_timeout.total_seconds())
                 self.wait_installed_libraries(timeout=wait_timeout)
             else:
                 self.clusters_client().start(cluster_id=self.cluster_id)

yggdrasil/databricks/sql/engine.py

@@ -309,7 +309,6 @@ class SQLEngine(WorkspaceService):
             engine=self,
             statement_id=response.statement_id,
             _response=response,
-            _response_refresh_time=time.time(),
             disposition=disposition,
         )
 

yggdrasil/databricks/sql/statement_result.py

@@ -44,6 +44,15 @@ if TYPE_CHECKING:
     from .engine import SQLEngine
 
 
+DONE_STATES = {
+    StatementState.CANCELED, StatementState.CLOSED, StatementState.FAILED,
+    StatementState.SUCCEEDED
+}
+
+FAILED_STATES = {
+    StatementState.FAILED, StatementState.CANCELED
+}
+
 __all__ = [
     "StatementResult"
 ]
@@ -57,7 +66,6 @@ class StatementResult:
     disposition: "Disposition"
 
     _response: Optional[StatementResponse] = dataclasses.field(default=None, repr=False)
-    _response_refresh_time: float = dataclasses.field(default=0, repr=False)
 
     _spark_df: Optional[SparkDataFrame] = dataclasses.field(default=None, repr=False)
     _arrow_table: Optional[pa.Table] = dataclasses.field(default=None, repr=False)
@@ -101,8 +109,30 @@ class StatementResult:
         Returns:
             The current StatementResponse object.
         """
-        if self._response is None and not self.is_spark_sql:
-            self.response = self.workspace.sdk().statement_execution.get_statement(self.statement_id)
+        if self.is_spark_sql:
+            return StatementResponse(
+                statement_id=self.statement_id or "sparksql",
+                status=StatementStatus(
+                    state=StatementState.SUCCEEDED
+                )
+            )
+        elif not self.statement_id:
+            return StatementResponse(
+                statement_id="unknown",
+                status=StatementStatus(
+                    state=StatementState.PENDING
+                )
+            )
+
+        statement_execution = self.workspace.sdk().statement_execution
+
+        if self._response is None:
+            # Initialize
+            self._response = statement_execution.get_statement(self.statement_id)
+        elif self._response.status.state not in DONE_STATES:
+            # Refresh
+            self._response = statement_execution.get_statement(self.statement_id)
+
         return self._response
 
     @response.setter
@@ -113,27 +143,8 @@ class StatementResult:
             value: StatementResponse to cache.
         """
         self._response = value
-        self._response_refresh_time = time.time()
-
         self.statement_id = self._response.statement_id
 
-    def fresh_response(self, delay: float):
-        """Refresh the response if it is older than ``delay`` seconds.
-
-        Args:
-            delay: Minimum age in seconds before refreshing.
-
-        Returns:
-            The refreshed StatementResponse object.
-        """
-        if self.is_spark_sql:
-            return self._response
-
-        if self.statement_id and not self.done and time.time() - self._response_refresh_time > delay:
-            self.response = self.workspace.sdk().statement_execution.get_statement(self.statement_id)
-
-        return self._response
-
     def result_data_at(self, chunk_index: int):
         """Fetch a specific result chunk by index.
 
@@ -166,17 +177,7 @@ class StatementResult:
         Returns:
             A StatementStatus object.
         """
-        if self.persisted:
-            return StatementStatus(
-                state=StatementState.SUCCEEDED
-            )
-
-        if not self.statement_id:
-            return StatementStatus(
-                state=StatementState.PENDING
-            )
-
-        return self.fresh_response(delay=1).status
+        return self.response.status
 
     @property
     def state(self):
@@ -194,8 +195,6 @@ class StatementResult:
         Returns:
             The result manifest or None for Spark SQL results.
         """
-        if self.is_spark_sql:
-            return None
         return self.response.manifest
 
     @property
@@ -214,15 +213,7 @@ class StatementResult:
         Returns:
             True if the statement is done, otherwise False.
         """
-        if self.persisted:
-            return True
-
-        if self._response is None:
-            return False
-
-        return self._response.status.state in [
-            StatementState.CANCELED, StatementState.CLOSED, StatementState.FAILED, StatementState.SUCCEEDED
-        ]
+        return self.state in DONE_STATES
 
     @property
     def failed(self):
@@ -231,13 +222,7 @@ class StatementResult:
         Returns:
             True if the statement failed or was cancelled.
         """
-        if self.persisted:
-            return True
-
-        if self._response is None:
-            return False
-
-        return self._response.status.state in [StatementState.CANCELED, StatementState.FAILED]
+        return self.state in FAILED_STATES
 
     @property
     def persisted(self):

yggdrasil/version.py

@@ -1 +1 @@
-__version__ = "0.1.35"
+__version__ = "0.1.38"