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

zenml/zen_server/deploy/helm/templates/server-secret.yaml

@@ -1,7 +1,6 @@
 {{- define "secret.content" -}}
 {{- $prevServerSecret := (lookup "v1" "Secret" .Release.Namespace (include "zenml.fullname" .)) -}}
 data:
-  ZENML_DEFAULT_USER_PASSWORD: {{ .Values.zenml.defaultPassword | b64enc | quote }}
   {{- if .Values.zenml.jwtSecretKey }}
   ZENML_SERVER_JWT_SECRET_KEY: {{ .Values.zenml.jwtSecretKey | b64enc | quote }}
   {{- else if .Values.zenml.auth.jwtSecretKey }}
@@ -11,17 +10,19 @@ data:
   {{- else }}
   ZENML_SERVER_JWT_SECRET_KEY: {{ $prevServerSecret.data.ZENML_SERVER_JWT_SECRET_KEY | default (randAlphaNum 32 | b64enc | quote) }}
   {{- end }}
-  {{- if .Values.zenml.database.url }}
-  ZENML_STORE_URL: {{ .Values.zenml.database.url | b64enc | quote }}
-  {{- if .Values.zenml.database.sslCa }}
-  ZENML_STORE_SSL_CA: {{ .Files.Get .Values.zenml.database.sslCa | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslCert }}
-  ZENML_STORE_SSL_CERT: {{ .Files.Get .Values.zenml.database.sslCert | b64enc }}
+  
+  {{- if .Values.zenml.pro.enabled }}
+  {{- if .Values.zenml.pro.enrollmentKey }}
+  ZENML_SERVER_PRO_OAUTH2_CLIENT_SECRET: {{ .Values.zenml.pro.enrollmentKey | b64enc | quote }}
+  {{- else if or .Release.IsInstall (not $prevServerSecret) }}
+  ZENML_SERVER_PRO_OAUTH2_CLIENT_SECRET: {{ randAlphaNum 64 | b64enc | quote }}
+  {{- else }}
+  ZENML_SERVER_PRO_OAUTH2_CLIENT_SECRET: {{ $prevServerSecret.data.ZENML_SERVER_PRO_OAUTH2_CLIENT_SECRET | default (randAlphaNum 64 | b64enc | quote) }}
   {{- end }}
-  {{- if .Values.zenml.database.sslKey }}
-  ZENML_STORE_SSL_KEY: {{ .Files.Get .Values.zenml.database.sslKey | b64enc }}
   {{- end }}
+
+  {{- range $k, $v := include "zenml.storeSecretEnvVariables" . | fromYaml}}
+  {{ $k }}: {{ $v | b64enc | quote }}
   {{- end }}
   {{- range $k, $v := include "zenml.secretsStoreSecretEnvVariables" . | fromYaml}}
   {{ $k }}: {{ $v | b64enc | quote }}

zenml/zen_server/deploy/helm/values.yaml

@@ -10,12 +10,74 @@ zenml:
 
   replicaCount: 1
 
+  # The number of ZenML server worker threads to use. This controls the number
+  # of concurrent requests that each ZenML server pod/replica can handle at a
+  # time. If not specified, the default value is 40.
+  #
+  # NOTE: this value should be coordinated with the `zenml.database.poolSize`
+  # and `zenml.database.maxOverflow` values to ensure that the ZenML server
+  # workers do not block on database connections (i.e. the sum of the pool size
+  # and max overflow should be greater than or equal to the thread pool size).
+  #
+  # threadPoolSize: 40
+
   image:
     repository: zenmldocker/zenml-server
     pullPolicy: Always
     # Overrides the image tag whose default is the chart appVersion.
     tag:
 
