demo_uc_setup 0.1.0__py3-none-any.whl

demo_uc_setup/common.py

@@ -0,0 +1,49 @@
+import logging
+from typing import TypeVar, Generic
+from databricks.sdk import WorkspaceClient
+
+from demo_uc_setup.config import Config
+
+# Example of a type variable bound to our Config class
+T = TypeVar("T", bound=Config)
+
+class Task(Generic[T]):
+    """
+    A reusable Task base class that works both locally and in Databricks notebooks.
+    When running locally, requires databricks_host and databricks_token.
+    When running in a notebook, these parameters are optional.
+    """
+
+    def __init__(self, config_class: type[T]):
+        # Instantiate the typed configuration
+        self.config: T = config_class()
+        # Setup a basic logger
+        self.logger = logging.getLogger(self.__class__.__name__)
+        logging.basicConfig(level=logging.INFO)
+        
+        # Create a Databricks workspace client with or without credentials
+        if self.config.databricks_host and self.config.databricks_token:
+            # Local execution with credentials
+            self.workspace_client = WorkspaceClient(
+                host=self.config.databricks_host,
+                token=self.config.databricks_token
+            )
+        else:
+            # Notebook execution - no credentials needed
+            self.workspace_client = WorkspaceClient()
+
+    @classmethod
+    def entrypoint(cls, *args, **kwargs):
+        """
+        Creates an instance of the task and runs it. If you
+        want a consistent run pattern, place it here.
+        """
+        instance = cls(*args, **kwargs)
+        instance.run()
+
+    def run(self):
+        """
+        The main entrypoint for the task's execution.
+        Override this in subclasses to implement custom logic.
+        """
+        self.logger.info("Base Task run method. Override in subclasses.") 

demo_uc_setup/config.py

@@ -0,0 +1,23 @@
+from pydantic_settings import BaseSettings
+from typing import Optional
+
+class Config(BaseSettings):
+    """
+    Configuration class using Pydantic BaseSettings.
+    By default, each field can be overridden by environment
+    variables matching the field name (in uppercase).
+    For example, DATABRICKS_HOST, DATABRICKS_TOKEN, etc.
+    """
+
+    # Databricks connection settings - optional for notebook execution
+    databricks_host: Optional[str] = None
+    databricks_token: Optional[str] = None
+
+    # Default names for Unity Catalog demo objects
+    demo_catalog_name: str = "demo_catalog"
+    demo_schemas: list[str] = ["demo_schema_1", "demo_schema_2"]  # List of schemas
+    demo_volume_name: str = "demo_volume"  # This could also be a list if needed
+
+    class Config:
+        env_file = ".env"  # or any custom file, if desired
+        env_file_encoding = "utf-8" 

demo_uc_setup/unity_catalog_setup.py

