ygg 0.1.65__tar.gz → 0.1.66__tar.gz

{ygg-0.1.65 → ygg-0.1.66}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.65
+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.65 → ygg-0.1.66}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.65"
+version = "0.1.66"
 description = "Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }

{ygg-0.1.65 → ygg-0.1.66}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.65
+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.65 → ygg-0.1.66}/src/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:

{ygg-0.1.65 → ygg-0.1.66}/src/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
 

{ygg-0.1.65 → ygg-0.1.66}/src/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:
@@ -291,7 +294,7 @@ class SQLWarehouse(WorkspaceService):
                 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)

ygg-0.1.66/src/yggdrasil/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.66"

ygg-0.1.65/src/yggdrasil/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.65"