+  # ZenML Pro related options.
+  pro:
+    # Set `enabled` to true to enable ZenML Pro servers. If set, some of the
+    # configuration options in the `zenml` section will be overridden with
+    # values specific to ZenML Pro servers computed from the values set in the
+    # `pro` section.
+    enabled: false
+
+    # The URL where the ZenML Pro server API is reachable
+    apiURL: https://cloudapi.zenml.io
+
+    # The URL where the ZenML Pro dashboard is reachable.
+    dashboardURL: https://cloud.zenml.io
+
+    # Additional origins to allow in the CORS policy.
+    extraCorsOrigins:
+
+    # The ID of the ZenML Pro tenant to use.
+    tenantID:
+
+    # The name of the ZenML Pro tenant to use.
+    tenantName:
+
+    # The ID of the ZenML Pro organization to use.
+    organizationID:
+
+    # The name of the ZenML Pro organization to use.
+    organizationName:
+
+    # The enrollment key to use for the ZenML Pro tenant. If not specified,
+    # an enrollment key will be auto-generated.
+    enrollmentKey:
+
+  # The URL where the ZenML server API is reachable. If not specified, the
+  # clients will use the same URL used to connect them to the ZenML server.
+  #
+  # IMPORTANT: this value must be set for ZenML Pro servers.
+  serverURL:
+
+  # The URL where the ZenML dashboard is reachable.
+  # If not specified, the `serverURL` value is used. This should be
+  # configured if the dashboard is served from a different URL than the
+  # ZenML server.
+  #
+  # This is value is used to compute the dashboard URLs during the web login
+  # authentication workflow, to print dashboard URLs in log messages when
+  # running a pipeline and for other similar tasks.
+  #
+  # This value is overridden if the `zenml.pro.enabled` value is set.
+  dashboardURL:
+
   debug: true
 
   # Flag to enable/disable the tracking process of the analytics
@@ -23,6 +85,8 @@ zenml:
 
   # ZenML server deployment type. This field is used for telemetry purposes.
   # Example values are "local", "kubernetes", "aws", "gcp", "azure".
+  #
+  # This value is overridden if the `zenml.pro.enabled` value is set.
   deploymentType:
 
   # Authentication settings that control how the ZenML server authenticates
@@ -35,6 +99,8 @@ zenml:
     # HTTP_BASIC - HTTP Basic authentication
     # OAUTH2_PASSWORD_BEARER - OAuth2 password bearer
     # EXTERNAL - External authentication (e.g. via a remote authenticator)
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     authType: OAUTH2_PASSWORD_BEARER
 
     # The secret key used to sign JWT tokens. This should be set to
@@ -66,7 +132,6 @@ zenml:
     # ZenML Server ID.
     jwtTokenIssuer:
 
-
     # The audience of the JWT tokens. If not specified, the audience is set to
     # the ZenML Server ID.
     jwtTokenAudience:
@@ -77,6 +142,8 @@ zenml:
 
     # The expiration time of JWT tokens in minutes. If not specified, generated
     # JWT tokens will not be set to expire.
+    #
+    # This value is automatically set if the `zenml.pro.enabled` value is set.
     jwtTokenExpireMinutes: 
 
     # The name of the http-only cookie used to store the JWT tokens used to
@@ -92,34 +159,42 @@ zenml:
     # The origins allowed to make cross-origin requests to the ZenML server. If
     # not specified, all origins are allowed. Set this when the ZenML dashboard
     # is hosted on a different domain than the ZenML server.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     corsAllowOrigins:
       - "*"
 
     # The maximum number of failed authentication attempts allowed for an OAuth
     # 2.0 device before the device is locked.
+    #
+    # This value is ignored if the `zenml.auth.authType` is set to `EXTERNAL` or
+    # `NO_AUTH`.
     maxFailedDeviceAuthAttempts: 3
 
     # The timeout in seconds after which a pending OAuth 2.0 device
     # authorization request expires.
+    #
+    # This value is ignored if the `zenml.auth.authType` is set to `EXTERNAL` or
+    # `NO_AUTH`.
     deviceAuthTimeout: 300
 
     # The polling interval in seconds used by clients to poll the OAuth 2.0
     # device authorization endpoint for the status of a pending device
     # authorization request.
