ygg 0.1.64__py3-none-any.whl → 0.1.66__py3-none-any.whl

{ygg-0.1.64.dist-info → ygg-0.1.66.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.64
+Version: 0.1.66
 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.64.dist-info → ygg-0.1.66.dist-info}/RECORD

@@ -1,7 +1,7 @@
-ygg-0.1.64.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.66.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=4-ghPak2S6zfMqmnlxW2GCgPb5s79znpKa2hGEGXcE4,24
 yggdrasil/exceptions.py,sha256=NEpbDFn-8ZRsLiEgJicCwrTHNMWAGtdrTJzosfAeVJo,82
-yggdrasil/version.py,sha256=4MIEvUQSlMm4HeF033QtQTWKTLCJXVFMfXEyriOrkl4,22
+yggdrasil/version.py,sha256=y19u6jx_RZZOmyZU6XLMb8UCSG1X0auTRRU67so04sc,22
 yggdrasil/ai/__init__.py,sha256=YEOVsyuvEOvPaZT8XN9xNysS_WOpHTbKgXgnA8up7x0,52
 yggdrasil/ai/session.py,sha256=10ATAnw8FOCpfIg9sNR4meki_MRckUzKZ9Uft4IXwLA,2515
 yggdrasil/ai/sql_session.py,sha256=n92tQjHUBIey6c3EJProiEEwfAtQm07Dtmei4WXzeG0,10812
@@ -16,11 +16,11 @@ yggdrasil/databricks/compute/remote.py,sha256=sF99i7GXZcC0GiNgO9VO0I26rFbrtnDhK9
 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=PetgRp1jEj5K3TgN09FwNUVjVN8YYuGq0cDIOTqsbns,144
-yggdrasil/databricks/sql/engine.py,sha256=nFWeegs91CtjCLzxgZsJwOlAXNVI1v_lfecuFVfKFFY,49979
+yggdrasil/databricks/sql/engine.py,sha256=hFPQSoliXBXuAMJDC1UZd3Sr1xSKAHg9X_DqE2EVinc,50015
 yggdrasil/databricks/sql/exceptions.py,sha256=srMR3Y9LQm45rkyxfyCgpgcoGtRRvGKWBEoUHf4kxsg,1762
-yggdrasil/databricks/sql/statement_result.py,sha256=01DzFX1bGDIGHj0OW2ngfVVJ1w1KHlZEfAI934E35CU,15549
+yggdrasil/databricks/sql/statement_result.py,sha256=5Q7qZ2IOfFgJmaXrJ4dJgGMR6CnlFb9a0VXMr6j85fU,15601
 yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCrMhdA,6412
-yggdrasil/databricks/sql/warehouse.py,sha256=6rZAQ98avExuhllT-pKrjbIS3RFBMCC2RBVzXouNFo4,18976
+yggdrasil/databricks/sql/warehouse.py,sha256=gDdBvThnCLQ0KjgTWz5k__YQz-K4pDgPp3U5hp6Y3R8,21349
 yggdrasil/databricks/workspaces/__init__.py,sha256=dv2zotoFVhNFlTCdRq6gwf5bEzeZkOZszoNZMs0k59g,114
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
 yggdrasil/databricks/workspaces/io.py,sha256=RdgN5lmEYNF5phPRkRMCVHbUl-t3ZUGkKbzgYSTKpII,37420
@@ -67,8 +67,8 @@ yggdrasil/types/cast/registry.py,sha256=OOqIfbIjPH-a3figvu-zTvEtUDTEWhe2xIl3cCA4
 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.64.dist-info/METADATA,sha256=GXhjqjZQyvPanFP8JLjl3BvZU-pLZB_NM6pR2It7Q74,18506
-ygg-0.1.64.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-ygg-0.1.64.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.64.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.64.dist-info/RECORD,,
+ygg-0.1.66.dist-info/METADATA,sha256=GwVrDtjvLiqAqeGiopK2hv5zK3QJoCHuCZ5q1zgzssM,18506
+ygg-0.1.66.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+ygg-0.1.66.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.66.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.66.dist-info/RECORD,,

yggdrasil/databricks/sql/engine.py

@@ -186,7 +186,8 @@ class SQLEngine(WorkspaceService):
             self._warehouse = wh.find_warehouse(
                 warehouse_id=warehouse_id,
                 warehouse_name=warehouse_name,
-                raise_error=False
+                raise_error=False,
+                find_starter=True
             )
 
             if self._warehouse is None:

yggdrasil/databricks/sql/statement_result.py

