skypilot-nightly 1.0.0.dev20250609__py3-none-any.whl → 1.0.0.dev20250611__py3-none-any.whl

sky/global_user_state.py

@@ -6,11 +6,13 @@ Concepts:
 - Cluster handle: (non-user facing) an opaque backend handle for us to
   interact with a cluster.
 """
+import functools
 import json
 import os
 import pathlib
 import pickle
 import re
+import threading
 import time
 import typing
 from typing import Any, Dict, List, Optional, Set, Tuple
@@ -44,6 +46,9 @@ logger = sky_logging.init_logger(__name__)
 
 _ENABLED_CLOUDS_KEY_PREFIX = 'enabled_clouds_'
 
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_DB_INIT_LOCK = threading.Lock()
+
 Base = declarative.declarative_base()
 
 config_table = sqlalchemy.Table(
@@ -171,11 +176,11 @@ def create_table():
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (SQLALCHEMY_ENGINE.dialect.name
+    if (_SQLALCHEMY_ENGINE.dialect.name
             == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(SQLALCHEMY_ENGINE) as session:
+            with orm.Session(_SQLALCHEMY_ENGINE) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -185,12 +190,12 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    Base.metadata.create_all(bind=SQLALCHEMY_ENGINE)
+    Base.metadata.create_all(bind=_SQLALCHEMY_ENGINE)
 
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         # Add autostop column to clusters table
         db_utils.add_column_to_table_sqlalchemy(session,
                                                 'clusters',
@@ -298,30 +303,53 @@ def create_table():
         session.commit()
 
 
-conn_string = None
-if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
-    conn_string = skypilot_config.get_nested(('db',), None)
-if conn_string:
-    logger.debug(f'using db URI from {conn_string}')
-    SQLALCHEMY_ENGINE = sqlalchemy.create_engine(conn_string)
-else:
-    _DB_PATH = os.path.expanduser('~/.sky/state.db')
-    pathlib.Path(_DB_PATH).parents[0].mkdir(parents=True, exist_ok=True)
-    SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' + _DB_PATH)
-create_table()
-
-
+def initialize_and_get_db() -> sqlalchemy.engine.Engine:
+    global _SQLALCHEMY_ENGINE
+    if _SQLALCHEMY_ENGINE is not None:
+        return _SQLALCHEMY_ENGINE
+    with _DB_INIT_LOCK:
+        if _SQLALCHEMY_ENGINE is None:
+            conn_string = None
+            if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+                conn_string = skypilot_config.get_nested(('db',), None)
+            if conn_string:
+                logger.debug(f'using db URI from {conn_string}')
+                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(conn_string)
+            else:
+                db_path = os.path.expanduser('~/.sky/state.db')
+                pathlib.Path(db_path).parents[0].mkdir(parents=True,
+                                                       exist_ok=True)
+                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
+                                                              db_path)
+            create_table()
+    return _SQLALCHEMY_ENGINE
+
+
+def _init_db(func):
+    """Initialize the database."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        initialize_and_get_db()
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@_init_db
 def add_or_update_user(user: models.User) -> bool:
     """Store the mapping from user hash to user name for display purposes.
 
     Returns:
         Boolean: whether the user is newly added
     """
+    assert _SQLALCHEMY_ENGINE is not None
+
     if user.name is None:
         return False
 
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             # For SQLite, use INSERT OR IGNORE followed by UPDATE to detect new
             # vs existing
@@ -343,7 +371,7 @@ def add_or_update_user(user: models.User) -> bool:
             session.commit()
             return was_inserted
 
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             # For PostgreSQL, use INSERT ... ON CONFLICT with RETURNING to
             # detect insert vs update
@@ -371,20 +399,25 @@ def add_or_update_user(user: models.User) -> bool:
             raise ValueError('Unsupported database dialect')
 
 
+@_init_db
 def get_user(user_id: str) -> Optional[models.User]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(user_table).filter_by(id=user_id).first()
     if row is None:
         return None
     return models.User(id=row.id, name=row.name)
 
 
+@_init_db
 def get_all_users() -> List[models.User]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(user_table).all()
     return [models.User(id=row.id, name=row.name) for row in rows]
 
 
+@_init_db
 def add_or_update_cluster(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle',
                           requested_resources: Optional[Set[Any]],
@@ -405,6 +438,7 @@ def add_or_update_cluster(cluster_name: str,
         config_hash: Configuration hash for the cluster.
         task_config: The config of the task being launched.
     """
+    assert _SQLALCHEMY_ENGINE is not None
     # FIXME: launched_at will be changed when `sky launch -c` is called.
     handle = pickle.dumps(cluster_handle)
     cluster_launched_at = int(time.time()) if is_launch else None
@@ -457,7 +491,7 @@ def add_or_update_cluster(cluster_name: str,
             'config_hash': config_hash,
         })
 
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         # with_for_update() locks the row until commit() or rollback()
         # is called, or until the code escapes the with block.
         cluster_row = session.query(cluster_table).filter_by(
@@ -484,10 +518,10 @@ def add_or_update_cluster(cluster_name: str,
                 'last_creation_command': last_use,
             })
 
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -562,29 +596,34 @@ def _get_user_hash_or_current_user(user_hash: Optional[str]) -> str:
     return common_utils.get_user_hash()
 
 
+@_init_db
 def update_cluster_handle(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle'):
+    assert _SQLALCHEMY_ENGINE is not None
     handle = pickle.dumps(cluster_handle)
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         session.query(cluster_table).filter_by(name=cluster_name).update(
             {cluster_table.c.handle: handle})
         session.commit()
 
 
+@_init_db
 def update_last_use(cluster_name: str):
     """Updates the last used command for the cluster."""
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         session.query(cluster_table).filter_by(name=cluster_name).update(
             {cluster_table.c.last_use: common_utils.get_current_command()})
         session.commit()
 
 
+@_init_db
 def remove_cluster(cluster_name: str, terminate: bool) -> None:
     """Removes cluster_name mapping."""
+    assert _SQLALCHEMY_ENGINE is not None
     cluster_hash = _get_hash_for_existing_cluster(cluster_name)
     usage_intervals = _get_cluster_usage_intervals(cluster_hash)
 
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         # usage_intervals is not None and not empty
         if usage_intervals:
             assert cluster_hash is not None, cluster_name
@@ -614,24 +653,28 @@ def remove_cluster(cluster_name: str, terminate: bool) -> None:
         session.commit()
 
 
+@_init_db
 def get_handle_from_cluster_name(
         cluster_name: str) -> Optional['backends.ResourceHandle']:
+    assert _SQLALCHEMY_ENGINE is not None
     assert cluster_name is not None, 'cluster_name cannot be None'
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None:
         return None
     return pickle.loads(row.handle)
 
 
+@_init_db
 def get_glob_cluster_names(cluster_name: str) -> List[str]:
+    assert _SQLALCHEMY_ENGINE is not None
     assert cluster_name is not None, 'cluster_name cannot be None'
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             rows = session.query(cluster_table).filter(
                 cluster_table.c.name.op('GLOB')(cluster_name)).all()
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             rows = session.query(cluster_table).filter(
                 cluster_table.c.name.op('SIMILAR TO')(
@@ -641,10 +684,12 @@ def get_glob_cluster_names(cluster_name: str) -> List[str]:
     return [row.name for row in rows]
 
 
+@_init_db
 def set_cluster_status(cluster_name: str,
                        status: status_lib.ClusterStatus) -> None:
+    assert _SQLALCHEMY_ENGINE is not None
     current_time = int(time.time())
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.status: status.value,
@@ -656,9 +701,11 @@ def set_cluster_status(cluster_name: str,
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
+@_init_db
 def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
                                to_down: bool) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.autostop: idle_minutes,
@@ -670,24 +717,30 @@ def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
+@_init_db
 def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.launched_at is None:
         return None
     return int(row.launched_at)
 
 
+@_init_db
 def get_cluster_info(cluster_name: str) -> Optional[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.metadata is None:
         return None
     return json.loads(row.metadata)
 
 
+@_init_db
 def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update(
                 {cluster_table.c.metadata: json.dumps(metadata)})
@@ -697,18 +750,22 @@ def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
+@_init_db
 def get_cluster_storage_mounts_metadata(
         cluster_name: str) -> Optional[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.storage_mounts_metadata is None:
         return None
     return pickle.loads(row.storage_mounts_metadata)
 
 
+@_init_db
 def set_cluster_storage_mounts_metadata(
         cluster_name: str, storage_mounts_metadata: Dict[str, Any]) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.storage_mounts_metadata:
@@ -720,12 +777,14 @@ def set_cluster_storage_mounts_metadata(
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
+@_init_db
 def _get_cluster_usage_intervals(
         cluster_hash: Optional[str]
 ) -> Optional[List[Tuple[int, Optional[int]]]]:
+    assert _SQLALCHEMY_ENGINE is not None
     if cluster_hash is None:
         return None
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).first()
     if row is None or row.usage_intervals is None:
@@ -759,10 +818,12 @@ def _get_cluster_duration(cluster_hash: str) -> int:
     return total_duration
 
 
+@_init_db
 def _set_cluster_usage_intervals(
         cluster_hash: str, usage_intervals: List[Tuple[int,
                                                        Optional[int]]]) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).update({
                 cluster_history_table.c.usage_intervals:
@@ -774,12 +835,14 @@ def _set_cluster_usage_intervals(
         raise ValueError(f'Cluster hash {cluster_hash} not found.')
 
 
+@_init_db
 def set_owner_identity_for_cluster(cluster_name: str,
                                    owner_identity: Optional[List[str]]) -> None:
+    assert _SQLALCHEMY_ENGINE is not None
     if owner_identity is None:
         return
     owner_identity_str = json.dumps(owner_identity)
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update(
                 {cluster_table.c.owner: owner_identity_str})
@@ -789,17 +852,21 @@ def set_owner_identity_for_cluster(cluster_name: str,
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
+@_init_db
 def _get_hash_for_existing_cluster(cluster_name: str) -> Optional[str]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.cluster_hash is None:
         return None
     return row.cluster_hash
 
 
+@_init_db
 def get_launched_resources_from_cluster_hash(
         cluster_hash: str) -> Optional[Tuple[int, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).first()
     if row is None:
@@ -841,10 +908,12 @@ def _load_storage_mounts_metadata(
     return pickle.loads(record_storage_mounts_metadata)
 
 
+@_init_db
 @context_utils.cancellation_guard
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None:
         return None
@@ -878,8 +947,10 @@ def get_cluster_from_name(
     return record
 
 
+@_init_db
 def get_clusters() -> List[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(cluster_table).order_by(
             sqlalchemy.desc(cluster_table.c.launched_at)).all()
     records = []
@@ -915,8 +986,10 @@ def get_clusters() -> List[Dict[str, Any]]:
     return records
 
 
+@_init_db
 def get_clusters_from_history() -> List[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(
             cluster_history_table.join(cluster_table,
                                        cluster_history_table.c.cluster_hash ==
@@ -951,16 +1024,20 @@ def get_clusters_from_history() -> List[Dict[str, Any]]:
     return records
 
 
+@_init_db
 def get_cluster_names_start_with(starts_with: str) -> List[str]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(cluster_table).filter(
             cluster_table.c.name.like(f'{starts_with}%')).all()
     return [row.name for row in rows]
 
 
+@_init_db
 def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
                               workspace: str) -> List['clouds.Cloud']:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(config_table).filter_by(
             key=_get_enabled_clouds_key(cloud_capability, workspace)).first()
     ret = []
@@ -981,14 +1058,16 @@ def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
     return enabled_clouds
 
 
+@_init_db
 def set_enabled_clouds(enabled_clouds: List[str],
                        cloud_capability: 'cloud.CloudCapability',
                        workspace: str) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1008,9 +1087,11 @@ def _get_enabled_clouds_key(cloud_capability: 'cloud.CloudCapability',
     return _ENABLED_CLOUDS_KEY_PREFIX + workspace + '_' + cloud_capability.value
 
 
+@_init_db
 def add_or_update_storage(storage_name: str,
                           storage_handle: 'Storage.StorageMetadata',
                           storage_status: status_lib.StorageStatus):
+    assert _SQLALCHEMY_ENGINE is not None
     storage_launched_at = int(time.time())
     handle = pickle.dumps(storage_handle)
     last_use = common_utils.get_current_command()
@@ -1021,11 +1102,11 @@ def add_or_update_storage(storage_name: str,
     if not status_check(storage_status):
         raise ValueError(f'Error in updating global state. Storage Status '
                          f'{storage_status} is passed in incorrectly')
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1048,16 +1129,20 @@ def add_or_update_storage(storage_name: str,
         session.commit()
 
 
+@_init_db
 def remove_storage(storage_name: str):
     """Removes Storage from Database"""
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         session.query(storage_table).filter_by(name=storage_name).delete()
         session.commit()
 
 
+@_init_db
 def set_storage_status(storage_name: str,
                        status: status_lib.StorageStatus) -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(storage_table).filter_by(
             name=storage_name).update({storage_table.c.status: status.value})
         session.commit()
@@ -1066,18 +1151,22 @@ def set_storage_status(storage_name: str,
         raise ValueError(f'Storage {storage_name} not found.')
 
 
+@_init_db
 def get_storage_status(storage_name: str) -> Optional[status_lib.StorageStatus]:
+    assert _SQLALCHEMY_ENGINE is not None
     assert storage_name is not None, 'storage_name cannot be None'
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(storage_table).filter_by(name=storage_name).first()
     if row:
         return status_lib.StorageStatus[row.status]
     return None
 
 
+@_init_db
 def set_storage_handle(storage_name: str,
                        handle: 'Storage.StorageMetadata') -> None:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         count = session.query(storage_table).filter_by(
             name=storage_name).update(
                 {storage_table.c.handle: pickle.dumps(handle)})
@@ -1087,25 +1176,29 @@ def set_storage_handle(storage_name: str,
         raise ValueError(f'Storage{storage_name} not found.')
 
 
+@_init_db
 def get_handle_from_storage_name(
         storage_name: Optional[str]) -> Optional['Storage.StorageMetadata']:
+    assert _SQLALCHEMY_ENGINE is not None
     if storage_name is None:
         return None
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(storage_table).filter_by(name=storage_name).first()
     if row:
         return pickle.loads(row.handle)
     return None
 
 
+@_init_db
 def get_glob_storage_name(storage_name: str) -> List[str]:
+    assert _SQLALCHEMY_ENGINE is not None
     assert storage_name is not None, 'storage_name cannot be None'
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             rows = session.query(storage_table).filter(
                 storage_table.c.name.op('GLOB')(storage_name)).all()
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             rows = session.query(storage_table).filter(
                 storage_table.c.name.op('SIMILAR TO')(
@@ -1115,15 +1208,19 @@ def get_glob_storage_name(storage_name: str) -> List[str]:
     return [row.name for row in rows]
 
 
+@_init_db
 def get_storage_names_start_with(starts_with: str) -> List[str]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(storage_table).filter(
             storage_table.c.name.like(f'{starts_with}%')).all()
     return [row.name for row in rows]
 
 
+@_init_db
 def get_storage() -> List[Dict[str, Any]]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(storage_table).all()
     records = []
     for row in rows:
@@ -1138,8 +1235,10 @@ def get_storage() -> List[Dict[str, Any]]:
     return records
 
 
+@_init_db
 def get_ssh_keys(user_hash: str) -> Tuple[str, str, bool]:
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(ssh_key_table).filter_by(
             user_hash=user_hash).first()
     if row:
@@ -1147,12 +1246,14 @@ def get_ssh_keys(user_hash: str) -> Tuple[str, str, bool]:
     return '', '', False
 
 
+@_init_db
 def set_ssh_keys(user_hash: str, ssh_public_key: str, ssh_private_key: str):
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1171,6 +1272,7 @@ def set_ssh_keys(user_hash: str, ssh_public_key: str, ssh_private_key: str):
         session.commit()
 
 
+@_init_db
 def get_cluster_yaml_str(cluster_yaml_path: Optional[str]) -> Optional[str]:
     """Get the cluster yaml from the database or the local file system.
     If the cluster yaml is not in the database, check if it exists on the
@@ -1178,11 +1280,12 @@ def get_cluster_yaml_str(cluster_yaml_path: Optional[str]) -> Optional[str]:
 
     It is assumed that the cluster yaml file is named as <cluster_name>.yml.
     """
+    assert _SQLALCHEMY_ENGINE is not None
     if cluster_yaml_path is None:
         raise ValueError('Attempted to read a None YAML.')
     cluster_file_name = os.path.basename(cluster_yaml_path)
     cluster_name, _ = os.path.splitext(cluster_file_name)
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_yaml_table).filter_by(
             cluster_name=cluster_name).first()
     if row is None:
@@ -1210,13 +1313,15 @@ def get_cluster_yaml_dict(cluster_yaml_path: Optional[str]) -> Dict[str, Any]:
     return yaml.safe_load(yaml_str)
 
 
+@_init_db
 def set_cluster_yaml(cluster_name: str, yaml_str: str) -> None:
     """Set the cluster yaml in the database."""
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
-        if (SQLALCHEMY_ENGINE.dialect.name ==
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (SQLALCHEMY_ENGINE.dialect.name ==
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1230,8 +1335,10 @@ def set_cluster_yaml(cluster_name: str, yaml_str: str) -> None:
         session.commit()
 
 
+@_init_db
 def remove_cluster_yaml(cluster_name: str):
-    with orm.Session(SQLALCHEMY_ENGINE) as session:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
         session.query(cluster_yaml_table).filter_by(
             cluster_name=cluster_name).delete()
         session.commit()