From b685e0bdb0207587c51cfae8974917bd3b9a9a2a Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Fri, 15 Dec 2023 22:49:51 +0100 Subject: [PATCH] Tested all secrets stores and fixed remaining bugs --- .../zenml-self-hosted/deploy-with-helm.md | 14 +++++++- .../deploy/helm/templates/server-db-job.yaml | 4 +-- .../helm/templates/server-deployment.yaml | 4 +-- .../secrets_stores/gcp_secrets_store.py | 6 +--- .../service_connector_secrets_store.py | 32 +++++++++++++++++-- 5 files changed, 47 insertions(+), 13 deletions(-) diff --git a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md index 2614bd8fcb5..852938a014d 100644 --- a/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md +++ b/docs/book/deploying-zenml/zenml-self-hosted/deploy-with-helm.md @@ -386,7 +386,19 @@ The GCP Secrets Store uses the ZenML GCP Service Connector under the hood to aut # GCP credentials JSON to use to authenticate with the GCP Secrets # Manager instance. - google_application_credentials: '{...}' + google_application_credentials: | + { + "type": "service_account", + "project_id": "my-project", + "private_key_id": "...", + "private_key": "-----BEGIN PRIVATE KEY-----\n...=\n-----END PRIVATE KEY-----\n", + "client_email": "...", + "client_id": "...", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "..." + } serviceAccount: diff --git a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml index 5d47721425a..d3fbff15fc6 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml @@ -88,7 +88,7 @@ spec: {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} - {{- endif }} + {{- end }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} @@ -97,7 +97,7 @@ spec: {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} - {{- endif }} + {{- end }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json diff --git a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml index d16dca8726a..1ff1a3ad563 100644 --- a/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml +++ b/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml @@ -166,7 +166,7 @@ spec: {{- if .Values.zenml.secretsStore.aws.region_name }} - name: ZENML_SECRETS_STORE_REGION_NAME value: {{ .Values.zenml.secretsStore.aws.region_name | quote }} - {{- endif }} + {{- end }} - name: ZENML_SECRETS_STORE_SECRET_LIST_REFRESH_TIMEOUT value: {{ .Values.zenml.secretsStore.aws.secret_list_refresh_timeout | quote }} {{- else if eq .Values.zenml.secretsStore.type "gcp" }} @@ -175,7 +175,7 @@ spec: {{- if .Values.zenml.secretsStore.gcp.project_id }} - name: ZENML_SECRETS_STORE_PROJECT_ID value: {{ .Values.zenml.secretsStore.gcp.project_id | quote }} - {{- endif }} + {{- end }} {{- if .Values.zenml.secretsStore.gcp.google_application_credentials }} - name: GOOGLE_APPLICATION_CREDENTIALS value: /gcp-credentials/credentials.json diff --git a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py index 3bc049c5906..06d5b4a649f 100644 --- a/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py @@ -35,7 +35,6 @@ from google.api_core import exceptions as google_exceptions from google.cloud.secretmanager import SecretManagerServiceClient -from google.oauth2 import service_account as gcp_service_account from pydantic import root_validator from zenml.analytics.enums import AnalyticsEvent @@ -183,10 +182,7 @@ def _initialize_client_from_connector(self, client: Any) -> Any: Returns: The GCP Secrets Manager client. """ - assert isinstance(client, gcp_service_account.Credentials) - return SecretManagerServiceClient( - project=self.config.project_id, credentials=client - ) + return SecretManagerServiceClient(credentials=client) # ------ # Secrets diff --git a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py index c33bc38d6cf..32372ae9a28 100644 --- a/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py +++ b/src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py @@ -14,6 +14,7 @@ """Base secrets store class used for all secrets stores that use a service connector.""" +import json from abc import abstractmethod from threading import Lock from typing import ( @@ -25,7 +26,7 @@ ) from uuid import uuid4 -from pydantic import Field +from pydantic import Field, root_validator from zenml.config.secrets_store_config import SecretsStoreConfiguration from zenml.logger import get_logger @@ -54,6 +55,20 @@ class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration): auth_method: str auth_config: Dict[str, Any] = Field(default_factory=dict) + @root_validator(pre=True) + def validate_auth_config(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Convert the authentication configuration if given in JSON format. + + Args: + values: The configuration values. + + Returns: + The validated configuration values. + """ + if isinstance(values.get("auth_config"), str): + values["auth_config"] = json.loads(values["auth_config"]) + return values + class ServiceConnectorSecretsStore(BaseSecretsStore): """Base secrets store class for service connector-based secrets stores. @@ -78,10 +93,11 @@ class ServiceConnectorSecretsStore(BaseSecretsStore): _connector: Optional[ServiceConnector] = None _client: Optional[Any] = None - _lock: Lock = Field(default_factory=Lock) + _lock: Optional[Lock] = None def _initialize(self) -> None: """Initialize the secrets store.""" + self._lock = Lock() # Initialize the client early, just to catch any configuration or # authentication errors early, before the Secrets Store is used. _ = self.client @@ -126,6 +142,16 @@ def _get_client(self) -> Any: self._client = self._initialize_client_from_connector(client) return self._client + @property + def lock(self) -> Lock: + """Get the lock used to treat the client initialization as a critical section. + + Returns: + The lock instance. + """ + assert self._lock is not None + return self._lock + @property def client(self) -> Any: """Get the secrets store API client. @@ -137,7 +163,7 @@ def client(self) -> Any: # of different threads. We want to make sure that we only initialize # the client once, and then reuse it. We have to use a lock to treat # this method as a critical section. - with self._lock: + with self.lock: return self._get_client() @abstractmethod