@@ -0,0 +1,82 @@
+from typing import Type
+from databricks.sdk.service import catalog
+from demo_uc_setup.common import Task, T
+from demo_uc_setup.config import Config
+
+class UnityCatalogSetupTask(Task[Config]):
+    """
+    A task to ensure catalogs, schemas, and volumes exist
+    in the Databricks workspace. Uses typed config for
+    resource names, credentials, etc.
+    """
+
+    def __init__(self, config_class: Type[T] = Config):
+        super().__init__(config_class)
+
+    def run(self):
+        self.logger.info("Starting Unity Catalog setup...")
+
+        # 1) Ensure the catalog exists
+        catalog_name = self.config.demo_catalog_name
+        self.logger.info(f"Ensuring catalog '{catalog_name}'")
+        try:
+            self.workspace_client.catalogs.get(name=catalog_name)
+            self.logger.info(f"Catalog '{catalog_name}' already exists.")
+        except Exception:
+            self.logger.info(f"Catalog '{catalog_name}' not found; creating it.")
+            self.workspace_client.catalogs.create(
+                name=catalog_name,
+                comment="Demo Catalog for Databricks demos"
+            )
+
+        # 2) Ensure all schemas exist and create volumes within each schema
+        for schema_name in self.config.demo_schemas:
+            # Create schema
+            self.logger.info(f"Ensuring schema '{catalog_name}.{schema_name}'")
+            try:
+                self.workspace_client.schemas.get(
+                    name=schema_name,
+                    catalog_name=catalog_name
+                )
+                self.logger.info(f"Schema '{catalog_name}.{schema_name}' already exists.")
+            except Exception:
+                try:
+                    self.logger.info(f"Schema '{catalog_name}.{schema_name}' not found; creating it.")
+                    self.workspace_client.schemas.create(
+                        name=schema_name,
+                        catalog_name=catalog_name,
+                        comment=f"Demo Schema {schema_name} for Databricks demos"
+                    )
+                except Exception as e:
+                    if "already exists" in str(e):
+                        self.logger.info(f"Schema '{catalog_name}.{schema_name}' already exists (caught during creation).")
+                    else:
+                        raise e
+
+            # Create volume within this schema
+            volume_name = self.config.demo_volume_name
+            self.logger.info(f"Ensuring volume '{catalog_name}.{schema_name}.{volume_name}'")
+            try:
+                self.workspace_client.volumes.get(
+                    name=volume_name,
+                    catalog_name=catalog_name,
+                    schema_name=schema_name
+                )
+                self.logger.info(f"Volume '{catalog_name}.{schema_name}.{volume_name}' already exists.")
+            except Exception:
+                try:
+                    self.logger.info(f"Volume '{catalog_name}.{schema_name}.{volume_name}' not found; creating it.")
+                    self.workspace_client.volumes.create(
+                        name=volume_name,
+                        catalog_name=catalog_name,
+                        schema_name=schema_name,
+                        volume_type=catalog.VolumeType.MANAGED,
+                        comment=f"Demo Volume for schema {schema_name}"
+                    )
+                except Exception as e:
+                    if "already exists" in str(e):
+                        self.logger.info(f"Volume '{catalog_name}.{schema_name}.{volume_name}' already exists (caught during creation).")
+                    else:
+                        raise e
+
+        self.logger.info("Unity Catalog setup complete!") 

demo_uc_setup/unity_catalog_teardown.py

@@ -0,0 +1,28 @@
+from typing import Type
+from databricks.sdk import WorkspaceClient
+from demo_uc_setup.common import Task, T
+from demo_uc_setup.config import Config
+
+class UnityCatalogTeardownTask(Task[Config]):
+    """
+    A task to delete (teardown) the Unity Catalog resources in
+    the configured Databricks workspace. Uses typed config for
+    resource names, credentials, etc.
+    """
+
+    def __init__(self, config_class: Type[T] = Config):
+        super().__init__(config_class)
+
+    def run(self):
+        self.logger.info("Starting teardown of Unity Catalog resources...")
+
+        catalog_name = self.config.demo_catalog_name
+        self.logger.info(f"Deleting catalog '{catalog_name}' and its dependencies (force=True).")
+
+        try:
+            self.workspace_client.catalogs.delete(name=catalog_name, force=True)
+            self.logger.info(f"Catalog '{catalog_name}' (and its contents) successfully deleted.")
+        except Exception as e:
+            self.logger.error(f"Failed to delete catalog '{catalog_name}'. Reason: {e}")
+
+        self.logger.info("Unity Catalog teardown complete!") 

demo_uc_setup-0.1.0.dist-info/METADATA