+    #
+    # This value is ignored if the `zenml.auth.authType` is set to `EXTERNAL` or
+    # `NO_AUTH`.
     deviceAuthPollingInterval: 5
 
-    # The URL where the ZenML dashboard is hosted. Used to construct the OAuth
-    # 2.0 device authorization endpoint. If not set, a partial URL is returned
-    # to the client which is used to construct the full URL based on the
-    # server's root URL path.
-    dashboardURL:
-
     # The time in minutes that an OAuth 2.0 device is allowed to be used to
     # authenticate with the ZenML server. If not set or if
     # `zenml.auth.jwtTokenExpireMinutes` is not set, the devices are allowed to
     # be used indefinitely. This controls the expiration time of the JWT tokens
     # issued to clients after they have authenticated with the ZenML server
     # using an OAuth 2.0 device.
+    #
+    # This value is ignored if the `zenml.auth.authType` is set to `EXTERNAL` or
+    # `NO_AUTH`.
     deviceExpirationMinutes:
 
     # The time in minutes that a trusted OAuth 2.0 device is allowed to be used
@@ -128,33 +203,46 @@ zenml:
     # be used indefinitely. This controls the expiration time of the JWT tokens
     # issued to clients after they have authenticated with the ZenML server
     # using an OAuth 2.0 device that was previously trusted by the user.
+    #
+    # This value is ignored if the `zenml.auth.authType` is set to `EXTERNAL` or
+    # `NO_AUTH`.
     trustedDeviceExpirationMinutes:
 
     # The login URL of an external authenticator service to use with the
     # `EXTERNAL` authentication scheme. Only relevant if `zenml.auth.authType`
     # is set to `EXTERNAL`.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     externalLoginURL:
 
     # The user info URL of an external authenticator service to use with the
     # `EXTERNAL` authentication scheme. Only relevant if `zenml.auth.authType`
     # is set to `EXTERNAL`.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     externalUserInfoURL:
 
-    # The name of the http-only cookie used to store the bearer token used to
-    # authenticate with the external authenticator service. Only relevant if
-    # `zenml.auth.authType` is set to `EXTERNAL`.
-    externalCookieName:
-
     # The UUID of the ZenML server to use with the `EXTERNAL` authentication
     # scheme. If not specified, the regular ZenML server ID (deployment ID) is
     # used.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     externalServerID:
 
     # Source pointing to a class implementing the RBAC interface defined by
-    # `zenml.zen_server.rbac_interface.RBACInterface`. If not specified,
+    # `zenml.zen_server.rbac.rbac_interface.RBACInterface`. If not specified,
     # RBAC will not be enabled for this server.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
     rbacImplementationSource:
 
+    # Source pointing to a class implementing the feature gate interface defined
+    # by `zenml.zen_server.feature_gate.feature_gate_interface.FeatureGateInterface`.
+    # If not specified, feature gating will not be enabled for this server.
+    #
+    # This value is overridden if the `zenml.pro.enabled` value is set.
+    featureGateImplementationSource:
+
   # The root URL path to use when behind a proxy. This is useful when the
   # `rewrite-target` annotation is used in the ingress controller, e.g.:
   #
@@ -171,13 +259,8 @@ zenml:
   # ```
   rootUrlPath:
 
-
-
   defaultProject: default
-  defaultUsername: default
-  # Use your own password here
-  defaultPassword: zenml
-
+  
   # Implicit authentication methods featured by service connectors that support
   # them are disabled by default, for security reasons. This is because they
   # allow users to authenticate to the cloud provider where ZenML is running
@@ -188,13 +271,61 @@ zenml:
   # used, which will not be persisted across pod restarts.
   # NOTE: the certificate files need to be copied in the helm chart folder and
   # the paths configured here need to be relative to the root of the helm chart.
