feat(jupyterhub): make SecretStore singleton

This commit is contained in:
Masaki Yatsu
2025-08-31 22:33:22 +09:00
parent 106714d0ac
commit 972805aa65
2 changed files with 37 additions and 18 deletions

View File

@@ -3,9 +3,9 @@ set fallback := true
export JUPYTERHUB_NAMESPACE := env("JUPYTERHUB_NAMESPACE", "jupyter") export JUPYTERHUB_NAMESPACE := env("JUPYTERHUB_NAMESPACE", "jupyter")
export JUPYTERHUB_CHART_VERSION := env("JUPYTERHUB_CHART_VERSION", "4.2.0") export JUPYTERHUB_CHART_VERSION := env("JUPYTERHUB_CHART_VERSION", "4.2.0")
export JUPYTERHUB_OIDC_CLIENT_ID := env("JUPYTERHUB_OIDC_CLIENT_ID", "jupyterhub") export JUPYTERHUB_OIDC_CLIENT_ID := env("JUPYTERHUB_OIDC_CLIENT_ID", "jupyterhub")
export JUPYTERHUB_ENABLE_NFS_PV := env("JUPYTERHUB_ENABLE_NFS_PV", "") export JUPYTERHUB_NFS_PV_ENABLED := env("JUPYTERHUB_NFS_PV_ENABLED", "")
export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "false") export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "")
export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-5") export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-6")
export KERNEL_IMAGE_BUUN_STACK_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_REPOSITORY", "buun-stack-notebook") export KERNEL_IMAGE_BUUN_STACK_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_REPOSITORY", "buun-stack-notebook")
export KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY", "buun-stack-cuda-notebook") export KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY", "buun-stack-cuda-notebook")
export JUPYTER_PROFILE_MINIMAL_ENABLED := env("JUPYTER_PROFILE_MINIMAL_ENABLED", "false") export JUPYTER_PROFILE_MINIMAL_ENABLED := env("JUPYTER_PROFILE_MINIMAL_ENABLED", "false")
@@ -64,14 +64,14 @@ install:
export JUPYTER_PYTHON_KERNEL_TAG=${JUPYTER_PYTHON_KERNEL_TAG} export JUPYTER_PYTHON_KERNEL_TAG=${JUPYTER_PYTHON_KERNEL_TAG}
export JUPYTER_FSGID=${JUPYTER_FSGID:-100} export JUPYTER_FSGID=${JUPYTER_FSGID:-100}
export PVC_NAME="" export PVC_NAME=""
if [ -z "${JUPYTERHUB_ENABLE_NFS_PV}" ]; then if [ -z "${JUPYTERHUB_NFS_PV_ENABLED}" ]; then
if gum confirm "Are you going to use NFS PV?"; then if gum confirm "Are you going to use NFS PV?"; then
JUPYTERHUB_ENABLE_NFS_PV=true JUPYTERHUB_NFS_PV_ENABLED=true
else else
JUPYTERHUB_ENABLE_NFS_PV=false JUPYTERHUB_NFS_PV_ENABLED=false
fi fi
fi fi
if [ "${JUPYTERHUB_ENABLE_NFS_PV}" = "true" ]; then if [ "${JUPYTERHUB_NFS_PV_ENABLED}" = "true" ]; then
if ! helm status longhorn -n ${LONGHORN_NAMESPACE} &>/dev/null; then if ! helm status longhorn -n ${LONGHORN_NAMESPACE} &>/dev/null; then
echo "Longhorn is not installed. Please install Longhorn first." >&2 echo "Longhorn is not installed. Please install Longhorn first." >&2
exit 1 exit 1
@@ -104,7 +104,13 @@ install:
# wait deployments manually because `helm upgrade --wait` does not work for JupyterHub # wait deployments manually because `helm upgrade --wait` does not work for JupyterHub
just k8s::wait-deployments-ready ${JUPYTERHUB_NAMESPACE} hub proxy just k8s::wait-deployments-ready ${JUPYTERHUB_NAMESPACE} hub proxy
# Setup Vault integration if enabled if [ -z "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" ]; then
if gum confirm "Are you going to enable Vault integration?"; then
JUPYTERHUB_VAULT_INTEGRATION_ENABLED=true
else
JUPYTERHUB_VAULT_INTEGRATION_ENABLED=false
fi
fi
if [ "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" = "true" ]; then if [ "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" = "true" ]; then
just setup-vault-jwt-auth just setup-vault-jwt-auth
fi fi

View File

