From e925c589216f0be20343377111946a15b41b61d6 Mon Sep 17 00:00:00 2001 From: Pari Date: Fri, 25 Aug 2023 23:30:50 +0200 Subject: [PATCH] fix: generate the bundle path including prefix/suffix if available (#20) * generate the bundle path including prefix/suffix if available --- brickflow/codegen/databricks_bundle.py | 16 + .../local_bundle_prefix_suffix.yml | 276 ++++++++++++++++++ tests/codegen/test_databricks_bundle.py | 49 ++++ 3 files changed, 341 insertions(+) create mode 100644 tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml diff --git a/brickflow/codegen/databricks_bundle.py b/brickflow/codegen/databricks_bundle.py index 0fc8a42c..484a8205 100644 --- a/brickflow/codegen/databricks_bundle.py +++ b/brickflow/codegen/databricks_bundle.py @@ -617,6 +617,22 @@ def proj_to_bundle(self) -> DatabricksAssetBundles: / self.env ) + bundle_suffix = config( + BrickflowEnvVars.BRICKFLOW_WORKFLOW_SUFFIX.value, + default=None, + ) + + bundle_prefix = config( + BrickflowEnvVars.BRICKFLOW_WORKFLOW_PREFIX.value, + default=None, + ) + + if bundle_prefix is not None: + bundle_root_path = bundle_root_path / bundle_prefix + + if bundle_suffix is not None: + bundle_root_path = bundle_root_path / bundle_suffix + env_content = Environments( workspace=Workspace( root_path=str(bundle_root_path.as_posix()), diff --git a/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml b/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml new file mode 100644 index 00000000..94593ccf --- /dev/null +++ b/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml @@ -0,0 +1,276 @@ +bundle: + name: test-project +environments: + test-project-local: + resources: + jobs: + test: + email_notifications: null + notification_settings: null + webhook_notifications: null + git_source: null + job_clusters: [] + max_concurrent_runs: 1.0 + name: test_user_test + run_as: + user_name: abc@abc.com + permissions: + - level: IS_OWNER + user_name: abc@abc.com + - level: CAN_MANAGE + user_name: abc@abc.com + - level: CAN_MANAGE_RUN + user_name: abc@abc.com + - level: CAN_VIEW + user_name: abc@abc.com + schedule: + quartz_cron_expression: '* * * * *' + timezone_id: UTC + tags: + brickflow_project_name: test-project + brickflow_deployment_mode: Databricks Asset Bundles + deployed_by: test_user + brickflow_version: 1.0.0 + environment: local + test: test2 + tasks: + - depends_on: [] + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: custom_python_task_push + timeout_seconds: null + - depends_on: [] + pipeline_task: + pipeline_id: ${resources.pipelines.test_hello_world.id} + task_key: dlt_pipeline + - depends_on: [] + pipeline_task: + pipeline_id: ${resources.pipelines.test_hello_world.id} + task_key: dlt_pipeline_2 + - depends_on: [ ] + email_notifications: { } + existing_cluster_id: existing_cluster_id + libraries: [ ] + max_retries: null + min_retry_interval_millis: null + notebook_task: + notebook_path: notebooks/notebook_a + retry_on_timeout: null + task_key: notebook_task_a + timeout_seconds: null + - depends_on: [] + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + test: var + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function + timeout_seconds: null + - depends_on: + - task_key: task_function + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function_2 + timeout_seconds: null + - depends_on: + - task_key: task_function_2 + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function_3 + timeout_seconds: null + - depends_on: + - task_key: task_function_3 + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function_4 + timeout_seconds: null + - depends_on: [] + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + test: var + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function_no_deco_args + timeout_seconds: null + - depends_on: [] + email_notifications: {} + existing_cluster_id: existing_cluster_id + libraries: [] + max_retries: null + min_retry_interval_millis: null + notebook_task: + base_parameters: + all_tasks1: test + all_tasks3: '123' + brickflow_env: local + brickflow_internal_only_run_tasks: '' + brickflow_internal_task_name: '{{task_key}}' + brickflow_internal_workflow_name: test + brickflow_internal_workflow_prefix: '' + brickflow_internal_workflow_suffix: '' + brickflow_job_id: '{{job_id}}' + brickflow_parent_run_id: '{{parent_run_id}}' + brickflow_run_id: '{{run_id}}' + brickflow_start_date: '{{start_date}}' + brickflow_start_time: '{{start_time}}' + brickflow_task_key: '{{task_key}}' + brickflow_task_retry_count: '{{task_retry_count}}' + notebook_path: test_databricks_bundle.py + source: WORKSPACE + retry_on_timeout: null + task_key: task_function_nokwargs + timeout_seconds: null + pipelines: + test_hello_world: + catalog: null + channel: current + configuration: {} + continuous: false + development: true + edition: advanced + libraries: + - notebook: + path: scripts/spark_script_2.py + name: test_user_hello world + photon: false + storage: '123' + target: null + workspace: + file_path: /Users/${workspace.current_user.userName}/.brickflow_bundles/test-project/local/_prefix/_suffix/files + root_path: /Users/${workspace.current_user.userName}/.brickflow_bundles/test-project/local/_prefix/_suffix + state_path: /Users/${workspace.current_user.userName}/.brickflow_bundles/test-project/local/_prefix/_suffix/state +workspace: {} diff --git a/tests/codegen/test_databricks_bundle.py b/tests/codegen/test_databricks_bundle.py index 858ed0c1..dc43e534 100644 --- a/tests/codegen/test_databricks_bundle.py +++ b/tests/codegen/test_databricks_bundle.py @@ -102,6 +102,55 @@ def test_generate_bundle_local( if os.path.exists(BUNDLE_FILE_NAME): os.remove(BUNDLE_FILE_NAME) + @patch.dict( + os.environ, + { + BrickflowEnvVars.BRICKFLOW_MODE.value: Stage.deploy.value, + BrickflowEnvVars.BRICKFLOW_ENV.value: "local", + BrickflowEnvVars.BRICKFLOW_DEPLOYMENT_MODE.value: BrickflowDeployMode.BUNDLE.value, + BrickflowEnvVars.BRICKFLOW_WORKFLOW_PREFIX.value: "_prefix", + BrickflowEnvVars.BRICKFLOW_WORKFLOW_SUFFIX.value: "_suffix", + }, + ) + @patch("subprocess.check_output") + @patch("brickflow.context.ctx.get_parameter") + @patch("importlib.metadata.version") + def test_generate_bundle_local_prefix_suffix( + self, + bf_version_mock: Mock, + dbutils: Mock, + sub_proc_mock: Mock, + ): + dbutils.return_value = None + sub_proc_mock.return_value = b"" + bf_version_mock.return_value = "1.0.0" + workspace_client = get_workspace_client_mock() + # get caller part breaks here + with Project( + "test-project", + entry_point_path="test_databricks_bundle.py", + codegen_kwargs={ + "mutators": [ + DatabricksBundleTagsAndNameMutator( + databricks_client=workspace_client + ) + ] + }, # dont test import mutator + ) as f: + f.add_workflow(wf) + + with open(BUNDLE_FILE_NAME, "r", encoding="utf-8") as bundle: + bundle_content = bundle.read() + assert bundle_content is not None + assert len(bundle_content) > 0 + + actual = read_yaml_file(BUNDLE_FILE_NAME) + expected = get_expected_bundle_yaml("local_bundle_prefix_suffix.yml") + bf_version_mock.assert_called_once() + assert_equal_dicts(actual, expected) + if os.path.exists(BUNDLE_FILE_NAME): + os.remove(BUNDLE_FILE_NAME) + @patch.dict( os.environ, {