zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -19,6 +19,8 @@ services:
 - Explicit GCP service account key
 
 """
+
+import base64
 import datetime
 import json
 import os
@@ -32,14 +34,21 @@ import google.api_core.exceptions
 import google.auth
 import google.auth.exceptions
 import requests
+from google.auth import aws as gcp_aws
+from google.auth import external_account as gcp_external_account
 from google.auth import (
     impersonated_credentials as gcp_impersonated_credentials,
 )
+from google.auth._default import (
+    _AWS_SUBJECT_TOKEN_TYPE,
+    _get_external_account_credentials,
+)
 from google.auth.transport.requests import Request
-from google.cloud import container_v1, storage
+from google.cloud import artifactregistry_v1, container_v1, storage
+from google.cloud.location import locations_pb2
 from google.oauth2 import credentials as gcp_credentials
 from google.oauth2 import service_account as gcp_service_account
-from pydantic import Field, SecretStr, root_validator, validator
+from pydantic import Field, field_validator, model_validator
 
 from zenml.constants import (
     DOCKER_REGISTRY_RESOURCE_TYPE,
@@ -67,18 +76,21 @@ from zenml.service_connectors.service_connector import (
     ServiceConnector,
 )
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
 GKE_KUBE_API_TOKEN_EXPIRATION = 60
 DEFAULT_IMPERSONATE_TOKEN_EXPIRATION = 3600  # 1 hour
+GCP_SESSION_EXPIRATION_BUFFER = 15  # 15 minutes
 
 
 class GCPUserAccountCredentials(AuthenticationConfig):
     """GCP user account credentials."""
 
-    user_account_json: SecretStr = Field(
-        title="GCP User Account Credentials JSON",
+    user_account_json: PlainSerializedSecretStr = Field(
+        title="GCP User Account Credentials JSON optionally base64 encoded.",
     )
 
     generate_temporary_tokens: bool = Field(
@@ -89,30 +101,48 @@ class GCPUserAccountCredentials(AuthenticationConfig):
         "distribute the user account credentials JSON to clients instead.",
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
     def validate_user_account_dict(
-        cls, values: Dict[str, Any]
+        cls, data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Convert the user account credentials to JSON if given in dict format.
 
         Args:
-            values: The configuration values.
+            data: The configuration values.
 
         Returns:
             The validated configuration values.
+
+        Raises:
+            ValueError: If the user account credentials JSON is invalid.
         """
-        if isinstance(values.get("user_account_json"), dict):
-            values["user_account_json"] = json.dumps(
-                values["user_account_json"]
-            )
-        return values
+        user_account_json = data.get("user_account_json")
+        if isinstance(user_account_json, dict):
+            data["user_account_json"] = json.dumps(data["user_account_json"])
+        elif isinstance(user_account_json, str):
+            # Check if the user account JSON is base64 encoded and decode it
+            if re.match(r"^[A-Za-z0-9+/=]+$", user_account_json):
+                try:
+                    data["user_account_json"] = base64.b64decode(
+                        user_account_json
+                    ).decode("utf-8")
+                except Exception as e:
+                    raise ValueError(
+                        f"Failed to decode base64 encoded user account JSON: {e}"
+                    )
+        return data
 
-    @validator("user_account_json")
-    def validate_user_account_json(cls, v: SecretStr) -> SecretStr:
+    @field_validator("user_account_json")
+    @classmethod
+    def validate_user_account_json(
+        cls, value: PlainSerializedSecretStr
+    ) -> PlainSerializedSecretStr:
         """Validate the user account credentials JSON.
 
         Args:
-            v: The user account credentials JSON.
+            value: The user account credentials JSON.
 
         Returns:
             The validated user account credentials JSON.
@@ -121,7 +151,7 @@ class GCPUserAccountCredentials(AuthenticationConfig):
             ValueError: If the user account credentials JSON is invalid.
         """
         try:
-            user_account_info = json.loads(v.get_secret_value())
+            user_account_info = json.loads(value.get_secret_value())
         except json.JSONDecodeError as e:
             raise ValueError(
                 f"GCP user account credentials is not a valid JSON: {e}"
@@ -139,7 +169,7 @@ class GCPUserAccountCredentials(AuthenticationConfig):
         if missing_fields:
             raise ValueError(
                 f"GCP user account credentials JSON is missing required "
-                f'fields: {", ".join(list(missing_fields))}'
+                f"fields: {', '.join(list(missing_fields))}"
             )
 
         if user_account_info["type"] != "authorized_user":
@@ -149,14 +179,14 @@ class GCPUserAccountCredentials(AuthenticationConfig):
                 "instead of 'authorized_user'."
             )
 
-        return v
+        return value
 
 
 class GCPServiceAccountCredentials(AuthenticationConfig):
     """GCP service account credentials."""
 