@@ -30,6 +30,9 @@ class SecretStore:
automatic OIDC token refresh via Keycloak integration and provides both automatic OIDC token refresh via Keycloak integration and provides both
manual and background token management options. manual and background token management options.
This class implements the singleton pattern to ensure only one instance
exists per user session, preventing duplicate background refresh threads.
Attributes Attributes
---------- ----------
auto_token_refresh : bool auto_token_refresh : bool
@@ -58,6 +61,14 @@ class SecretStore:
'sk-123' 'sk-123'
""" """
_instance = None
_initialized = False
def __new__(cls, *args, **kwargs):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__( def __init__(
self, self,
auto_token_refresh: bool = True, auto_token_refresh: bool = True,
@@ -67,18 +78,22 @@ class SecretStore:
""" """
Initialize SecretStore with authentication and configuration. Initialize SecretStore with authentication and configuration.
Note: Due to singleton pattern, parameters are only used on the first
instantiation. Subsequent calls return the existing instance with
its original configuration.
Parameters Parameters
---------- ----------
auto_token_refresh : bool, optional auto_token_refresh : bool, optional
Enable automatic token refresh using Keycloak OIDC, by default True. Enable automatic token refresh using Keycloak OIDC, by default True.
Requires KEYCLOAK_HOST, KEYCLOAK_REALM, and JUPYTERHUB_OIDC_REFRESH_TOKEN Requires KEYCLOAK_HOST, KEYCLOAK_REALM, and JUPYTERHUB_OIDC_REFRESH_TOKEN
environment variables. environment variables. Only used on first instantiation.
refresh_buffer_seconds : int, optional refresh_buffer_seconds : int, optional
Seconds before token expiry to trigger refresh, by default 300. Seconds before token expiry to trigger refresh, by default 300.
Only used when auto_token_refresh is True. Only used when auto_token_refresh is True. Only used on first instantiation.
background_refresh_interval : int, optional background_refresh_interval : int, optional
Seconds between background refresh checks, by default 1800. Seconds between background refresh checks, by default 1800.
Only used when background refresh is started. Only used when background refresh is started. Only used on first instantiation.
Raises Raises
------ ------
@@ -105,37 +120,33 @@ class SecretStore:
... background_refresh_interval=3600 ... background_refresh_interval=3600
... ) ... )
""" """
if self._initialized:
return
self.auto_token_refresh = auto_token_refresh self.auto_token_refresh = auto_token_refresh
self.refresh_buffer_seconds = refresh_buffer_seconds self.refresh_buffer_seconds = refresh_buffer_seconds
self.background_refresh_interval = background_refresh_interval self.background_refresh_interval = background_refresh_interval
# User and environment info
self.username = os.getenv("JUPYTERHUB_USER") self.username = os.getenv("JUPYTERHUB_USER")
self.vault_addr = os.getenv("VAULT_ADDR") self.vault_addr = os.getenv("VAULT_ADDR")
# Keycloak configuration (only needed if auto_token_refresh is enabled)
if self.auto_token_refresh: if self.auto_token_refresh:
self.keycloak_host = os.getenv("KEYCLOAK_HOST") self.keycloak_host = os.getenv("KEYCLOAK_HOST")
self.keycloak_realm = os.getenv("KEYCLOAK_REALM") self.keycloak_realm = os.getenv("KEYCLOAK_REALM")
self.keycloak_client_id = os.getenv("KEYCLOAK_CLIENT_ID", "jupyterhub") self.keycloak_client_id = os.getenv("KEYCLOAK_CLIENT_ID", "jupyterhub")
self.refresh_token = os.getenv("JUPYTERHUB_OIDC_REFRESH_TOKEN") self.refresh_token = os.getenv("JUPYTERHUB_OIDC_REFRESH_TOKEN")
# Token management
self.access_token = os.getenv("JUPYTERHUB_OIDC_ACCESS_TOKEN") self.access_token = os.getenv("JUPYTERHUB_OIDC_ACCESS_TOKEN")
self.token_expiry = ( self.token_expiry = (
self._get_token_expiry(self.access_token) if self.access_token else None self._get_token_expiry(self.access_token) if self.access_token else None
) )
# Initialize Vault client
self.client = hvac.Client(url=self.vault_addr, verify=False) self.client = hvac.Client(url=self.vault_addr, verify=False)
# Background refresher
self._background_refresher = None self._background_refresher = None
# Authenticate initially
self._authenticate_vault() self._authenticate_vault()
# Set base path for user storage
self.base_path = f"jupyter/users/{self.username}" self.base_path = f"jupyter/users/{self.username}"
logger.info(f"SecretStore initialized for user: {self.username}") logger.info(f"SecretStore initialized for user: {self.username}")
@@ -146,6 +157,8 @@ class SecretStore:
if self.auto_token_refresh and self.token_expiry: if self.auto_token_refresh and self.token_expiry:
logger.info(f"Token expires at: {self.token_expiry}") logger.info(f"Token expires at: {self.token_expiry}")
self._initialized = True
def _get_token_expiry(self, token: str) -> datetime | None: def _get_token_expiry(self, token: str) -> datetime | None:
"""Extract expiry time from JWT token""" """Extract expiry time from JWT token"""
if not token: if not token: