From 64dfc3c8bf5eea541b684db97b57100a1cedf70a Mon Sep 17 00:00:00 2001
From: rubenbblazquez <ruben.barroso@t2o.com>
Date: Wed, 28 Aug 2024 11:38:34 +0200
Subject: [PATCH] [feat]: Adapt GCPCluster to be able to initialize the cluster
 by passing credentials without obtaining from the environment

---
 dask_cloudprovider/gcp/instances.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
index 5ef9c0da..cc5517cb 100644
--- a/dask_cloudprovider/gcp/instances.py
+++ b/dask_cloudprovider/gcp/instances.py
@@ -4,6 +4,7 @@
 import json
 
 import sqlite3
+from typing import Optional, Any
 
 import dask
 from dask.utils import tmpfile
@@ -106,7 +107,6 @@ def __init__(
         self.instance_labels = _instance_labels
 
         self.general_zone = "-".join(self.zone.split("-")[:2])  # us-east1-c -> us-east1
-
         self.service_account = service_account or self.config.get("service_account")
 
     def create_gcp_config(self):
@@ -494,6 +494,8 @@ class GCPCluster(VMCluster):
     service_account: str
         Service account that all VMs will run under.
         Defaults to the default Compute Engine service account for your GCP project.
+    service_account_credentials: Optional[dict[str, Any]]
+        Service account credentials to create the compute engine Vms
 
     Examples
     --------
@@ -587,9 +589,10 @@ def __init__(
         debug=False,
         instance_labels=None,
         service_account=None,
+        service_account_credentials: Optional[dict[str, Any]]=None,
         **kwargs,
     ):
-        self.compute = GCPCompute()
+        self.compute = GCPCompute(service_account_credentials)
 
         self.config = dask.config.get("cloudprovider.gcp", {})
         self.auto_shutdown = (
@@ -641,9 +644,17 @@ def __init__(
 
 
 class GCPCompute:
-    """Wrapper for the ``googleapiclient`` compute object."""
+    """
+    Wrapper for the ``googleapiclient`` compute object.
+
+    Attributes
+    ----------
+    service_account_credentials: Optional[dict]
+        Service account credentials to create the compute engine Vms
+    """
 
-    def __init__(self):
+    def __init__(self, service_account_credentials: Optional[dict[str, Any]] = None):
+        self.service_account_credentials = service_account_credentials or {}
         self._compute = self.refresh_client()
 
     def refresh_client(self):
@@ -654,6 +665,13 @@ def refresh_client(self):
                 os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
                 scopes=["https://www.googleapis.com/auth/cloud-platform"],
             )
+        elif self.service_account_credentials:
+            import google.oauth2.service_account  # google-auth
+
+            creds = google.oauth2.service_account.Credentials.from_service_account_info(
+                self.service_account_credentials,
+                scopes=["https://www.googleapis.com/auth/cloud-platform"],
+            )
         else:
             import google.auth.credentials  # google-auth