-  database: {}
+  database:
     # url: "mysql://admin:password@zenml-mysql:3306/database"
     # sslCa: /path/to/ca.pem
     # sslCert: /path/to/client-cert.pem
     # sslKey: /path/to/client-key.pem
     # sslVerifyServerCert: True
 
+    # Connection pool settings (only relevant for MySQL databases).
+    #
+    # NOTE: these values should be coordinated with the `zenml.threadPoolSize`
+    # to ensure that the ZenML server workers do not block on database
+    # connections (i.e. the sum of the pool size and max overflow should be
+    # greater than or equal to the thread pool size).
+    #
+    # poolSize: 20
+    # maxOverflow: 20
+
+    # ZenML supports backing up the database before DB migrations are performed
+    # and restoring it in case of a DB migration failure. For more information,
+    # see the following documentation:
+    # https://docs.zenml.io/getting-started/deploying-zenml/deploy-with-helm#database-backup-and-recovery
+    #
+    # Several backup strategies are supported:
+    #
+    # disabled - no backup is performed
+    # in-memory - the database schema and data are stored in memory. This is
+    #             the fastest backup strategy, but the backup is not persisted
+    #             across pod restarts, so no manual intervention is possible
+    #             in case the automatic DB recovery fails after a failed DB
+    #             migration. Adequate memory resources should be allocated to
+    #             the ZenML server pod when using this backup strategy with
+    #             large databases.
+    #             This is the default backup strategy.
+    # dump-file - the database schema and data are dumped to a local file.
+    #             Users may optionally configure a persistent volume where
+    #             the dump file will be stored by setting the
+    #             `backupPVStorageSize` and optionally the
+    #             `backupPVStorageClass` options. If a
+    #             persistent volume is not configured, the dump file will be
+    #             stored in an emptyDir volume, which is not persisted.
+    #             NOTE: you should set the podSecurityContext.fsGroup option
+    #             if you are using a persistent volume to store the dump file.
+    # database - the database is copied to a backup database in the same
+    #            database server. This requires the `backupDatabase`
+    #            option to be set to the name of the backup database.
+    #            This backup strategy is only supported for MySQL
+    #            compatible databases and the user specified in the
+    #            database URL must have permissions to manage (create, drop, and
+    #            modify) the backup database in addition to the main
+    #            database.
+    backupStrategy: in-memory
+    # backupPVStorageClass: standard
+    # backupPVStorageSize: 1Gi
+    # backupDatabase: "zenml_backup"
+
 
   # Secrets store settings. This is used to store centralized secrets.
   secretsStore:
@@ -260,7 +391,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/aws-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -299,7 +430,8 @@ zenml:
         aws_access_key_id:
         aws_secret_access_key:
         # The AWS role ARN to use to authenticate with the AWS Secrets Manager
-        # instance. Required if the `authMethod` is set to `iam-role`.
+        # instance. Required if the `authMethod` is set to `iam-role`. Optional
+        # if the `authMethod` is set to `implicit`.
         role_arn:
         
       # The AWS region to use. This must be set to the region where the AWS
@@ -332,7 +464,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/gcp-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -404,7 +536,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/azure-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -556,7 +688,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/aws-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/aws-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -610,7 +742,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/gcp-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/gcp-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -630,6 +762,8 @@ zenml:
       # - user-account - Use a GCP user account
       # - service-account - Use a GCP service account
       # - impersonation - Use the GCP service account impersonation feature
+      # - external-account - Use workload identity federation to authenticate
+      #                      with the GCP Secrets Manager API
       authMethod: service-account
 
       # The GCP Service Connector authentication configuration. This should
@@ -656,6 +790,11 @@ zenml:
         # `impersonation`.
         target_principal:
 
+        # The GCP workload identity credentials to use to authenticate with the
+        # GCP Secrets Manager instance. Required if the `authMethod` is set to
+        # `external-account`.
+        external_account_json:
+
     # Azure backup Key Vault secrets store configuration. Only relevant if the
     # `azure` backup secrets store type is configured.
     #
