Merge pull request #28 from nirupama-dev/main

Composer Scheduler: Dataproc multi tenant cluster enablement

Author: harsha-accenture

scheduler_jupyter_plugin/dagTemplates/pysparkJobTemplate-v1.txt

@@ -157,6 +157,9 @@ submit_pyspark_job = DataprocSubmitJobOperator(
             'args' : notebook_args
         },
     },
+    {% if multi_tenant_service_account %}
+    impersonation_chain=['{{multi_tenant_service_account}}'],
+    {% endif %}
     gcp_conn_id='google_cloud_default',  # Reference to the GCP connection
     dag=dag,
 )

scheduler_jupyter_plugin/services/executor.py

@@ -41,6 +41,7 @@
     UTF8,
     PAYLOAD_JSON_FILE_PATH,
     HTTP_STATUS_OK,
+    DATAPROC_SERVICE_NAME,
 )
 from scheduler_jupyter_plugin.models.models import DescribeJob
 from scheduler_jupyter_plugin.services import airflow
@@ -147,7 +148,50 @@ async def upload_to_gcs(
             self.log.exception(f"Error uploading file to GCS: {str(error)}")
             raise IOError(str(error))
 
-    def prepare_dag(self, job, gcs_dag_bucket, dag_file, project_id, region_id):
+    async def get_cluster_details(self, cluster_name):
+        try:
+            dataproc_url = await urls.gcp_service_url(DATAPROC_SERVICE_NAME)
+            api_endpoint = f"{dataproc_url}/v1/projects/{self.project_id}/regions/{self.region_id}/clusters/{cluster_name}"
+            async with self.client_session.get(
+                api_endpoint, headers=self.create_headers()
+            ) as response:
+                if response.status == HTTP_STATUS_OK:
+                    resp = await response.json()
+                    return resp
+                else:
+                    return {
+                        "error": f"Failed to fetch clusters: {response.status} {await response.text()}"
+                    }
+
+        except Exception as e:
+            self.log.exception("Error fetching cluster list")
+            return {"error": str(e)}
+
+    async def multi_tenant_user_service_account(self, cluster_name):
+        cluster_data = await self.get_cluster_details(cluster_name)
+        if cluster_data:
+            multi_tenant = (
+                cluster_data.get("config", {})
+                .get("softwareConfig", {})
+                .get("properties", {})
+                .get("dataproc:dataproc.dynamic.multi.tenancy.enabled", "false")
+            )
+            if multi_tenant == "true":
+                cmd = "config get account"
+                process = await async_run_gcloud_subcommand(cmd)
+                user_email = process.strip()
+                service_account = (
+                    cluster_data.get("config", {})
+                    .get("securityConfig", {})
+                    .get("identityConfig", {})
+                    .get("userServiceAccountMapping", {})
+                    .get(user_email, "")
+                )
+                if service_account:
+                    return service_account
+        return ""
+
+    async def prepare_dag(self, job, gcs_dag_bucket, dag_file, project_id, region_id):
         self.log.info("Generating dag file")
         DAG_TEMPLATE_CLUSTER_V1 = "pysparkJobTemplate-v1.txt"
         DAG_TEMPLATE_SERVERLESS_V1 = "pysparkBatchTemplate-v1.txt"
@@ -181,6 +225,11 @@ def prepare_dag(self, job, gcs_dag_bucket, dag_file, project_id, region_id):
             parameters = ""
         if job.local_kernel is False:
             if job.mode_selected == "cluster":
+                multi_tenant_service_account = (
+                    await self.multi_tenant_user_service_account(
+                        cluster_name=job.cluster_name
+                    )
+                )
                 template = environment.get_template(DAG_TEMPLATE_CLUSTER_V1)
                 if not job.input_filename.startswith(GCS):
                     input_notebook = f"gs://{gcs_dag_bucket}/dataproc-notebooks/{job.name}/input_notebooks/{job.input_filename}"
@@ -198,6 +247,7 @@ def prepare_dag(self, job, gcs_dag_bucket, dag_file, project_id, region_id):
                     start_date=start_date,
                     parameters=parameters,
                     time_zone=time_zone,
+                    multi_tenant_service_account=multi_tenant_service_account,
                 )
             else:
                 template = environment.get_template(DAG_TEMPLATE_SERVERLESS_V1)
@@ -402,7 +452,7 @@ async def execute(self, input_data, project_id, region_id):
                 destination_dir=f"dataproc-notebooks/{job_name}/dag_details",
             )
 
-            file_path = self.prepare_dag(
+            file_path = await self.prepare_dag(
                 job, gcs_dag_bucket, dag_file, project_id, region_id
             )
             await self.upload_to_gcs(

scheduler_jupyter_plugin/tests/test_dataproc.py

@@ -33,7 +33,7 @@ async def test_list_clusters(monkeypatch, jp_fetch):
     payload = json.loads(response.body)
     assert (
         payload["api_endpoint"]
-        == f"https://dataproc.googleapis.com//v1/projects/credentials-project/regions/{mock_region_id}/clusters?pageSize={mock_page_size}&pageToken={mock_page_token}"
+        == f"https://dataproc.googleapis.com//v1/projects/{mock_project_id}/regions/{mock_region_id}/clusters?pageSize={mock_page_size}&pageToken={mock_page_token}"
     )
     assert payload["headers"]["Authorization"] == f"Bearer mock-token"
 

src/controls/RegionDropdown.tsx

@@ -38,6 +38,7 @@ type Props = {
   /** Initial loading flag for region */
   loaderRegion?: boolean;
   setLoaderRegion?: (value: boolean) => void;
+  label?: string;
 };
 
 /**
@@ -53,7 +54,8 @@ export function RegionDropdown(props: Props) {
     regionDisable,
     fromPage,
     loaderRegion,
-    setLoaderRegion
+    setLoaderRegion,
+    label
   } = props;
   let regionStrList: string[] = [];
 
@@ -77,7 +79,7 @@ export function RegionDropdown(props: Props) {
       renderInput={params => (
         <TextField
           {...params}
-          label={'Region*'}
+          label={label || 'Region*'}
           InputProps={{
             ...params.InputProps,
             endAdornment: (

src/scheduler/composer/CreateNotebookScheduler.tsx

@@ -717,7 +717,7 @@ const CreateNotebookScheduler = ({
                   setProjectId(projectId ?? '')
                 }
                 fetchFunc={projectListAPI}
-                label="Project ID*"
+                label="Composer Project ID*"
                 // Always show the clear indicator and hide the dropdown arrow
                 // make it very clear that this is an autocomplete.
                 sx={{
@@ -740,6 +740,7 @@ const CreateNotebookScheduler = ({
                 editMode={editMode}
                 loaderRegion={loaderRegion}
                 setLoaderRegion={setLoaderRegion}
+                label={'Composer Region*'}
               />
             </div>
             {!region && <ErrorMessage message="Region is required" />}

src/scheduler/composer/ListNotebookScheduler.tsx

@@ -829,7 +829,7 @@ function ListNotebookScheduler({
                   handleProjectIdChange(projectId);
                 }}
                 fetchFunc={projectListAPI}
-                label="Project ID*"
+                label="Composer Project ID*"
                 // Always show the clear indicator and hide the dropdown arrow
                 // make it very clear that this is an autocomplete.
                 sx={{
@@ -855,6 +855,7 @@ function ListNotebookScheduler({
                 onRegionChange={region => handleRegionChange(region)}
                 loaderRegion={loaderRegion}
                 setLoaderRegion={setLoaderRegion}
+                label={'Composer Region*'}
               />
             </div>
             {!region && (