-    service_account_json: SecretStr = Field(
-        title="GCP Service Account Key JSON",
+    service_account_json: PlainSerializedSecretStr = Field(
+        title="GCP Service Account Key JSON optionally base64 encoded.",
     )
 
     generate_temporary_tokens: bool = Field(
@@ -167,30 +197,51 @@ class GCPServiceAccountCredentials(AuthenticationConfig):
         "distribute the service account key JSON to clients instead.",
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
     def validate_service_account_dict(
-        cls, values: Dict[str, Any]
+        cls, data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Convert the service account credentials to JSON if given in dict format.
 
         Args:
-            values: The configuration values.
+            data: The configuration values.
 
         Returns:
             The validated configuration values.
+
+        Raises:
+            ValueError: If the service account credentials JSON is invalid.
         """
-        if isinstance(values.get("service_account_json"), dict):
-            values["service_account_json"] = json.dumps(
-                values["service_account_json"]
+        service_account_json = data.get("service_account_json")
+        if isinstance(service_account_json, dict):
+            data["service_account_json"] = json.dumps(
+                data["service_account_json"]
             )
-        return values
+        elif isinstance(service_account_json, str):
+            # Check if the service account JSON is base64 encoded and decode it
+            if re.match(r"^[A-Za-z0-9+/=]+$", service_account_json):
+                try:
+                    data["service_account_json"] = base64.b64decode(
+                        service_account_json
+                    ).decode("utf-8")
+                except Exception as e:
+                    raise ValueError(
+                        f"Failed to decode base64 encoded service account JSON: {e}"
+                    )
 
-    @validator("service_account_json")
-    def validate_service_account_json(cls, v: SecretStr) -> SecretStr:
+        return data
+
+    @field_validator("service_account_json")
+    @classmethod
+    def validate_service_account_json(
+        cls, value: PlainSerializedSecretStr
+    ) -> PlainSerializedSecretStr:
         """Validate the service account credentials JSON.
 
         Args:
-            v: The service account credentials JSON.
+            value: The service account credentials JSON.
 
         Returns:
             The validated service account credentials JSON.
@@ -199,7 +250,7 @@ class GCPServiceAccountCredentials(AuthenticationConfig):
             ValueError: If the service account credentials JSON is invalid.
         """
         try:
-            service_account_info = json.loads(v.get_secret_value())
+            service_account_info = json.loads(value.get_secret_value())
         except json.JSONDecodeError as e:
             raise ValueError(
                 f"GCP service account credentials is not a valid JSON: {e}"
@@ -225,7 +276,7 @@ class GCPServiceAccountCredentials(AuthenticationConfig):
         if missing_fields:
             raise ValueError(
                 f"GCP service account credentials JSON is missing required "
-                f'fields: {", ".join(list(missing_fields))}'
+                f"fields: {', '.join(list(missing_fields))}"
             )
 
         if service_account_info["type"] != "service_account":
@@ -235,13 +286,113 @@ class GCPServiceAccountCredentials(AuthenticationConfig):
                 "instead of 'service_account'."
             )
 
-        return v
+        return value
+
+
+class GCPExternalAccountCredentials(AuthenticationConfig):
+    """GCP external account credentials."""
+
+    external_account_json: PlainSerializedSecretStr = Field(
+        title="GCP External Account JSON optionally base64 encoded.",
+    )
+
+    generate_temporary_tokens: bool = Field(
+        default=True,
+        title="Generate temporary OAuth 2.0 tokens",
+        description="Whether to generate temporary OAuth 2.0 tokens from the "
+        "external account key JSON. If set to False, the connector will "
+        "distribute the external account JSON to clients instead.",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def validate_service_account_dict(
+        cls, data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Convert the external account credentials to JSON if given in dict format.
+
+        Args:
+            data: The configuration values.
+
+        Returns:
+            The validated configuration values.
+
+        Raises:
+            ValueError: If the external account credentials JSON is invalid.
+        """
+        external_account_json = data.get("external_account_json")
+        if isinstance(external_account_json, dict):
+            data["external_account_json"] = json.dumps(
+                data["external_account_json"]
+            )
+        elif isinstance(external_account_json, str):
+            # Check if the external account JSON is base64 encoded and decode it
+            if re.match(r"^[A-Za-z0-9+/=]+$", external_account_json):
+                try:
+                    data["external_account_json"] = base64.b64decode(
+                        external_account_json
+                    ).decode("utf-8")
+                except Exception as e:
+                    raise ValueError(
+                        f"Failed to decode base64 encoded external account JSON: {e}"
+                    )
+
+        return data
+
+    @field_validator("external_account_json")
+    @classmethod
+    def validate_external_account_json(
+        cls, value: PlainSerializedSecretStr
+    ) -> PlainSerializedSecretStr:
+        """Validate the external account credentials JSON.
+
+        Args:
+            value: The external account credentials JSON.
+
+        Returns:
+            The validated external account credentials JSON.
+
+        Raises:
+            ValueError: If the external account credentials JSON is invalid.
+        """
+        try:
+            external_account_info = json.loads(value.get_secret_value())
+        except json.JSONDecodeError as e:
+            raise ValueError(
+                f"GCP external account credentials is not a valid JSON: {e}"
+            )
+
+        # Check that all fields are present
+        required_fields = [
+            "type",
+            "subject_token_type",
+            "token_url",
+        ]
+        # Compute missing fields
+        missing_fields = set(required_fields) - set(
+            external_account_info.keys()
+        )
+        if missing_fields:
+            raise ValueError(
+                f"GCP external account credentials JSON is missing required "
+                f"fields: {', '.join(list(missing_fields))}"
+            )
+
+        if external_account_info["type"] != "external_account":
+            raise ValueError(
+                "The JSON does not contain GCP external account credentials. "
+                f'The "type" field is set to {external_account_info["type"]} '
+                "instead of 'external_account'."
+            )
+
+        return value
 
 
 class GCPOAuth2Token(AuthenticationConfig):
     """GCP OAuth 2.0 token credentials."""
 
-    token: SecretStr = Field(
+    token: PlainSerializedSecretStr = Field(
         title="GCP OAuth 2.0 Token",
     )
 
@@ -249,20 +400,72 @@ class GCPOAuth2Token(AuthenticationConfig):
 class GCPBaseConfig(AuthenticationConfig):
     """GCP base configuration."""
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        This method must be implemented by subclasses to ensure that the GCP
+        project ID is always available.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError
+
+
+class GCPBaseProjectIDConfig(GCPBaseConfig):
+    """GCP base configuration with included project ID."""
+
     project_id: str = Field(
         title="GCP Project ID where the target resource is located.",
     )
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
 
-class GCPUserAccountConfig(GCPBaseConfig, GCPUserAccountCredentials):
+        Returns:
+            The GCP project ID.
+        """
+        return self.project_id
+
+
+class GCPUserAccountConfig(GCPBaseProjectIDConfig, GCPUserAccountCredentials):
     """GCP user account configuration."""
 
 
 class GCPServiceAccountConfig(GCPBaseConfig, GCPServiceAccountCredentials):
     """GCP service account configuration."""
 
+    _project_id: Optional[str] = None
 
-class GCPOAuth2TokenConfig(GCPBaseConfig, GCPOAuth2Token):
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        When a service account JSON is provided, the project ID can be extracted
+        from it instead of being provided explicitly.
+
+        Returns:
+            The GCP project ID.
+        """
+        if self._project_id is None:
+            self._project_id = json.loads(
+                self.service_account_json.get_secret_value()
+            )["project_id"]
+            # Guaranteed by the field validator
+            assert self._project_id is not None
+
+        return self._project_id
+
+
+class GCPExternalAccountConfig(
+    GCPBaseProjectIDConfig, GCPExternalAccountCredentials
+):
+    """GCP external account configuration."""
+
+
+class GCPOAuth2TokenConfig(GCPBaseProjectIDConfig, GCPOAuth2Token):
     """GCP OAuth 2.0 configuration."""
 
     service_account_email: Optional[str] = Field(
@@ -288,10 +491,177 @@ class GCPAuthenticationMethods(StrEnum):
     IMPLICIT = "implicit"
     USER_ACCOUNT = "user-account"
     SERVICE_ACCOUNT = "service-account"
+    EXTERNAL_ACCOUNT = "external-account"
     OAUTH2_TOKEN = "oauth2-token"
     IMPERSONATION = "impersonation"
 
 
+try:
+    from google.auth.aws import _DefaultAwsSecurityCredentialsSupplier
+
+    class ZenMLAwsSecurityCredentialsSupplier(
+        _DefaultAwsSecurityCredentialsSupplier  # type: ignore[misc]
+    ):
+        """An improved version of the GCP external account credential supplier for AWS.
+
+        The original GCP external account credential supplier only provides
+        rudimentary support for extracting AWS credentials from environment
+        variables or the AWS metadata service. This version improves on that by
+        using the boto3 library itself (if available), which uses the entire range
+        of implicit authentication features packed into it.
+
+        Without this improvement, `sts.AssumeRoleWithWebIdentity` authentication is
+        not supported for EKS pods and the EC2 attached role credentials are
+        used instead (see: https://medium.com/@derek10cloud/gcp-workload-identity-federation-doesnt-yet-support-eks-irsa-in-aws-a3c71877671a).
+        """
+
+        def get_aws_security_credentials(
+            self, context: Any, request: Any
+        ) -> gcp_aws.AwsSecurityCredentials:
+            """Get the security credentials from the local environment.
+
+            This method is a copy of the original method from the
+            `google.auth.aws._DefaultAwsSecurityCredentialsSupplier` class. It has
+            been modified to use the boto3 library to extract the AWS credentials
+            from the local environment.
+
+            Args:
+                context: The context to use to get the security credentials.
+                request: The request to use to get the security credentials.
+
+            Returns:
+                The AWS temporary security credentials.
+            """
+            try:
+                import boto3
+
+                session = boto3.Session()
+                credentials = session.get_credentials()
+                if credentials is not None:
+                    creds = credentials.get_frozen_credentials()
+                    return gcp_aws.AwsSecurityCredentials(
+                        creds.access_key,
+                        creds.secret_key,
+                        creds.token,
+                    )
+            except ImportError:
+                pass
+
+            logger.debug(
+                "Failed to extract AWS credentials from the local environment "
+                "using the boto3 library. Falling back to the original "
+                "implementation."
+            )
+
+            return super().get_aws_security_credentials(context, request)
+
+        def get_aws_region(self, context: Any, request: Any) -> str:
+            """Get the AWS region from the local environment.
+
+            This method is a copy of the original method from the
+            `google.auth.aws._DefaultAwsSecurityCredentialsSupplier` class. It has
+            been modified to use the boto3 library to extract the AWS
+            region from the local environment.
+
+            Args:
+                context: The context to use to get the security credentials.
+                request: The request to use to get the security credentials.
+
+            Returns:
+                The AWS region.
+            """
+            try:
+                import boto3
+
+                session = boto3.Session()
+                if session.region_name:
+                    return session.region_name  # type: ignore[no-any-return]
+            except ImportError:
+                pass
+
+            logger.debug(
+                "Failed to extract AWS region from the local environment "
+                "using the boto3 library. Falling back to the original "
+                "implementation."
+            )
+
+            return super().get_aws_region(  # type: ignore[no-any-return]
+                context, request
+            )
+
+except ImportError:
+    # The `google.auth.aws._DefaultAwsSecurityCredentialsSupplier`
+    # class has been introduced in the `google-auth` library version 2.29.0.
+    # Before that, the AWS logic was part of the `google.auth.awsCredentials`
+    # class itself.
+    ZenMLAwsSecurityCredentialsSupplier = None  # type: ignore[assignment,misc]
+    pass
+
+
+class ZenMLGCPAWSExternalAccountCredentials(gcp_aws.Credentials):  # type: ignore[misc]
+    """An improved version of the GCP external account credential for AWS.
+
+    The original GCP external account credential only provides rudimentary
+    support for extracting AWS credentials from environment variables or the
+    AWS metadata service. This version improves on that by using the boto3
+    library itself (if available), which uses the entire range of implicit
+    authentication features packed into it.
+
+    Without this improvement, `sts.AssumeRoleWithWebIdentity` authentication is
+    not supported for EKS pods and the EC2 attached role credentials are
+    used instead (see: https://medium.com/@derek10cloud/gcp-workload-identity-federation-doesnt-yet-support-eks-irsa-in-aws-a3c71877671a).
+
+    IMPORTANT: subclassing this class only works with the `google-auth` library
+    version lower than 2.29.0. Starting from version 2.29.0, the AWS logic
+    has been moved to a separate `google.auth.aws._DefaultAwsSecurityCredentialsSupplier`
+    class that can be subclassed instead and supplied as the
+    `aws_security_credentials_supplier` parameter to the
+    `google.auth.aws.Credentials` class.
+    """
+
+    def _get_security_credentials(
+        self, request: Any, imdsv2_session_token: Any
+    ) -> Dict[str, Any]:
+        """Get the security credentials from the local environment.
+
+        This method is a copy of the original method from the
+        `google.auth._default` module. It has been modified to use the boto3
+        library to extract the AWS credentials from the local environment.
+
+        Args:
+            request: The request to use to get the security credentials.
+            imdsv2_session_token: The IMDSv2 session token to use to get the
+                security credentials.
+
+        Returns:
+            The AWS temporary security credentials.
+        """
+        try:
+            import boto3
+
+            session = boto3.Session()
+            credentials = session.get_credentials()
+            if credentials is not None:
+                creds = credentials.get_frozen_credentials()
+                return {
+                    "access_key_id": creds.access_key,
+                    "secret_access_key": creds.secret_key,
+                    "security_token": creds.token,
+                }
+        except ImportError:
+            pass
+
+        logger.debug(
+            "Failed to extract AWS credentials from the local environment "
+            "using the boto3 library. Falling back to the original "
+            "implementation."
+        )
+
+        return super()._get_security_credentials(  # type: ignore[no-any-return]
+            request, imdsv2_session_token
+        )
+
+
 GCP_SERVICE_CONNECTOR_TYPE_SPEC = ServiceConnectorTypeModel(
     name="GCP Service Connector",
     connector_type=GCP_CONNECTOR_TYPE,
@@ -378,7 +748,7 @@ resources in the specified project. When used remotely in a GCP workload, the
 configured project has to be the same as the project of the attached service
 account.
 """,
-            config_class=GCPBaseConfig,
+            config_class=GCPBaseProjectIDConfig,
         ),
         AuthenticationMethodModel(
             name="GCP User Account",
@@ -435,6 +805,46 @@ picked up when auto-configuration is used.
 """,
             config_class=GCPServiceAccountConfig,
         ),
+        AuthenticationMethodModel(
+            name="GCP External Account",
+            auth_method=GCPAuthenticationMethods.EXTERNAL_ACCOUNT,
+            description="""
+Use [GCP workload identity federation](https://cloud.google.com/iam/docs/workload-identity-federation)
+to authenticate to GCP services using AWS IAM credentials, Azure Active
+Directory credentials or generic OIDC tokens.
+
+This authentication method only requires a GCP workload identity external
+account JSON file that only contains the configuration for the external account
+without any sensitive credentials. It allows implementing a two layer
+authentication scheme that keeps the set of permissions associated with implicit
+credentials down to the bare minimum and grants permissions to the
+privilege-bearing GCP service account instead.
+
+This authentication method can be used to authenticate to GCP services using
+credentials from other cloud providers or identity providers. When used with
+workloads running on AWS or Azure, it involves automatically picking up
+credentials from the AWS IAM or Azure AD identity associated with the workload
+and using them to authenticate to GCP services. This means that the result
+depends on the environment where the ZenML server is deployed and is thus not
+fully reproducible.
+
+By default, the GCP connector generates temporary OAuth 2.0 tokens from the
+external account credentials and distributes them to clients. The tokens have a
+limited lifetime of 1 hour. This behavior can be disabled by setting the
+`generate_temporary_tokens` configuration option to `False`, in which case, the
+connector will distribute the external account credentials JSON to clients
+instead (not recommended).
+
+A GCP project is required and the connector may only be used to access GCP
+resources in the specified project. This project must be the same as the one
+for which the external account was configured.
+
+If you already have the GOOGLE_APPLICATION_CREDENTIALS environment variable
+configured to point to an external account key JSON file, it will be
+automatically picked up when auto-configuration is used.
+""",
+            config_class=GCPExternalAccountConfig,
+        ),
         AuthenticationMethodModel(
             name="GCP Oauth 2.0 Token",
             auth_method=GCPAuthenticationMethods.OAUTH2_TOKEN,
@@ -584,14 +994,53 @@ GKE clusters in the GCP project that it is configured to use.
             emoji=":cyclone:",
         ),
         ResourceTypeModel(
-            name="GCP GCR container registry",
+            name="GCP GAR container registry",
             resource_type=DOCKER_REGISTRY_RESOURCE_TYPE,
             description="""
-Allows Stack Components to access a GCR registry as a standard
+Allows Stack Components to access a Google Artifact Registry as a standard
 Docker registry resource. When used by Stack Components, they are provided a
 pre-authenticated Python Docker client instance.
 
-The configured credentials must have at least the following [GCP permissions](https://cloud.google.com/iam/docs/permissions-reference):
+The configured credentials must have at least the following [GCP permissions](https://cloud.google.com/iam/docs/understanding-roles#artifact-registry-roles):
+
+- `artifactregistry.repositories.createOnPush`
+- `artifactregistry.repositories.downloadArtifacts`
+- `artifactregistry.repositories.get`
+- `artifactregistry.repositories.list`
+- `artifactregistry.repositories.readViaVirtualRepository`
+- `artifactregistry.repositories.uploadArtifacts`
+- `artifactregistry.locations.list`
+
+The Artifact Registry Create-on-Push Writer role includes all of the above
+permissions.
+
+This resource type also includes legacy GCR container registry support.
+
+**Important Notice: Google Container Registry** [**is being replaced by Artifact Registry**](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr).
+Please start using Artifact Registry for your containers. As per Google's
+documentation, *"after May 15, 2024, Artifact Registry will host images for the
+gcr.io domain in Google Cloud projects without previous Container Registry
+usage. After March 18, 2025, Container Registry will be shut down."*.
+
+Support for legacy GCR registries is still included in the GCP service
+connector. Users that already have GCP service connectors configured to access
+GCR registries may continue to use them without taking any action. However, it
+is recommended to transition to Google Artifact Registries as soon as possible
+by following [the GCP guide on this subject](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr)
+and making the following updates to ZenML GCP Service Connectors that are used
+to access GCR resources:
+
+* add the IAM permissions documented here to the GCP Service Connector
+credentials to enable them to access the Artifact Registries.
+* users may keep the gcr.io GCR URLs already configured in the GCP Service
+Connectors as well as those used in linked Container Registry stack components
+given that these domains are redirected by Google to GAR as covered in the GCR
+transition guide. Alternatively, users may update the GCP Service Connector
+configuration and/or the Container Registry stack components to use the
+replacement Artifact Registry URLs.
+
+When used with GCR registries, the configured credentials must have at least the
+following [GCP permissions](https://cloud.google.com/iam/docs/understanding-roles#cloud-storage-roles):
 
 - `storage.buckets.get`
 - `storage.multipartUploads.abort`
@@ -605,17 +1054,21 @@ The configured credentials must have at least the following [GCP permissions](ht
 The Storage Legacy Bucket Writer role includes all of the above permissions
 while at the same time restricting access to only the GCR buckets.
 
-The resource name associated with this resource type identifies the GCR
-container registry associated with the GCP project (the repository name is
-optional):
+If set, the resource name must identify a GAR or GCR registry using one of the
+following formats:
+
+- Google Artifact Registry repository URI: `[https://]<region>-docker.pkg.dev/<project-id>/<registry-id>[/<repository-name>]`
+- Google Artifact Registry name: `projects/<project-id>/locations/<location>/repositories/<repository-id>`
+- (legacy) GCR repository URI: `[https://][us.|eu.|asia.]gcr.io/<project-id>[/<repository-name>]`
 
-- GCR repository URI: `[https://]gcr.io/{project-id}[/{repository-name}]
+The connector can only be used to access GAR and GCR registries in the GCP
+project that it is configured to use.
 """,
             auth_methods=GCPAuthenticationMethods.values(),
-            # Does not support instances, given that the connector
-            # provides access to the entire GCR container registry
-            # for the configured GCP project.
-            supports_instances=False,
+            # The connector provides access to the entire GCR container registry
+            # for the configured GCP project as well as any number of artifact
+            # registry repositories.
+            supports_instances=True,
             logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/container_registry/docker.png",
             emoji=":whale:",
         ),
@@ -650,7 +1103,7 @@ class GCPServiceConnector(ServiceConnector):
         auth_method: str,
         resource_type: Optional[str] = None,
         resource_id: Optional[str] = None,
-    ) -> Tuple[gcp_service_account.Credentials, Optional[datetime.datetime]]:
+    ) -> Tuple[gcp_credentials.Credentials, Optional[datetime.datetime]]:
         """Get a GCP session object with credentials for the specified resource.
 
         Args:
@@ -673,7 +1126,10 @@ class GCPServiceConnector(ServiceConnector):
             # Refresh expired sessions
             now = datetime.datetime.now(datetime.timezone.utc)
             expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
-            if expires_at > now:
+            # check if the token expires in the near future
+            if expires_at > now + datetime.timedelta(
+                minutes=GCP_SESSION_EXPIRATION_BUFFER
+            ):
                 return session, expires_at
 
         logger.debug(
@@ -711,7 +1167,10 @@ class GCPServiceConnector(ServiceConnector):
         auth_method: str,
         resource_type: Optional[str] = None,
         resource_id: Optional[str] = None,
-    ) -> Tuple[gcp_service_account.Credentials, Optional[datetime.datetime]]:
+    ) -> Tuple[
+        gcp_credentials.Credentials,
+        Optional[datetime.datetime],
+    ]:
         """Authenticate to GCP and return a session with credentials.
 
         Args:
@@ -762,8 +1221,49 @@ class GCPServiceConnector(ServiceConnector):
                         scopes=scopes,
                     )
                 )
+            elif auth_method == GCPAuthenticationMethods.EXTERNAL_ACCOUNT:
+                self._check_implicit_auth_method_allowed()
+
+                assert isinstance(cfg, GCPExternalAccountConfig)
+
+                # As a special case, for the AWS external account credential,
+                # we use a custom credential class that supports extracting
+                # the AWS credentials from the local environment, metadata
+                # service or IRSA (if running on AWS EKS).
+                account_info = json.loads(
+                    cfg.external_account_json.get_secret_value()
+                )
+                if (
+                    account_info.get("subject_token_type")
+                    == _AWS_SUBJECT_TOKEN_TYPE
+                ):
+                    if ZenMLAwsSecurityCredentialsSupplier is not None:
+                        account_info["aws_security_credentials_supplier"] = (
+                            ZenMLAwsSecurityCredentialsSupplier(
+                                account_info.pop("credential_source"),
+                            )
+                        )
+                    credentials = (
+                        ZenMLGCPAWSExternalAccountCredentials.from_info(
+                            account_info,
+                            scopes=scopes,
+                        )
+                    )
+                else:
+                    credentials, _ = _get_external_account_credentials(
+                        json.loads(
+                            cfg.external_account_json.get_secret_value()
+                        ),
+                        filename="",  # Not used
+                        scopes=scopes,
+                    )
+
             else:
+                # Service account or impersonation (which is a special case of
+                # service account authentication)
+
                 assert isinstance(cfg, GCPServiceAccountConfig)
+
                 credentials = (
                     gcp_service_account.Credentials.from_service_account_info(
                         json.loads(
@@ -773,22 +1273,24 @@ class GCPServiceConnector(ServiceConnector):
                     )
                 )
 
-            if auth_method == GCPAuthenticationMethods.IMPERSONATION:
-                assert isinstance(cfg, GCPServiceAccountImpersonationConfig)
-
-                try:
-                    credentials = gcp_impersonated_credentials.Credentials(
-                        source_credentials=credentials,
-                        target_principal=cfg.target_principal,
-                        target_scopes=scopes,
-                        lifetime=self.expiration_seconds,
-                    )
-                except google.auth.exceptions.GoogleAuthError as e:
-                    raise AuthorizationException(
-                        f"Failed to impersonate service account "
-                        f"'{cfg.target_principal}': {e}"
+                if auth_method == GCPAuthenticationMethods.IMPERSONATION:
+                    assert isinstance(
+                        cfg, GCPServiceAccountImpersonationConfig
                     )
 
+                    try:
+                        credentials = gcp_impersonated_credentials.Credentials(
+                            source_credentials=credentials,
+                            target_principal=cfg.target_principal,
+                            target_scopes=scopes,
+                            lifetime=self.expiration_seconds,
+                        )
+                    except google.auth.exceptions.GoogleAuthError as e:
+                        raise AuthorizationException(
+                            f"Failed to impersonate service account "
+                            f"'{cfg.target_principal}': {e}"
+                        )
+
         if not credentials.valid:
             try:
                 with requests.Session() as session:
@@ -849,60 +1351,101 @@ class GCPServiceConnector(ServiceConnector):
 
         return bucket_name
 
-    def _parse_gcr_resource_id(
+    def _parse_gar_resource_id(
         self,
         resource_id: str,
-    ) -> str:
-        """Validate and convert an GCR resource ID to an GCR registry ID.
+    ) -> Tuple[str, Optional[str]]:
+        """Validate and convert a GAR resource ID to a Google Artifact Registry ID and name.
 
         Args:
             resource_id: The resource ID to convert.
 
         Returns:
-            The GCR registry ID.
+            The Google Artifact Registry ID and name. The name is omitted if the
+            resource ID is a GCR repository URI.
 
         Raises:
-            ValueError: If the provided resource ID is not a valid GCR
-                repository URI.
+            ValueError: If the provided resource ID is not a valid GAR
+                or GCR repository URI.
         """
         # The resource ID could mean different things:
         #
-        # - an GCR repository URI
+        # - a GAR repository URI
+        # - a GAR repository name
+        # - a GCR repository URI (backwards-compatibility)
         #
         # We need to extract the project ID and registry ID from
         # the provided resource ID
-        config_project_id = self.config.project_id
+        config_project_id = self.config.gcp_project_id
         project_id: Optional[str] = None
-        # A GCR repository URI uses one of several hostnames (gcr.io, us.gcr.io,
-        # eu.gcr.io, asia.gcr.io etc.) and the project ID is the first part of
-        # the URL path
-        if re.match(
-            r"^(https://)?([a-z]+.)*gcr.io/[a-z0-9-]+(/.+)*$",
+        canonical_url: str
+        registry_name: Optional[str] = None
+
+        # A Google Artifact Registry URI uses the <location>-docker-pkg.dev
+        # domain format with the project ID as the first part of the URL path
+        # and the registry name as the second part of the URL path
+        if match := re.match(
+            r"^(https://)?(([a-z0-9-]+)-docker\.pkg\.dev/([a-z0-9-]+)/([a-z0-9-.]+))(/.+)*$",
             resource_id,
         ):
-            # The resource ID is a GCR repository URI
-            if resource_id.startswith("https://"):
-                project_id = resource_id.split("/")[3]
-            else:
-                project_id = resource_id.split("/")[1]
+            # The resource ID is a Google Artifact Registry URI
+            project_id = match[4]
+            location = match[3]
+            repository = match[5]
+
+            # Return the GAR URL without the image name and without the protocol
+            canonical_url = match[2]
+            registry_name = f"projects/{project_id}/locations/{location}/repositories/{repository}"
+
+        # Alternatively, the Google Artifact Registry name uses the
+        # projects/<project-id>/locations/<location>/repositories/<repository-id>
+        # format
+        elif match := re.match(
+            r"^projects/([a-z0-9-]+)/locations/([a-z0-9-]+)/repositories/([a-z0-9-.]+)$",
+            resource_id,
+        ):
+            # The resource ID is a Google Artifact Registry name
+            project_id = match[1]
+            location = match[2]
+            repository = match[3]
+
+            # Return the GAR URL
+            canonical_url = (
+                f"{location}-docker.pkg.dev/{project_id}/{repository}"
+            )
+            registry_name = resource_id
+
+        # A legacy GCR repository URI uses one of several hostnames (gcr.io,
+        # us.gcr.io, eu.gcr.io, asia.gcr.io) and the project ID is the
+        # first part of the URL path
+        elif match := re.match(
+            r"^(https://)?(((us|eu|asia)\.)?gcr\.io/[a-z0-9-]+)(/.+)*$",
+            resource_id,
+        ):
+            # The resource ID is a legacy GCR repository URI.
+            # Return the GAR URL without the image name and without the protocol
+            canonical_url = match[2]
+
         else:
             raise ValueError(
-                f"Invalid resource ID for a GCR registry: {resource_id}. "
-                f"Supported formats are:\n"
+                f"Invalid resource ID for a Google Artifact Registry: "
+                f"{resource_id}. Supported formats are:\n"
+                f"Google Artifact Registry URI: [https://]<region>-docker.pkg.dev/<project-id>/<registry-id>[/<repository-name>]\n"
+                f"Google Artifact Registry name: projects/<project-id>/locations/<location>/repositories/<repository-id>\n"
                 f"GCR repository URI: [https://][us.|eu.|asia.]gcr.io/<project-id>[/<repository-name>]"
             )
 
         # If the connector is configured with a project and the resource ID
-        # is an GCR repository URI that specifies a different project,
+        # is a GAR repository URI that specifies a different project,
         # we raise an error
         if project_id and project_id != config_project_id:
             raise ValueError(
-                f"The GCP project for the {resource_id} GCR repository "
-                f"'{project_id}' does not match the project configured in "
-                f"the connector: '{config_project_id}'."
+                f"The GCP project for the {resource_id} Google Artifact "
+                f"Registry '{project_id}' does not match the project "
+                f"configured in the connector: '{config_project_id}'."
             )
 
-        return f"gcr.io/{project_id}"
+        return canonical_url, registry_name
 
     def _parse_gke_resource_id(self, resource_id: str) -> str:
         """Validate and convert an GKE resource ID to a GKE cluster name.
@@ -951,7 +1494,7 @@ class GCPServiceConnector(ServiceConnector):
             cluster_name = self._parse_gke_resource_id(resource_id)
             return cluster_name
         elif resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            registry_id = self._parse_gcr_resource_id(
+            registry_id, _ = self._parse_gar_resource_id(
                 resource_id,
             )
             return registry_id
@@ -975,9 +1518,7 @@ class GCPServiceConnector(ServiceConnector):
                 authorized.
         """
         if resource_type == GCP_RESOURCE_TYPE:
-            return self.config.project_id
-        elif resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            return f"gcr.io/{self.config.project_id}"
+            return self.config.gcp_project_id
 
         raise RuntimeError(
             f"Default resource ID not supported for '{resource_type}' resource "
@@ -1034,7 +1575,7 @@ class GCPServiceConnector(ServiceConnector):
 
             # Create an GCS client for the bucket
             client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             return client
 
@@ -1074,7 +1615,7 @@ class GCPServiceConnector(ServiceConnector):
 
             # There is no way to configure the local gcloud CLI to use
             # temporary OAuth 2.0 tokens. However, we can configure it to use
-            # the service account credentials
+            # the service account or external account credentials
             if self.auth_method == GCPAuthenticationMethods.SERVICE_ACCOUNT:
                 assert isinstance(self.config, GCPServiceAccountConfig)
                 # Use the service account credentials JSON to configure the
@@ -1082,6 +1623,13 @@ class GCPServiceConnector(ServiceConnector):
                 gcloud_config_json = (
                     self.config.service_account_json.get_secret_value()
                 )
+            elif self.auth_method == GCPAuthenticationMethods.EXTERNAL_ACCOUNT:
+                assert isinstance(self.config, GCPExternalAccountConfig)
+                # Use the external account credentials JSON to configure the
+                # local gcloud CLI
+                gcloud_config_json = (
+                    self.config.external_account_json.get_secret_value()
+                )
 
             if gcloud_config_json:
                 from google.auth import _cloud_sdk
@@ -1108,6 +1656,7 @@ class GCPServiceConnector(ServiceConnector):
                                 "gcloud",
                                 "auth",
                                 "login",
+                                "--quiet",
                                 "--cred-file",
                                 adc_path,
                             ],
@@ -1132,7 +1681,7 @@ class GCPServiceConnector(ServiceConnector):
                             "config",
                             "set",
                             "project",
-                            self.config.project_id,
+                            self.config.gcp_project_id,
                         ],
                         check=True,
                         stderr=subprocess.STDOUT,
@@ -1163,9 +1712,10 @@ class GCPServiceConnector(ServiceConnector):
             raise NotImplementedError(
                 f"Local gcloud client configuration for resource type "
                 f"{resource_type} is only supported if the "
-                f"'{GCPAuthenticationMethods.SERVICE_ACCOUNT}' authentication "
-                f"method is used and only if the generation of temporary OAuth "
-                f"2.0 tokens is disabled by setting the "
+                f"'{GCPAuthenticationMethods.SERVICE_ACCOUNT}' or "
+                f"'{GCPAuthenticationMethods.EXTERNAL_ACCOUNT}' "
+                f"authentication method is used and only if the generation of "
+                f"temporary OAuth 2.0 tokens is disabled by setting the "
                 f"'generate_temporary_tokens' option to 'False' in the "
                 f"service connector configuration."
             )
@@ -1235,9 +1785,7 @@ class GCPServiceConnector(ServiceConnector):
             )
 
         if auth_method == GCPAuthenticationMethods.IMPLICIT:
-            cls._check_implicit_auth_method_allowed()
-
-            auth_config = GCPBaseConfig(
+            auth_config = GCPBaseProjectIDConfig(
                 project_id=project_id,
             )
         elif auth_method == GCPAuthenticationMethods.OAUTH2_TOKEN:
@@ -1339,6 +1887,51 @@ class GCPServiceConnector(ServiceConnector):
                     project_id=project_id,
                     service_account_json=service_account_json,
                 )
+            # Check if external account credentials are available
+            elif isinstance(credentials, gcp_external_account.Credentials):
+                if auth_method not in [
+                    GCPAuthenticationMethods.EXTERNAL_ACCOUNT,
+                    None,
+                ]:
+                    raise NotImplementedError(
+                        f"Could not perform auto-configuration for "
+                        f"authentication method {auth_method}. Only "
+                        f"GCP external account credentials have been detected."
+                    )
+
+                auth_method = GCPAuthenticationMethods.EXTERNAL_ACCOUNT
+                external_account_json_file = os.environ.get(
+                    "GOOGLE_APPLICATION_CREDENTIALS"
+                )
+                if external_account_json_file is None:
+                    # No explicit service account JSON file was specified in the
+                    # environment, meaning that the credentials were loaded from
+                    # the GCP application default credentials (ADC) file.
+                    from google.auth import _cloud_sdk
+
+                    # Use the location of the gcloud application default
+                    # credentials file
+                    external_account_json_file = (
+                        _cloud_sdk.get_application_default_credentials_path()
+                    )
+
+                if not external_account_json_file or not os.path.isfile(
+                    external_account_json_file
+                ):
+                    raise AuthorizationException(
+                        "No GCP service account credentials were found in the "
+                        "environment or the application default credentials "
+                        "path. Please set the GOOGLE_APPLICATION_CREDENTIALS "
+                        "environment variable to the path of the external "
+                        "account JSON file or run 'gcloud auth application-"
+                        "default login' to generate a new ADC file."
+                    )
+                with open(external_account_json_file, "r") as f:
+                    external_account_json = f.read()
+                auth_config = GCPExternalAccountConfig(
+                    project_id=project_id,
+                    external_account_json=external_account_json,
+                )
             else:
                 raise AuthorizationException(
                     "No valid GCP credentials could be detected."
@@ -1401,7 +1994,7 @@ class GCPServiceConnector(ServiceConnector):
 
         if resource_type == GCS_RESOURCE_TYPE:
             gcs_client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             if not resource_id:
                 # List all GCS buckets
@@ -1426,11 +2019,87 @@ class GCPServiceConnector(ServiceConnector):
                     raise AuthorizationException(msg) from e
 
         if resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            assert resource_id is not None
+            # Get a GAR client
+            gar_client = artifactregistry_v1.ArtifactRegistryClient(
+                credentials=credentials
+            )
 
-            # No way to verify a GCR registry without attempting to
-            # connect to it via Docker/OCI, so just return the resource ID.
-            return [resource_id]
+            if resource_id:
+                registry_id, registry_name = self._parse_gar_resource_id(
+                    resource_id
+                )
+
+                if registry_name is None:
+                    # This is a legacy GCR repository URI. We can't verify
+                    # the repository access without attempting to connect to it
+                    # via Docker/OCI, so just return the resource ID.
+                    return [registry_id]
+
+                # Check if the specified GAR registry exists
+                try:
+                    repository = gar_client.get_repository(
+                        name=registry_name,
+                    )
+                    if repository.format_.name != "DOCKER":
+                        raise AuthorizationException(
+                            f"Google Artifact Registry '{resource_id}' is not a "
+                            "Docker registry."
+                        )
+                    return [registry_id]
+                except google.api_core.exceptions.GoogleAPIError as e:
+                    msg = f"Failed to fetch Google Artifact Registry '{registry_id}': {e}"
+                    logger.error(msg)
+                    raise AuthorizationException(msg) from e
+
+            # For backwards compatibility, we initialize the list of resource
+            # IDs with all GCR supported registries for the configured GCP
+            # project
+            resource_ids: List[str] = [
+                f"{location}gcr.io/{self.config.gcp_project_id}"
+                for location in ["", "us.", "eu.", "asia."]
+            ]
+
+            # List all Google Artifact Registries
+            try:
+                # First, we need to fetch all the Artifact Registry supported
+                # locations
+                locations = gar_client.list_locations(
+                    request=locations_pb2.ListLocationsRequest(
+                        name=f"projects/{self.config.gcp_project_id}"
+                    )
+                )
+                location_names = [
+                    locations.locations[i].location_id
+                    for i in range(len(locations.locations))
+                ]
+
+                # Then, we need to fetch all the repositories in each location
+                repository_names: List[str] = []
+                for location in location_names:
+                    repositories = gar_client.list_repositories(
+                        parent=f"projects/{self.config.gcp_project_id}/locations/{location}"
+                    )
+                    repository_names.extend(
+                        [
+                            repository.name
+                            for repository in repositories
+                            if repository.format_.name == "DOCKER"
+                        ]
+                    )
+
+                for repository_name in repository_names:
+                    # Convert the repository name to a canonical GAR URL
+                    resource_ids.append(
+                        self._parse_gar_resource_id(repository_name)[0]
+                    )
+
+            except google.api_core.exceptions.GoogleAPIError as e:
+                msg = f"Failed to list Google Artifact Registries: {e}"
+                logger.error(msg)
+                # TODO: enable when GCR is no longer supported:
+                # raise AuthorizationException(msg) from e
+
+            return resource_ids
 
         if resource_type == KUBERNETES_CLUSTER_RESOURCE_TYPE:
             gke_client = container_v1.ClusterManagerClient(
@@ -1440,7 +2109,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_names = [cluster.name for cluster in clusters.clusters]
             except google.api_core.exceptions.GoogleAPIError as e:
@@ -1514,7 +2183,7 @@ class GCPServiceConnector(ServiceConnector):
             # object
             auth_method: str = GCPAuthenticationMethods.OAUTH2_TOKEN
             config: GCPBaseConfig = GCPOAuth2TokenConfig(
-                project_id=self.config.project_id,
+                project_id=self.config.gcp_project_id,
                 token=credentials.token,
                 service_account_email=credentials.signer_email
                 if hasattr(credentials, "signer_email")
@@ -1535,6 +2204,12 @@ class GCPServiceConnector(ServiceConnector):
                     config = self.config
                     auth_method = self.auth_method
                     expires_at = None
+            elif self.auth_method == GCPAuthenticationMethods.EXTERNAL_ACCOUNT:
+                assert isinstance(self.config, GCPExternalAccountConfig)
+                if not self.config.generate_temporary_tokens:
+                    config = self.config
+                    auth_method = self.auth_method
+                    expires_at = None
 
             # Create a client-side GCP connector instance that is fully formed
             # and ready to use to connect to the specified resource (i.e. has
@@ -1553,7 +2228,7 @@ class GCPServiceConnector(ServiceConnector):
         if resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
             assert resource_id is not None
 
-            registry_id = self._parse_gcr_resource_id(resource_id)
+            registry_id, _ = self._parse_gar_resource_id(resource_id)
 
             # Create a client-side Docker connector instance with the temporary
             # Docker credentials
@@ -1582,7 +2257,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_map = {
                     cluster.name: cluster for cluster in clusters.clusters
@@ -1626,7 +2301,7 @@ class GCPServiceConnector(ServiceConnector):
                 auth_method=KubernetesAuthenticationMethods.TOKEN,
                 resource_type=resource_type,
                 config=KubernetesTokenConfig(
-                    cluster_name=f"gke_{self.config.project_id}_{cluster_name}",
+                    cluster_name=f"gke_{self.config.gcp_project_id}_{cluster_name}",
                     certificate_authority=cluster_ca_cert,
                     server=f"https://{cluster_server}",
                     token=bearer_token,