@@ -668,7 +807,7 @@ zenml:
     # 
     # For a list of supported authentication methods and their configuration
     # options, see the following documentation:
-    # https://docs.zenml.io/stacks-and-components/auth-management/azure-service-connector#authentication-methods
+    # https://docs.zenml.io/how-to/auth-management/azure-service-connector#authentication-methods
     # 
     # You can also use the ZenML CLI to get the list of supported authentication
     # methods and their configuration options, e.g.:
@@ -737,6 +876,49 @@ zenml:
       # variable in the `zenml.secretEnvironment` variable.
       class_path: my.custom.secrets.store.MyCustomSecretsStore
 
+  # The ZenML server's secure headers configuration. This can be used to
+  # enable, disable or set custom values for security headers in the ZenML
+  # server's HTTP responses. The following values can be set for any of the
+  # supported secure headers configuration options:
+  #
+  # - `enabled`, `on`, `true` or `yes` - enables the secure header with the
+  # default value.
+  # - `disabled`, `off`, `false`, `none` or `no` - disables the secure header
+  # entirely, so that it is not set in the ZenML server's HTTP responses.
+  # - any other value - sets the secure header to the specified value.  
+  secure_headers:
+    # The `Server` HTTP header value used to identify the server. The default
+    # value is the ZenML server ID.
+    server: enabled
+    # The `Strict-Transport-Security` HTTP header value. The default value is
+    # `max-age=63072000; includeSubDomains`.
+    hsts: enabled
+    # The `X-Frame-Options` HTTP header value. The default value is `SAMEORIGIN`.
+    xfo: enabled
+    # The `X-XSS-Protection` HTTP header value. The default value is `0`.
+    # NOTE: this header is deprecated and should not be customized anymore. The
+    # `Content-Security-Policy` header should be used instead.
+    xxp: enabled
+    # The `X-Content-Type-Options` HTTP header value. The default value is
+    # `nosniff`.
+    content: enabled
+    # The `Content-Security-Policy` HTTP header value. This is by default set
+    # to a strict CSP policy that only allows content from the origins required
+    # by the ZenML dashboard.
+    # NOTE: customizing this header is discouraged, as it may cause the ZenML
+    # dashboard to malfunction.
+    csp: enabled
+    # The `Referrer-Policy` HTTP header value. The default value is
+    # `no-referrer-when-downgrade`.
+    referrer: enabled
+    # The `Cache-Control` HTTP header value. The default value is
+    # `no-store, no-cache, must-revalidate`.
+    cache: enabled
+    # The `Permissions-Policy` HTTP header value. The default value is
+    # `accelerometer=(), camera=(), geolocation=(), gyroscope=(), magnetometer=(), microphone=(), payment=(), usb=()`.
+    permissions: enabled
+
+
   # Extra environment variables to set in the ZenML server container.
   environment: {}
 
@@ -787,7 +969,7 @@ serviceAccount:
 podAnnotations: {}
 
 podSecurityContext: {}
-  # fsGroup: 2000
+  # fsGroup: 1000 # if you're using a PVC for backup, this should necessarily be set.
 
 securityContext:
   runAsNonRoot: true

zenml/zen_server/exceptions.py

@@ -20,13 +20,17 @@ from pydantic import BaseModel
 
 from zenml.exceptions import (
     AuthorizationException,
+    CredentialsNotValid,
     DoesNotExistException,
     DuplicateRunNameError,
+    EntityCreationError,
     EntityExistsError,
     IllegalOperationError,
+    MethodNotAllowedError,
     SecretExistsError,
     StackComponentExistsError,
     StackExistsError,
+    SubscriptionUpgradeRequiredError,
     ValidationError,
     ZenKeyError,
 )
@@ -38,7 +42,7 @@ if TYPE_CHECKING:
 class ErrorModel(BaseModel):
     """Base class for error responses."""
 
-    detail: Any
+    detail: Optional[Any] = None
 
 
 error_response = dict(model=ErrorModel)