@@ -0,0 +1,105 @@
+Metadata-Version: 2.4
+Name: demo_uc_setup
+Version: 0.1.0
+Summary: A reusable task-based framework for managing Unity Catalog in Databricks
+Requires-Python: >=3.11
+Requires-Dist: databricks-connect>=16.1.1
+Requires-Dist: databricks-sdk>=0.44.1
+Requires-Dist: pydantic-settings>=2.8.0
+Requires-Dist: pydantic>=2.10.6
+Description-Content-Type: text/markdown
+
+# Databricks Unity Catalog Setup Demo
+
+A Python package that demonstrates automated setup and teardown of Databricks Unity Catalog resources using the Databricks SDK. This package provides a reusable framework for managing Unity Catalog resources programmatically, both from local environments and within Databricks notebooks.
+
+## Features
+
+- Automated creation of Unity Catalog resources:
+  - Catalogs
+  - Schemas
+  - Volumes
+- Configurable resource naming via environment variables
+- Support for both local execution and Databricks notebook execution
+- Type-safe configuration management using Pydantic
+- Clean teardown functionality
+
+## Prerequisites
+
+- Python 3.8+
+- A Databricks workspace with Unity Catalog enabled
+- Appropriate permissions to create/manage Unity Catalog resources
+
+## Installation
+
+```bash
+pip install demo-uc-setup
+```
+
+## Configuration
+
+The package uses environment variables for configuration. You can set these either in your environment or in a `.env` file:
+
+```env
+# Required for local execution (optional in Databricks notebooks)
+DATABRICKS_HOST=your-workspace-url
+DATABRICKS_TOKEN=your-pat-token
+
+# Optional - override default resource names
+DEMO_CATALOG_NAME=custom_catalog_name
+DEMO_SCHEMAS=["schema1", "schema2"]
+DEMO_VOLUME_NAME=custom_volume_name
+```
+
+## Usage
+
+### Local Execution
+
+```python
+from demo_uc_setup.unity_catalog_setup import UnityCatalogSetupTask
+from demo_uc_setup.unity_catalog_teardown import UnityCatalogTeardownTask
+
+# Setup Unity Catalog resources
+UnityCatalogSetupTask.entrypoint()
+
+# Teardown Unity Catalog resources
+UnityCatalogTeardownTask.entrypoint()
+```
+
+### Databricks Notebook Execution
+
+```python
+%pip install demo-uc-setup
+
+from demo_uc_setup.unity_catalog_setup import UnityCatalogSetupTask
+UnityCatalogSetupTask.entrypoint()
+```
+
+## Default Resource Names
+
+If not overridden via environment variables, the package will create:
+- A catalog named `demo_catalog`
+- Two schemas: `demo_schema_1` and `demo_schema_2`
+- A volume named `demo_volume` in each schema
+
+## Extending the Framework
+
+The package provides a reusable `Task` base class that can be extended for custom Unity Catalog operations:
+
+```python
+from demo_uc_setup.common import Task
+from demo_uc_setup.config import Config
+
+class CustomTask(Task[Config]):
+    def run(self):
+        self.logger.info("Starting custom task...")
+        # Your custom logic here
+```
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

demo_uc_setup-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+demo_uc_setup/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+demo_uc_setup/common.py,sha256=r5v-nUh7EoGsR4ecQi-pA9NJVRyJlRgOtc_9JxQWjxs,1754
+demo_uc_setup/config.py,sha256=Jwl5anXQ6Ceq-hRyjZ9gUtBaC3iqqJXZCkLjNYjny1o,873
+demo_uc_setup/unity_catalog_setup.py,sha256=3jfIlS9_nxJpXcoGbXJ2Pj4EefqAlpo-GTV8MShoa1w,3680
+demo_uc_setup/unity_catalog_teardown.py,sha256=jHcT4iWeyUIlFDhvUGSfTW2RAYL7d-0sW4VU_uLuRd4,1095
+demo_uc_setup-0.1.0.dist-info/METADATA,sha256=aVkYiBtNuljFYng3CL2JPMLbzkBupHhFYMRiEBDLRjc,2934
+demo_uc_setup-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+demo_uc_setup-0.1.0.dist-info/RECORD,,

demo_uc_setup-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any