@@ -305,14 +305,16 @@ class StatementResult:
             The current StatementResult instance.
         """
         wait = WaitingConfig.check_arg(wait)
-        iteration, start = 0, time.time()
 
-        if not self.done:
-            wait.sleep(iteration=iteration, start=start)
-            iteration += 1
+        if wait.timeout:
+            iteration, start = 0, time.time()
 
-        if raise_error:
-            self.raise_for_status()
+            if not self.done:
+                wait.sleep(iteration=iteration, start=start)
+                iteration += 1
+
+            if raise_error:
+                self.raise_for_status()
 
         return self
 

yggdrasil/databricks/sql/warehouse.py

@@ -18,8 +18,9 @@ try:
         EndpointTags, EndpointTagPair, EndpointInfoWarehouseType,
         GetWarehouseResponse, GetWarehouseResponseWarehouseType,
         Disposition, Format,
-        ExecuteStatementRequestOnWaitTimeout, StatementParameterListItem
-    )
+        ExecuteStatementRequestOnWaitTimeout, StatementParameterListItem,
+        WarehouseAccessControlRequest, WarehousePermissionLevel
+)
 
     _CREATE_ARG_NAMES = {_ for _ in inspect.signature(WarehousesAPI.create).parameters.keys()}
     _EDIT_ARG_NAMES = {_ for _ in inspect.signature(WarehousesAPI.edit).parameters.keys()}
@@ -36,6 +37,7 @@ except ImportError:
     Format = DatabricksDummyClass
     ExecuteStatementRequestOnWaitTimeout = DatabricksDummyClass
     StatementParameterListItem = DatabricksDummyClass
+    WarehouseAccessControlRequest = DatabricksDummyClass
 
 
 __all__ = [
@@ -242,6 +244,7 @@ class SQLWarehouse(WorkspaceService):
         warehouse_id: Optional[str] = None,
         warehouse_name: Optional[str] = None,
         raise_error: bool = True,
+        find_starter: bool = False
     ):
         if warehouse_id:
             if warehouse_id == self.warehouse_id:
@@ -288,10 +291,10 @@ class SQLWarehouse(WorkspaceService):
 
                     return warehouse
 
-                elif warehouse.warehouse_name == starter_warehouse:
+                elif warehouse.warehouse_name == starter_name:
                     starter_warehouse = warehouse
 
-        if starter_warehouse is not None:
+        if find_starter and starter_warehouse is not None:
             return starter_warehouse
 
         if raise_error:
@@ -401,6 +404,7 @@ class SQLWarehouse(WorkspaceService):
             warehouse_id=warehouse_id,
             warehouse_name=name,
             raise_error=False,
+            find_starter=False
         )
 
         if found is not None:
@@ -443,13 +447,17 @@ class SQLWarehouse(WorkspaceService):
 
             info = EndpointInfo(**update_details)
 
-        return SQLWarehouse(
+        created = SQLWarehouse(
             workspace=self.workspace,
             warehouse_id=info.id,
             warehouse_name=info.name,
             _details=info
         )
 
+        created.update_permissions()
+
+        return created
+
     def update(
         self,
         wait: Optional[WaitingConfigArg] = None,
@@ -510,6 +518,76 @@ class SQLWarehouse(WorkspaceService):
 
         return self
 
+    def update_permissions(
+        self,
+        access_control_list: Optional[List[WarehouseAccessControlRequest]] = None,
+        wait: Optional[WaitingConfigArg] = None
+    ):
+        if self.warehouse_id:
+            client = self.warehouse_client()
+
+            access_control_list = self._check_access_control_list(
+                access_control_list=access_control_list
+            )
+
+            if access_control_list:
+                client.update_permissions(
+                    warehouse_id=self.warehouse_id,
+                    access_control_list=access_control_list
+                )
+
+    def default_access_control_list(self, for_all: bool):
+        if for_all:
+            base = [
+                WarehouseAccessControlRequest(
+                    group_name="users",
+                    permission_level=WarehousePermissionLevel.CAN_USE
+                )
+            ]
+        else:
+            base = []
+
+        groups = self.workspace.current_user.groups
+
+        if groups:
+            base.extend(
+                WarehouseAccessControlRequest(
+                    group_name=group.display,
+                    permission_level=WarehousePermissionLevel.CAN_MANAGE
+                )
+                for group in groups
+                if group.display != "users"
+            )
+
+        return base
+
+    def _check_access_control_list(
+        self,
+        access_control_list: Optional[List[WarehouseAccessControlRequest]] = None
+    ):
+        if access_control_list is None:
+            access_control_list = []
+
+        access_control_list.extend(self.default_access_control_list(
+            for_all=self.warehouse_name.startswith("yggdrasil") if self.warehouse_name else False
+        ))
+
+        return access_control_list
+
+    def delete(self):
+        if self.warehouse_id:
+            LOGGER.debug(
+                "Deleting %s",
+                self
+            )
+
+            self.warehouse_client().delete(id=self.warehouse_id)
+
+            LOGGER.info(
+                "Deleted %s",
+                self
+            )
+
     def execute(
         self,
         statement: Optional[str] = None,
@@ -599,4 +677,4 @@ class SQLWarehouse(WorkspaceService):
             execution
         )
 
-        return execution.wait() if wait is not None else execution
+        return execution.wait(wait=wait)

yggdrasil/version.py

@@ -1 +1 @@
-__version__ = "0.1.64"
+__version__ = "0.1.66"