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.
Files changed (32) hide show
  1. zipher-sdk-0.3.0/PKG-INFO +139 -0
  2. zipher-sdk-0.3.0/README.md +122 -0
  3. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/setup.py +9 -4
  4. zipher-sdk-0.3.0/tests/test_acl/__init__.py +0 -0
  5. zipher-sdk-0.3.0/tests/test_acl/sample_job.py +15 -0
  6. zipher-sdk-0.3.0/tests/test_acl/test_acl.py +353 -0
  7. zipher-sdk-0.3.0/zipher/acl/__init__.py +0 -0
  8. zipher-sdk-0.3.0/zipher/acl/models.py +49 -0
  9. zipher-sdk-0.3.0/zipher/acl/zipher_installer.py +376 -0
  10. zipher-sdk-0.3.0/zipher/cli.py +153 -0
  11. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher/client.py +9 -7
  12. zipher-sdk-0.3.0/zipher/exceptions.py +16 -0
  13. zipher-sdk-0.3.0/zipher/models.py +26 -0
  14. zipher-sdk-0.3.0/zipher/utils.py +67 -0
  15. zipher-sdk-0.3.0/zipher_sdk.egg-info/PKG-INFO +139 -0
  16. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher_sdk.egg-info/SOURCES.txt +9 -1
  17. zipher-sdk-0.3.0/zipher_sdk.egg-info/entry_points.txt +3 -0
  18. zipher-sdk-0.3.0/zipher_sdk.egg-info/requires.txt +3 -0
  19. zipher-sdk-0.2.0/PKG-INFO +0 -75
  20. zipher-sdk-0.2.0/README.md +0 -58
  21. zipher-sdk-0.2.0/zipher/conf.py +0 -6
  22. zipher-sdk-0.2.0/zipher/exceptions.py +0 -4
  23. zipher-sdk-0.2.0/zipher/models.py +0 -8
  24. zipher-sdk-0.2.0/zipher_sdk.egg-info/PKG-INFO +0 -75
  25. zipher-sdk-0.2.0/zipher_sdk.egg-info/requires.txt +0 -3
  26. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/LICENSE +0 -0
  27. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/setup.cfg +0 -0
  28. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/tests/__init__.py +0 -0
  29. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/tests/test_client.py +0 -0
  30. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher/__init__.py +0 -0
  31. {zipher-sdk-0.2.0 → zipher-sdk-0.3.0}/zipher_sdk.egg-info/dependency_links.txt +0 -0
  32. {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.2.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.1.1",
20
- "requests==2.31.0",
21
- "databricks-sdk==0.28.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