zipher-sdk 0.2.0__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zipher-sdk-0.3.0/PKG-INFO +139 -0
- zipher-sdk-0.3.0/README.md +122 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/setup.py +9 -4
- zipher-sdk-0.3.0/tests/test_acl/__init__.py +0 -0
- zipher-sdk-0.3.0/tests/test_acl/sample_job.py +15 -0
- zipher-sdk-0.3.0/tests/test_acl/test_acl.py +353 -0
- zipher-sdk-0.3.0/zipher/acl/__init__.py +0 -0
- zipher-sdk-0.3.0/zipher/acl/models.py +49 -0
- zipher-sdk-0.3.0/zipher/acl/zipher_installer.py +376 -0
- zipher-sdk-0.3.0/zipher/cli.py +153 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher/client.py +9 -7
- zipher-sdk-0.3.0/zipher/exceptions.py +16 -0
- zipher-sdk-0.3.0/zipher/models.py +26 -0
- zipher-sdk-0.3.0/zipher/utils.py +67 -0
- zipher-sdk-0.3.0/zipher_sdk.egg-info/PKG-INFO +139 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher_sdk.egg-info/SOURCES.txt +9 -1
- zipher-sdk-0.3.0/zipher_sdk.egg-info/entry_points.txt +3 -0
- zipher-sdk-0.3.0/zipher_sdk.egg-info/requires.txt +3 -0
- zipher-sdk-0.2.0/PKG-INFO +0 -75
- zipher-sdk-0.2.0/README.md +0 -58
- zipher-sdk-0.2.0/zipher/conf.py +0 -6
- zipher-sdk-0.2.0/zipher/exceptions.py +0 -4
- zipher-sdk-0.2.0/zipher/models.py +0 -8
- zipher-sdk-0.2.0/zipher_sdk.egg-info/PKG-INFO +0 -75
- zipher-sdk-0.2.0/zipher_sdk.egg-info/requires.txt +0 -3
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/LICENSE +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/setup.cfg +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/tests/__init__.py +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/tests/test_client.py +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher/__init__.py +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher_sdk.egg-info/dependency_links.txt +0 -0
- {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher_sdk.egg-info/top_level.txt +0 -0
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: zipher-sdk
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Python SDK for programmatic access to the Zipher API
|
|
5
|
+
Home-page: UNKNOWN
|
|
6
|
+
Author: Zipher Inc.
|
|
7
|
+
Author-email: yoav@zipher.cloud
|
|
8
|
+
License: UNKNOWN
|
|
9
|
+
Platform: UNKNOWN
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Operating System :: OS Independent
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
# Zipher SDK
|
|
18
|
+
|
|
19
|
+
The Zipher SDK is a Python library for interacting with Zipher's APIs.
|
|
20
|
+
|
|
21
|
+
- [Package Installation](#package-installation)
|
|
22
|
+
- [Granting access to Zipher](#providing-zipher-with-access-to-databricks-workspace)
|
|
23
|
+
- [Usage](#sdk-usage)
|
|
24
|
+
|
|
25
|
+
## Package Installation
|
|
26
|
+
|
|
27
|
+
You can install the Zipher SDK using pip:
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install zipher-sdk
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Providing Zipher with access to Databricks workspace
|
|
34
|
+
|
|
35
|
+
After installing `zipher-sdk` package a cli tool to automatically create all necessary resources for Zipher is available.
|
|
36
|
+
|
|
37
|
+
### Setting up credentials
|
|
38
|
+
|
|
39
|
+
You need to provide credentials that will be used to create all necessary resources and permissions for Zipher.
|
|
40
|
+
|
|
41
|
+
Here are ways to set up credentials:
|
|
42
|
+
* `.databrickscfg` config (cli tool supports profile choice as an argument)
|
|
43
|
+
* `ZIPHER_DATABRICKS_HOST` and `ZIPHER_DATABRICKS_TOKEN` or `ZIPHER_DATABRICKS_CLIENT_ID`, `ZIPHER_DATABRICKS_CLIENT_SECRET` environment variables
|
|
44
|
+
* providing credentials as arguments to cli tool
|
|
45
|
+
|
|
46
|
+
### Cli tool usage examples
|
|
47
|
+
|
|
48
|
+
#### Providing Zipher with access to a list of jobs
|
|
49
|
+
```
|
|
50
|
+
zipher setup --jobs-list 12345678,87654321,12344321,21436587
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
#### Providing Zipher with access to n jobs from the workspace
|
|
54
|
+
```
|
|
55
|
+
zipher setup --max-jobs 50
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
#### Providing Zipher with readonly access to a list of jobs
|
|
59
|
+
```
|
|
60
|
+
zipher setup --readonly --jobs-list 12345678,87654321,12344321,21436587
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
### Full cli tool specification
|
|
64
|
+
```
|
|
65
|
+
usage: zipher setup [-h] [--workspace-host WORKSPACE_HOST] [--access-token ACCESS_TOKEN] [--client-id CLIENT_ID] [--client-secret CLIENT_SECRET] [--profile PROFILE] [--verbose] [--jobs-list JOBS_LIST] [--max-jobs MAX_JOBS]
|
|
66
|
+
[--max-runs MAX_RUNS] [--days-back DAYS_BACK] [--readonly] [--pat] [--skip-approval]
|
|
67
|
+
|
|
68
|
+
options:
|
|
69
|
+
-h, --help show this help message and exit
|
|
70
|
+
--workspace-host WORKSPACE_HOST
|
|
71
|
+
Databricks workspace host URL.
|
|
72
|
+
--access-token ACCESS_TOKEN
|
|
73
|
+
Databricks workspace access token.
|
|
74
|
+
--client-id CLIENT_ID
|
|
75
|
+
Databricks workspace OAuth client id.
|
|
76
|
+
--client-secret CLIENT_SECRET
|
|
77
|
+
Databricks workspace OAuth client secret.
|
|
78
|
+
--profile PROFILE Profile name from .databrickscfg.
|
|
79
|
+
--verbose Print full error message on fail.
|
|
80
|
+
--jobs-list JOBS_LIST
|
|
81
|
+
Comma-separated list of jobs ids to provide access to.
|
|
82
|
+
--max-jobs MAX_JOBS Maximum number of jobs to consider when iterating over jobs to grant permissions (default: 2000).
|
|
83
|
+
--max-runs MAX_RUNS Maximum number of runs to consider when iterating over runs to grant permissions to relative jobs (default: 2000).
|
|
84
|
+
--days-back DAYS_BACK
|
|
85
|
+
How many days back to fetch relevant job runs for permission updates (default: 7).
|
|
86
|
+
--readonly Provide Zipher with only CAN_VIEW permissions on listed jobs. When not provided will default to CAN_MANAGE permissions.
|
|
87
|
+
--pat Generate Personal Access Token for Zipher instead of default OAuth client creds.
|
|
88
|
+
--skip-approval Skip user input approval.
|
|
89
|
+
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
## SDK Usage
|
|
94
|
+
|
|
95
|
+
Here are some basic examples of how you can use the Zipher SDK to optimize your databricks clusters using Zipher's ML-powered optimization engine:
|
|
96
|
+
|
|
97
|
+
### Update Existing Configuration
|
|
98
|
+
|
|
99
|
+
You can update an existing configuration by initializing a zipher Client and sending a JSON payload to the `update_existing_conf` function. Here's how you can do it:
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
from zipher import Client
|
|
103
|
+
|
|
104
|
+
client = Client(customer_id="my_customer_id") # assuming the zipher API key is stored in ZIPHER_API_KEY environment variable
|
|
105
|
+
|
|
106
|
+
# Your existing cluster config:
|
|
107
|
+
config_payload = {
|
|
108
|
+
"new_cluster": {
|
|
109
|
+
"autoscale": {
|
|
110
|
+
"min_workers": 1,
|
|
111
|
+
"max_workers": 30
|
|
112
|
+
},
|
|
113
|
+
"cluster_name": "my-cluster",
|
|
114
|
+
"spark_version": "10.4.x-scala2.12",
|
|
115
|
+
"spark_conf": {
|
|
116
|
+
"spark.driver.maxResultSize": "4g"
|
|
117
|
+
},
|
|
118
|
+
"aws_attributes": {
|
|
119
|
+
"first_on_demand": 0,
|
|
120
|
+
"availability": "SPOT",
|
|
121
|
+
"zone_id": "auto",
|
|
122
|
+
"spot_bid_price_percent": 100,
|
|
123
|
+
"ebs_volume_count": 0
|
|
124
|
+
},
|
|
125
|
+
"node_type_id": "rd-fleet.2xlarge",
|
|
126
|
+
"driver_node_type_id": "rd-fleet.xlarge",
|
|
127
|
+
"spark_env_vars": {},
|
|
128
|
+
"enable_elastic_disk": "false"
|
|
129
|
+
}
|
|
130
|
+
}
|
|
131
|
+
|
|
132
|
+
# Update configuration
|
|
133
|
+
optimized_cluster = client.update_existing_conf(job_id="my-job-id", existing_conf=config_payload)
|
|
134
|
+
|
|
135
|
+
# Continue with sending the optimized configuration to Databricks via the Databricks python SDK, Airflow operator, etc.
|
|
136
|
+
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
|
|
2
|
+
# Zipher SDK
|
|
3
|
+
|
|
4
|
+
The Zipher SDK is a Python library for interacting with Zipher's APIs.
|
|
5
|
+
|
|
6
|
+
- [Package Installation](#package-installation)
|
|
7
|
+
- [Granting access to Zipher](#providing-zipher-with-access-to-databricks-workspace)
|
|
8
|
+
- [Usage](#sdk-usage)
|
|
9
|
+
|
|
10
|
+
## Package Installation
|
|
11
|
+
|
|
12
|
+
You can install the Zipher SDK using pip:
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
pip install zipher-sdk
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
## Providing Zipher with access to Databricks workspace
|
|
19
|
+
|
|
20
|
+
After installing `zipher-sdk` package a cli tool to automatically create all necessary resources for Zipher is available.
|
|
21
|
+
|
|
22
|
+
### Setting up credentials
|
|
23
|
+
|
|
24
|
+
You need to provide credentials that will be used to create all necessary resources and permissions for Zipher.
|
|
25
|
+
|
|
26
|
+
Here are ways to set up credentials:
|
|
27
|
+
* `.databrickscfg` config (cli tool supports profile choice as an argument)
|
|
28
|
+
* `ZIPHER_DATABRICKS_HOST` and `ZIPHER_DATABRICKS_TOKEN` or `ZIPHER_DATABRICKS_CLIENT_ID`, `ZIPHER_DATABRICKS_CLIENT_SECRET` environment variables
|
|
29
|
+
* providing credentials as arguments to cli tool
|
|
30
|
+
|
|
31
|
+
### Cli tool usage examples
|
|
32
|
+
|
|
33
|
+
#### Providing Zipher with access to a list of jobs
|
|
34
|
+
```
|
|
35
|
+
zipher setup --jobs-list 12345678,87654321,12344321,21436587
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
#### Providing Zipher with access to n jobs from the workspace
|
|
39
|
+
```
|
|
40
|
+
zipher setup --max-jobs 50
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
#### Providing Zipher with readonly access to a list of jobs
|
|
44
|
+
```
|
|
45
|
+
zipher setup --readonly --jobs-list 12345678,87654321,12344321,21436587
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
### Full cli tool specification
|
|
49
|
+
```
|
|
50
|
+
usage: zipher setup [-h] [--workspace-host WORKSPACE_HOST] [--access-token ACCESS_TOKEN] [--client-id CLIENT_ID] [--client-secret CLIENT_SECRET] [--profile PROFILE] [--verbose] [--jobs-list JOBS_LIST] [--max-jobs MAX_JOBS]
|
|
51
|
+
[--max-runs MAX_RUNS] [--days-back DAYS_BACK] [--readonly] [--pat] [--skip-approval]
|
|
52
|
+
|
|
53
|
+
options:
|
|
54
|
+
-h, --help show this help message and exit
|
|
55
|
+
--workspace-host WORKSPACE_HOST
|
|
56
|
+
Databricks workspace host URL.
|
|
57
|
+
--access-token ACCESS_TOKEN
|
|
58
|
+
Databricks workspace access token.
|
|
59
|
+
--client-id CLIENT_ID
|
|
60
|
+
Databricks workspace OAuth client id.
|
|
61
|
+
--client-secret CLIENT_SECRET
|
|
62
|
+
Databricks workspace OAuth client secret.
|
|
63
|
+
--profile PROFILE Profile name from .databrickscfg.
|
|
64
|
+
--verbose Print full error message on fail.
|
|
65
|
+
--jobs-list JOBS_LIST
|
|
66
|
+
Comma-separated list of jobs ids to provide access to.
|
|
67
|
+
--max-jobs MAX_JOBS Maximum number of jobs to consider when iterating over jobs to grant permissions (default: 2000).
|
|
68
|
+
--max-runs MAX_RUNS Maximum number of runs to consider when iterating over runs to grant permissions to relative jobs (default: 2000).
|
|
69
|
+
--days-back DAYS_BACK
|
|
70
|
+
How many days back to fetch relevant job runs for permission updates (default: 7).
|
|
71
|
+
--readonly Provide Zipher with only CAN_VIEW permissions on listed jobs. When not provided will default to CAN_MANAGE permissions.
|
|
72
|
+
--pat Generate Personal Access Token for Zipher instead of default OAuth client creds.
|
|
73
|
+
--skip-approval Skip user input approval.
|
|
74
|
+
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
## SDK Usage
|
|
79
|
+
|
|
80
|
+
Here are some basic examples of how you can use the Zipher SDK to optimize your databricks clusters using Zipher's ML-powered optimization engine:
|
|
81
|
+
|
|
82
|
+
### Update Existing Configuration
|
|
83
|
+
|
|
84
|
+
You can update an existing configuration by initializing a zipher Client and sending a JSON payload to the `update_existing_conf` function. Here's how you can do it:
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
from zipher import Client
|
|
88
|
+
|
|
89
|
+
client = Client(customer_id="my_customer_id") # assuming the zipher API key is stored in ZIPHER_API_KEY environment variable
|
|
90
|
+
|
|
91
|
+
# Your existing cluster config:
|
|
92
|
+
config_payload = {
|
|
93
|
+
"new_cluster": {
|
|
94
|
+
"autoscale": {
|
|
95
|
+
"min_workers": 1,
|
|
96
|
+
"max_workers": 30
|
|
97
|
+
},
|
|
98
|
+
"cluster_name": "my-cluster",
|
|
99
|
+
"spark_version": "10.4.x-scala2.12",
|
|
100
|
+
"spark_conf": {
|
|
101
|
+
"spark.driver.maxResultSize": "4g"
|
|
102
|
+
},
|
|
103
|
+
"aws_attributes": {
|
|
104
|
+
"first_on_demand": 0,
|
|
105
|
+
"availability": "SPOT",
|
|
106
|
+
"zone_id": "auto",
|
|
107
|
+
"spot_bid_price_percent": 100,
|
|
108
|
+
"ebs_volume_count": 0
|
|
109
|
+
},
|
|
110
|
+
"node_type_id": "rd-fleet.2xlarge",
|
|
111
|
+
"driver_node_type_id": "rd-fleet.xlarge",
|
|
112
|
+
"spark_env_vars": {},
|
|
113
|
+
"enable_elastic_disk": "false"
|
|
114
|
+
}
|
|
115
|
+
}
|
|
116
|
+
|
|
117
|
+
# Update configuration
|
|
118
|
+
optimized_cluster = client.update_existing_conf(job_id="my-job-id", existing_conf=config_payload)
|
|
119
|
+
|
|
120
|
+
# Continue with sending the optimized configuration to Databricks via the Databricks python SDK, Airflow operator, etc.
|
|
121
|
+
|
|
122
|
+
```
|
|
@@ -8,7 +8,7 @@ with open('README.md', encoding='utf-8') as f:
|
|
|
8
8
|
|
|
9
9
|
setup(
|
|
10
10
|
name='zipher-sdk',
|
|
11
|
-
version='0.
|
|
11
|
+
version='0.3.0',
|
|
12
12
|
description='Python SDK for programmatic access to the Zipher API',
|
|
13
13
|
long_description=long_description,
|
|
14
14
|
long_description_content_type='text/markdown',
|
|
@@ -16,10 +16,15 @@ setup(
|
|
|
16
16
|
author_email='yoav@zipher.cloud',
|
|
17
17
|
packages=find_packages(),
|
|
18
18
|
install_requires=[
|
|
19
|
-
"pydantic==2.
|
|
20
|
-
"requests==2.
|
|
21
|
-
"databricks-sdk==0.
|
|
19
|
+
"pydantic==2.11.3",
|
|
20
|
+
"requests==2.32.3",
|
|
21
|
+
"databricks-sdk==0.52.0"
|
|
22
22
|
],
|
|
23
|
+
entry_points={
|
|
24
|
+
"console_scripts": [
|
|
25
|
+
"zipher=zipher.cli:main",
|
|
26
|
+
]
|
|
27
|
+
},
|
|
23
28
|
classifiers=[
|
|
24
29
|
'Programming Language :: Python :: 3',
|
|
25
30
|
'License :: OSI Approved :: MIT License',
|
|
File without changes
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from time import sleep
|
|
3
|
+
|
|
4
|
+
print("<<<<<<<< >>>>>>>>> HELLO WORLD <<<<<<<< >>>>>>>>>")
|
|
5
|
+
sleep_time = 20
|
|
6
|
+
|
|
7
|
+
sample_secret_env_var = os.getenv('ZIPHER_API_KEY')
|
|
8
|
+
if sample_secret_env_var and sample_secret_env_var == 'test_key':
|
|
9
|
+
print(f"Successfully retrieved secret from environment variable: {sample_secret_env_var}")
|
|
10
|
+
sleep_time = 90
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
print(f"About to sleep for {sleep_time} seconds")
|
|
14
|
+
sleep(sleep_time)
|
|
15
|
+
print("<<<<<<<< >>>>>>>>> GOODBYE WORLD <<<<<<<< >>>>>>>>>")
|
|
@@ -0,0 +1,353 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import time
|
|
3
|
+
import json
|
|
4
|
+
import base64
|
|
5
|
+
from uuid import uuid4
|
|
6
|
+
from copy import deepcopy
|
|
7
|
+
from unittest import TestCase
|
|
8
|
+
|
|
9
|
+
from databricks.sdk.service import iam
|
|
10
|
+
from databricks.sdk import WorkspaceClient
|
|
11
|
+
from databricks.sdk.service.sql import CreateWarehouseRequestWarehouseType, CreateQueryRequestQuery, \
|
|
12
|
+
WarehouseAccessControlRequest, WarehousePermissionLevel
|
|
13
|
+
from databricks.sdk.service.workspace import ImportFormat
|
|
14
|
+
from databricks.sdk.service.settings import TokenAccessControlRequest, TokenPermissionLevel
|
|
15
|
+
from databricks.sdk.service.iam import ServicePrincipal, AccessControlRequest, PermissionLevel
|
|
16
|
+
from databricks.sdk.service.compute import ClusterSpec, AwsAttributes, InitScriptInfo, WorkspaceStorageInfo, \
|
|
17
|
+
ClusterLogConf, DbfsStorageInfo
|
|
18
|
+
from databricks.sdk.service.jobs import JobPermissions, JobPermissionLevel, \
|
|
19
|
+
JobCluster, RunLifeCycleState, RunResultState, Task, SparkPythonTask, SqlTask, SqlTaskQuery
|
|
20
|
+
|
|
21
|
+
from zipher.acl.zipher_installer import ZipherInstaller
|
|
22
|
+
from zipher.acl.models import InstallerConfig, InstallerUserConfig
|
|
23
|
+
from zipher.utils import get_workspace_client, create_oauth_secret_for_sp
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class TestACL(TestCase):
|
|
27
|
+
@classmethod
|
|
28
|
+
def setUpClass(cls):
|
|
29
|
+
cls.host = os.environ['DATABRICKS_HOST_TEST']
|
|
30
|
+
cls.admin_dbx: WorkspaceClient = WorkspaceClient(host=cls.host, token=os.environ['DATABRICKS_TOKEN_TEST'])
|
|
31
|
+
|
|
32
|
+
cls.test_uuid = uuid4()
|
|
33
|
+
cls.dbfs_test_folder = 'zipher_acl_test'
|
|
34
|
+
|
|
35
|
+
# Create a simple service principal that should reflect a non admin user in the client's workspace
|
|
36
|
+
# that created a job which we want to be able to run / update etc...
|
|
37
|
+
cls.sp = cls._create_user_sp()
|
|
38
|
+
cls.job_id = str(cls._create_and_run_sample_job_using_sp(sp=cls.sp))
|
|
39
|
+
|
|
40
|
+
cls.installer_config = InstallerConfig(
|
|
41
|
+
zipher_group_name=f'zipher_{cls.test_uuid}',
|
|
42
|
+
zipher_service_principal_name=f'zipher_sp_{cls.test_uuid}',
|
|
43
|
+
zipher_workspace_dir_path=f'/zipher_script_{cls.test_uuid}',
|
|
44
|
+
zipher_secret_scope_name=f'zipher_{cls.test_uuid}'
|
|
45
|
+
)
|
|
46
|
+
cls.installer = ZipherInstaller(host=cls.host, token=os.environ['DATABRICKS_TOKEN_TEST'], config=cls.installer_config)
|
|
47
|
+
zipher_creds = cls.installer.setup(user_config=InstallerUserConfig(jobs_list=[cls.job_id]), skip_approval=True)
|
|
48
|
+
|
|
49
|
+
cls.zipher_dbx = get_workspace_client(**zipher_creds)
|
|
50
|
+
|
|
51
|
+
@classmethod
|
|
52
|
+
def _create_user_sp(cls):
|
|
53
|
+
print('Creating user-level service principal to run test job with.')
|
|
54
|
+
entitlements = [iam.ComplexValue(value=ent) for ent in ['allow-cluster-create']]
|
|
55
|
+
groups = [iam.ComplexValue(value=g.id) for g in cls.admin_dbx.groups.list() if g.display_name in ['users']]
|
|
56
|
+
sp = cls.admin_dbx.service_principals.create(display_name=f"test_acl_user_sp_{cls.test_uuid}",
|
|
57
|
+
groups=groups,
|
|
58
|
+
entitlements=entitlements)
|
|
59
|
+
|
|
60
|
+
permissions_to_grant = [TokenAccessControlRequest(service_principal_name=sp.application_id,
|
|
61
|
+
permission_level=TokenPermissionLevel.CAN_USE)]
|
|
62
|
+
cls.admin_dbx.token_management.update_permissions(access_control_list=permissions_to_grant)
|
|
63
|
+
print(f'Done creating user-level service principal with name "{sp.display_name}".')
|
|
64
|
+
return sp
|
|
65
|
+
|
|
66
|
+
@classmethod
|
|
67
|
+
def _create_and_run_sample_job_using_sp(cls, sp: ServicePrincipal):
|
|
68
|
+
print('Creating test job.')
|
|
69
|
+
oauth_secret = create_oauth_secret_for_sp(dbx_client=cls.admin_dbx, sp=sp)
|
|
70
|
+
user_dbx = WorkspaceClient(host=cls.host,
|
|
71
|
+
client_id=sp.application_id,
|
|
72
|
+
client_secret=oauth_secret['secret'])
|
|
73
|
+
|
|
74
|
+
with open('sample_job.py') as sample_job:
|
|
75
|
+
dbfs_job_script_path = f"dbfs:/{cls.dbfs_test_folder}/myspark_{cls.test_uuid}.py"
|
|
76
|
+
contents = base64.b64encode(sample_job.read().encode("utf8"))
|
|
77
|
+
contents = str(contents)[2:-1]
|
|
78
|
+
user_dbx.dbfs.put(dbfs_job_script_path, contents=contents, overwrite=True)
|
|
79
|
+
|
|
80
|
+
cls.policy = cls.admin_dbx.cluster_policies.create(
|
|
81
|
+
name=f'acl_test_policy_{cls.test_uuid}',
|
|
82
|
+
definition=json.dumps({
|
|
83
|
+
"instance_pool_id": {
|
|
84
|
+
"type": "forbidden",
|
|
85
|
+
"hidden": True
|
|
86
|
+
}
|
|
87
|
+
})
|
|
88
|
+
)
|
|
89
|
+
cls.admin_dbx.permissions.update(
|
|
90
|
+
request_object_type='cluster-policies',
|
|
91
|
+
request_object_id=cls.policy.policy_id,
|
|
92
|
+
access_control_list=[AccessControlRequest(
|
|
93
|
+
service_principal_name=sp.application_id,
|
|
94
|
+
permission_level=PermissionLevel.CAN_USE
|
|
95
|
+
)]
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
cls.query = cls.admin_dbx.queries.create(
|
|
99
|
+
query=CreateQueryRequestQuery(
|
|
100
|
+
display_name=f'acl_test_query_{cls.test_uuid}',
|
|
101
|
+
query_text='select 1'
|
|
102
|
+
)
|
|
103
|
+
)
|
|
104
|
+
cls.admin_dbx.permissions.update(
|
|
105
|
+
request_object_type='queries',
|
|
106
|
+
request_object_id=cls.query.id,
|
|
107
|
+
access_control_list=[AccessControlRequest(
|
|
108
|
+
service_principal_name=sp.application_id,
|
|
109
|
+
permission_level=PermissionLevel.CAN_RUN
|
|
110
|
+
)]
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
cls.warehouse = cls.admin_dbx.warehouses.create(
|
|
114
|
+
name=f'acl_test_warehouse_{cls.test_uuid}',
|
|
115
|
+
cluster_size='2X-Small',
|
|
116
|
+
enable_serverless_compute=True,
|
|
117
|
+
max_num_clusters=1,
|
|
118
|
+
auto_stop_mins=10,
|
|
119
|
+
warehouse_type=CreateWarehouseRequestWarehouseType.PRO
|
|
120
|
+
)
|
|
121
|
+
cls.admin_dbx.warehouses.wait_get_warehouse_running(id=cls.warehouse.id)
|
|
122
|
+
cls.admin_dbx.warehouses.update_permissions(warehouse_id=cls.warehouse.id,
|
|
123
|
+
access_control_list=[WarehouseAccessControlRequest(
|
|
124
|
+
service_principal_name=sp.application_id,
|
|
125
|
+
permission_level=WarehousePermissionLevel.CAN_USE
|
|
126
|
+
)])
|
|
127
|
+
|
|
128
|
+
if 'azuredatabricks' in cls.admin_dbx.config.host:
|
|
129
|
+
spec = ClusterSpec(num_workers=1, spark_version="14.3.x-scala2.12",
|
|
130
|
+
policy_id=cls.policy.policy_id, node_type_id="Standard_F4")
|
|
131
|
+
else:
|
|
132
|
+
spec = ClusterSpec(aws_attributes=AwsAttributes(ebs_volume_count=1, ebs_volume_size=32),
|
|
133
|
+
num_workers=1, spark_version="14.3.x-scala2.12", node_type_id="m6g.large",
|
|
134
|
+
policy_id=cls.policy.policy_id)
|
|
135
|
+
job = user_dbx.jobs.create(
|
|
136
|
+
name=f"test_acl_{cls.test_uuid}",
|
|
137
|
+
job_clusters=[
|
|
138
|
+
JobCluster(
|
|
139
|
+
job_cluster_key="test_cluster_1",
|
|
140
|
+
new_cluster=spec
|
|
141
|
+
)
|
|
142
|
+
],
|
|
143
|
+
tasks=[
|
|
144
|
+
Task(
|
|
145
|
+
task_key="test_cluster_1",
|
|
146
|
+
job_cluster_key="test_cluster_1",
|
|
147
|
+
spark_python_task=SparkPythonTask(python_file=dbfs_job_script_path)
|
|
148
|
+
),
|
|
149
|
+
Task(
|
|
150
|
+
task_key="test_sql",
|
|
151
|
+
sql_task=SqlTask(
|
|
152
|
+
query=SqlTaskQuery(query_id=cls.query.id),
|
|
153
|
+
warehouse_id=cls.warehouse.id
|
|
154
|
+
)
|
|
155
|
+
)
|
|
156
|
+
]
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
print(f'Running job {job.job_id}')
|
|
160
|
+
run_wait = user_dbx.jobs.run_now(job_id=int(job.job_id))
|
|
161
|
+
print(f'Waiting for job {job.job_id} to finish...')
|
|
162
|
+
run = run_wait.result()
|
|
163
|
+
print('Done creating test job.')
|
|
164
|
+
return job.job_id
|
|
165
|
+
|
|
166
|
+
def _test_can_manage(self):
|
|
167
|
+
permissions: JobPermissions = self.admin_dbx.jobs.get_permissions(job_id=self.job_id)
|
|
168
|
+
for entity in permissions.access_control_list:
|
|
169
|
+
if entity.group_name == self.installer_config.zipher_group_name:
|
|
170
|
+
for perm in entity.all_permissions:
|
|
171
|
+
if perm.permission_level == JobPermissionLevel.CAN_MANAGE:
|
|
172
|
+
return
|
|
173
|
+
assert False, f"Job ID {self.job_id} didn't grant CAN_MANAGE permissions to {self.installer_config.zipher_group_name}"
|
|
174
|
+
|
|
175
|
+
def _test_fetch_historical_run(self):
|
|
176
|
+
run = next(self.zipher_dbx.jobs.list_runs(job_id=int(self.job_id), expand_tasks=True))
|
|
177
|
+
|
|
178
|
+
cluster_id = run.tasks[0].cluster_instance.cluster_id
|
|
179
|
+
|
|
180
|
+
self.zipher_dbx.clusters.get(cluster_id=cluster_id)
|
|
181
|
+
print(f"Successfully tested get historical cluster_id: {cluster_id}")
|
|
182
|
+
|
|
183
|
+
self.zipher_dbx.clusters.events(cluster_id=cluster_id)
|
|
184
|
+
print(f"Successfully tested get events historical cluster_id: {cluster_id}")
|
|
185
|
+
|
|
186
|
+
def _test_update_job_settings(self):
|
|
187
|
+
job = self.zipher_dbx.jobs.get(job_id=int(self.job_id))
|
|
188
|
+
|
|
189
|
+
new_settings = deepcopy(job.settings)
|
|
190
|
+
new_settings.job_clusters[0].new_cluster.num_workers = 2
|
|
191
|
+
|
|
192
|
+
self.zipher_dbx.jobs.update(job_id=int(self.job_id), new_settings=new_settings)
|
|
193
|
+
|
|
194
|
+
new_job = self.zipher_dbx.jobs.get(job_id=int(self.job_id))
|
|
195
|
+
assert new_job.settings.job_clusters[0].new_cluster.num_workers == 2, 'Job was not updated with new settings'
|
|
196
|
+
|
|
197
|
+
new_settings.job_clusters[0].new_cluster.num_workers = 1
|
|
198
|
+
self.zipher_dbx.jobs.update(job_id=int(self.job_id), new_settings=new_settings)
|
|
199
|
+
|
|
200
|
+
def _update_job_settings(self):
|
|
201
|
+
job = self.zipher_dbx.jobs.get(job_id=int(self.job_id))
|
|
202
|
+
new_settings = deepcopy(job.settings)
|
|
203
|
+
|
|
204
|
+
# Upload zipher secret and put a path to it into job settings
|
|
205
|
+
self.zipher_dbx.secrets.put_secret(scope=self.installer_config.zipher_secret_scope_name,
|
|
206
|
+
key='zipher_api_key',
|
|
207
|
+
string_value='test_key')
|
|
208
|
+
new_settings.job_clusters[0].new_cluster.spark_env_vars = {
|
|
209
|
+
'ZIPHER_API_KEY': f"{{{{secrets/{self.installer_config.zipher_secret_scope_name}/zipher_api_key}}}}"
|
|
210
|
+
}
|
|
211
|
+
|
|
212
|
+
# Upload init_script to workspace and put a path to it into job settings
|
|
213
|
+
with open("init_script.sh", "rb") as init_script:
|
|
214
|
+
self.zipher_dbx.workspace.upload(path=f'{self.installer_config.zipher_workspace_dir_path}/init_script.sh',
|
|
215
|
+
content=init_script,
|
|
216
|
+
overwrite=True, format=ImportFormat.AUTO)
|
|
217
|
+
new_settings.job_clusters[0].new_cluster.init_scripts = [InitScriptInfo(
|
|
218
|
+
workspace=WorkspaceStorageInfo(
|
|
219
|
+
destination=f'{self.installer_config.zipher_workspace_dir_path}/init_script.sh'
|
|
220
|
+
)
|
|
221
|
+
)]
|
|
222
|
+
|
|
223
|
+
# Update log path
|
|
224
|
+
new_settings.job_clusters[0].new_cluster.cluster_log_conf = ClusterLogConf(
|
|
225
|
+
dbfs=DbfsStorageInfo(destination=f'dbfs:/{self.dbfs_test_folder}')
|
|
226
|
+
)
|
|
227
|
+
|
|
228
|
+
self.zipher_dbx.jobs.update(job_id=int(self.job_id), new_settings=new_settings)
|
|
229
|
+
print(f"Successfully updated secret and init_script to job_id {self.job_id}")
|
|
230
|
+
|
|
231
|
+
def _wait_for_cluster_state_running(self, cluster_id):
|
|
232
|
+
while True:
|
|
233
|
+
cluster_info = self.zipher_dbx.clusters.get(cluster_id=cluster_id)
|
|
234
|
+
state = cluster_info.state.value
|
|
235
|
+
print(f"Cluster state: {state}")
|
|
236
|
+
|
|
237
|
+
if state == 'RUNNING':
|
|
238
|
+
print("Cluster is up and running!")
|
|
239
|
+
break
|
|
240
|
+
elif state == 'TERMINATED':
|
|
241
|
+
raise RuntimeError("Cluster terminated unexpectedly.")
|
|
242
|
+
elif state == 'ERROR':
|
|
243
|
+
raise RuntimeError("Cluster failed to start.")
|
|
244
|
+
|
|
245
|
+
time.sleep(10)
|
|
246
|
+
|
|
247
|
+
def _resize_job(self, cluster_id):
|
|
248
|
+
self.zipher_dbx.clusters.resize(cluster_id=cluster_id, num_workers=2)
|
|
249
|
+
|
|
250
|
+
assert self.zipher_dbx.clusters.get(cluster_id=cluster_id).num_workers == 2, "Didn't manage to resize cluster"
|
|
251
|
+
|
|
252
|
+
def _wait_for_run_success_end(self, run_id):
|
|
253
|
+
while True:
|
|
254
|
+
run_info = self.zipher_dbx.jobs.get_run(run_id=run_id)
|
|
255
|
+
life_cycle_state = run_info.state.life_cycle_state
|
|
256
|
+
result_state = run_info.state.result_state
|
|
257
|
+
|
|
258
|
+
print(f"Job run state: {life_cycle_state}")
|
|
259
|
+
|
|
260
|
+
# Check if the job has finished
|
|
261
|
+
if life_cycle_state == RunLifeCycleState.TERMINATED:
|
|
262
|
+
if result_state == RunResultState.SUCCESS:
|
|
263
|
+
print("Job completed successfully.")
|
|
264
|
+
return
|
|
265
|
+
elif result_state == RunResultState.FAILED:
|
|
266
|
+
raise RuntimeError("Job failed with status: FAILED")
|
|
267
|
+
elif result_state == RunResultState.TIMEDOUT:
|
|
268
|
+
raise RuntimeError("Job failed with status: TIMEDOUT")
|
|
269
|
+
elif result_state == RunResultState.CANCELED:
|
|
270
|
+
raise RuntimeError("Job was canceled.")
|
|
271
|
+
|
|
272
|
+
# Sleep before polling again
|
|
273
|
+
time.sleep(10)
|
|
274
|
+
|
|
275
|
+
def _test_init_script_log(self, cluster_id):
|
|
276
|
+
init_script_stdout_files = [
|
|
277
|
+
file
|
|
278
|
+
for file in self.zipher_dbx.dbfs.list(path=f'dbfs:/{self.dbfs_test_folder}/{cluster_id}/init_scripts',
|
|
279
|
+
recursive=True)
|
|
280
|
+
if file.path.endswith('stdout.log')
|
|
281
|
+
]
|
|
282
|
+
|
|
283
|
+
assert init_script_stdout_files, 'No init script logs found'
|
|
284
|
+
|
|
285
|
+
script_stdout_content = [
|
|
286
|
+
base64.b64decode(self.zipher_dbx.dbfs.read(file.path).data).decode("utf-8")
|
|
287
|
+
for file in init_script_stdout_files
|
|
288
|
+
]
|
|
289
|
+
|
|
290
|
+
assert any([
|
|
291
|
+
'Zipher init_script test succeeded!' in stdout
|
|
292
|
+
for stdout in script_stdout_content
|
|
293
|
+
])
|
|
294
|
+
|
|
295
|
+
def _test_run_output(self, run_id):
|
|
296
|
+
task_run_id = self.zipher_dbx.jobs.get_run(run_id=run_id).tasks[0].run_id
|
|
297
|
+
task_output = self.zipher_dbx.jobs.get_run_output(run_id=task_run_id).logs
|
|
298
|
+
|
|
299
|
+
assert 'Successfully retrieved secret' in task_output
|
|
300
|
+
|
|
301
|
+
def _test_new_job_run(self):
|
|
302
|
+
self._update_job_settings()
|
|
303
|
+
|
|
304
|
+
run = self.zipher_dbx.jobs.run_now(job_id=int(self.job_id))
|
|
305
|
+
time.sleep(20)
|
|
306
|
+
cluster_id = self.zipher_dbx.jobs.get_run(run.run_id).tasks[0].cluster_instance.cluster_id
|
|
307
|
+
|
|
308
|
+
self._wait_for_cluster_state_running(cluster_id=cluster_id)
|
|
309
|
+
self._resize_job(cluster_id=cluster_id)
|
|
310
|
+
self._wait_for_run_success_end(run_id=run.run_id)
|
|
311
|
+
|
|
312
|
+
time.sleep(20)
|
|
313
|
+
self._test_init_script_log(cluster_id=cluster_id)
|
|
314
|
+
self._test_run_output(run_id=run.run_id)
|
|
315
|
+
|
|
316
|
+
def _test_can_read_policy(self):
|
|
317
|
+
policy_id = self.zipher_dbx.jobs.get(job_id=int(self.job_id)).settings.job_clusters[0].new_cluster.policy_id
|
|
318
|
+
policy = self.zipher_dbx.cluster_policies.get(policy_id=policy_id)
|
|
319
|
+
|
|
320
|
+
assert policy
|
|
321
|
+
|
|
322
|
+
def _test_cancel_run(self):
|
|
323
|
+
run = self.zipher_dbx.jobs.run_now(job_id=int(self.job_id))
|
|
324
|
+
time.sleep(20)
|
|
325
|
+
cancel_response = self.zipher_dbx.jobs.cancel_run_and_wait(run_id=run.run_id)
|
|
326
|
+
|
|
327
|
+
assert cancel_response.status.termination_details.message == "Run cancelled by user", f"Didn't manage to cancel run {run.run_id}"
|
|
328
|
+
|
|
329
|
+
def _test_can_store_secret(self):
|
|
330
|
+
self.zipher_dbx.secrets.put_secret(
|
|
331
|
+
scope=self.installer_config.zipher_secret_scope_name,
|
|
332
|
+
key='test_key',
|
|
333
|
+
string_value='test_value'
|
|
334
|
+
)
|
|
335
|
+
|
|
336
|
+
def test(self):
|
|
337
|
+
self._test_can_manage()
|
|
338
|
+
self._test_can_read_policy()
|
|
339
|
+
self._test_can_store_secret()
|
|
340
|
+
self._test_fetch_historical_run()
|
|
341
|
+
self._test_update_job_settings()
|
|
342
|
+
self._test_new_job_run()
|
|
343
|
+
self._test_cancel_run()
|
|
344
|
+
|
|
345
|
+
@classmethod
|
|
346
|
+
def tearDownClass(cls):
|
|
347
|
+
cls.installer.clean_resources()
|
|
348
|
+
cls.admin_dbx.cluster_policies.delete(policy_id=cls.policy.policy_id)
|
|
349
|
+
cls.admin_dbx.service_principals.delete(id=cls.sp.id)
|
|
350
|
+
cls.admin_dbx.jobs.delete(job_id=int(cls.job_id))
|
|
351
|
+
cls.admin_dbx.dbfs.delete(path=f'dbfs:/{cls.dbfs_test_folder}', recursive=True)
|
|
352
|
+
cls.admin_dbx.queries.delete(cls.query.id)
|
|
353
|
+
cls.admin_dbx.warehouses.delete(cls.warehouse.id)
|
|
File without changes
|