@@ -76,7 +80,10 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     # 403 Forbidden
     (IllegalOperationError, 403),
     # 401 Unauthorized
+    (CredentialsNotValid, 401),
     (AuthorizationException, 401),
+    # 402 Payment required
+    (SubscriptionUpgradeRequiredError, 402),
     # 404 Not Found
     (DoesNotExistException, 404),
     (ZenKeyError, 404),
@@ -87,9 +94,12 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     # 422 Unprocessable Entity
     (ValueError, 422),
     # 500 Internal Server Error
+    (EntityCreationError, 500),
     (RuntimeError, 500),
     # 501 Not Implemented,
     (NotImplementedError, 501),
+    # 405 Method Not Allowed
+    (MethodNotAllowedError, 405),
 ]
 
 
@@ -247,4 +257,13 @@ def exception_from_response(
 
         exception = default_exc
 
+    # There is one special case where we want to return a specific exception:
+    # 401 Unauthorized exceptions thrown directly by FastAPI in the course of
+    # authentication are interpreted as AuthorizationException, but we want to
+    # return CredentialsNotValid instead.
+    if response.status_code == 401:
+        if not isinstance(exception(), CredentialsNotValid):
+            if response.headers.get("WWW-Authenticate"):
+                return CredentialsNotValid(exc_msg)
+
     return exception(exc_msg)

zenml/zen_server/feature_gate/__init__.py

@@ -0,0 +1,13 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.

zenml/zen_server/feature_gate/endpoint_utils.py

@@ -0,0 +1,61 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""All endpoint utils for the feature gate implementations."""
+
+from uuid import UUID
+
+from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.utils import feature_gate, server_config
+
+
+def check_entitlement(resource_type: ResourceType) -> None:
+    """Queries the feature gate to see if the operation falls within the tenants entitlements.
+
+    Raises an exception if the user is not entitled to create an instance of the
+    resource. Otherwise, simply returns.
+
+    Args:
+        resource_type: The type of resource to check for.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    return feature_gate().check_entitlement(resource=resource_type)
+
+
+def report_usage(resource_type: ResourceType, resource_id: UUID) -> None:
+    """Reports the creation/usage of a feature/resource.
+
+    Args:
+        resource_type: The type of resource to report a usage for
+        resource_id: ID of the resource that was created.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    feature_gate().report_event(
+        resource=resource_type, resource_id=resource_id
+    )
+
+
+def report_decrement(resource_type: ResourceType, resource_id: UUID) -> None:
+    """Reports the deletion/deactivation of a feature/resource.
+
+    Args:
+        resource_type: The type of resource to report a decrement in count for.
+        resource_id: ID of the resource that was deleted.
+    """
+    if not server_config().feature_gate_enabled:
+        return
+    feature_gate().report_event(
+        resource=resource_type, resource_id=resource_id, is_decrement=True
+    )

zenml/zen_server/feature_gate/feature_gate_interface.py

@@ -0,0 +1,49 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Definition of the feature gate interface."""
+
+from abc import ABC, abstractmethod
+from uuid import UUID
+
+from zenml.zen_server.rbac.models import ResourceType
+
+
+class FeatureGateInterface(ABC):
+    """Feature gate interface definition."""
+
+    @abstractmethod
+    def check_entitlement(self, resource: ResourceType) -> None:
+        """Checks if a user is entitled to create a resource.
+
+        Args:
+            resource: The resource the user wants to create
+
+        Raises:
+            UpgradeRequiredError in case a subscription limit is reached
+        """
+
+    @abstractmethod
+    def report_event(
+        self,
+        resource: ResourceType,
+        resource_id: UUID,
+        is_decrement: bool = False,
+    ) -> None:
+        """Reports the usage of a feature to the aggregator backend.
+
+        Args:
+            resource: The resource the user created
+            resource_id: ID of the resource that was created/deleted.
+            is_decrement: In case this event reports an actual decrement of usage
+        """