skypilot-nightly 1.0.0.dev20250523__py3-none-any.whl → 1.0.0.dev20250526__py3-none-any.whl

sky/global_user_state.py

@@ -10,14 +10,20 @@ import json
 import os
 import pathlib
 import pickle
-import sqlite3
 import time
 import typing
 from typing import Any, Dict, List, Optional, Set, Tuple
 import uuid
 
+import sqlalchemy
+from sqlalchemy import exc as sqlalchemy_exc
+from sqlalchemy import orm
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.ext import declarative
+
 from sky import models
 from sky import sky_logging
+from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import context_utils
 from sky.utils import db_utils
@@ -37,159 +43,215 @@ _ENABLED_CLOUDS_KEY_PREFIX = 'enabled_clouds_'
 _DB_PATH = os.path.expanduser('~/.sky/state.db')
 pathlib.Path(_DB_PATH).parents[0].mkdir(parents=True, exist_ok=True)
 
-
-def create_table(cursor, conn):
+_SQLALCHEMY_ENGINE = sqlalchemy.create_engine(f'sqlite:///{_DB_PATH}')
+
+Base = declarative.declarative_base()
+
+config_table = sqlalchemy.Table(
+    'config',
+    Base.metadata,
+    sqlalchemy.Column('key', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('value', sqlalchemy.Text),
+)
+
+user_table = sqlalchemy.Table(
+    'users',
+    Base.metadata,
+    sqlalchemy.Column('id', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('name', sqlalchemy.Text),
+)
+
+cluster_table = sqlalchemy.Table(
+    'clusters',
+    Base.metadata,
+    sqlalchemy.Column('name', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('launched_at', sqlalchemy.Integer),
+    sqlalchemy.Column('handle', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('last_use', sqlalchemy.Text),
+    sqlalchemy.Column('status', sqlalchemy.Text),
+    sqlalchemy.Column('autostop', sqlalchemy.Integer, server_default='-1'),
+    sqlalchemy.Column('to_down', sqlalchemy.Integer, server_default='0'),
+    sqlalchemy.Column('metadata', sqlalchemy.Text, server_default='{}'),
+    sqlalchemy.Column('owner', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('cluster_hash', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('storage_mounts_metadata',
+                      sqlalchemy.LargeBinary,
+                      server_default=None),
+    sqlalchemy.Column('cluster_ever_up', sqlalchemy.Integer,
+                      server_default='0'),
+    sqlalchemy.Column('status_updated_at',
+                      sqlalchemy.Integer,
+                      server_default=None),
+    sqlalchemy.Column('config_hash', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('user_hash', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('workspace',
+                      sqlalchemy.Text,
+                      server_default=constants.SKYPILOT_DEFAULT_WORKSPACE),
+)
+
+storage_table = sqlalchemy.Table(
+    'storage',
+    Base.metadata,
+    sqlalchemy.Column('name', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('launched_at', sqlalchemy.Integer),
+    sqlalchemy.Column('handle', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('last_use', sqlalchemy.Text),
+    sqlalchemy.Column('status', sqlalchemy.Text),
+)
+
+# Table for Cluster History
+# usage_intervals: List[Tuple[int, int]]
+#  Specifies start and end timestamps of cluster.
+#  When the last end time is None, the cluster is still UP.
+#  Example: [(start1, end1), (start2, end2), (start3, None)]
+
+# requested_resources: Set[resource_lib.Resource]
+#  Requested resources fetched from task that user specifies.
+
+# launched_resources: Optional[resources_lib.Resources]
+#  Actual launched resources fetched from handle for cluster.
+
+# num_nodes: Optional[int] number of nodes launched.
+cluster_history_table = sqlalchemy.Table(
+    'cluster_history',
+    Base.metadata,
+    sqlalchemy.Column('cluster_hash', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('name', sqlalchemy.Text),
+    sqlalchemy.Column('num_nodes', sqlalchemy.Integer),
+    sqlalchemy.Column('requested_resources', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('launched_resources', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('usage_intervals', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('user_hash', sqlalchemy.Text),
+)
+
+
+def create_table():
     # Enable WAL mode to avoid locking issues.
     # See: issue #1441 and PR #1509
     # 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 not common_utils.is_wsl():
+    if (_SQLALCHEMY_ENGINE.dialect.name
+            == db_utils.SQLAlchemyDialect.SQLITE.value and
+            not common_utils.is_wsl()):
         try:
-            cursor.execute('PRAGMA journal_mode=WAL')
-        except sqlite3.OperationalError as e:
+            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+                session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
+                session.commit()
+        except sqlalchemy_exc.OperationalError as e:
             if 'database is locked' not in str(e):
                 raise
             # If the database is locked, it is OK to continue, as the WAL mode
             # is not critical and is likely to be enabled by other processes.
 
-    # Table for Clusters
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS clusters (
-        name TEXT PRIMARY KEY,
-        launched_at INTEGER,
-        handle BLOB,
-        last_use TEXT,
-        status TEXT,
-        autostop INTEGER DEFAULT -1,
-        metadata TEXT DEFAULT '{}',
-        to_down INTEGER DEFAULT 0,
-        owner TEXT DEFAULT null,
-        cluster_hash TEXT DEFAULT null,
-        storage_mounts_metadata BLOB DEFAULT null,
-        cluster_ever_up INTEGER DEFAULT 0,
-        status_updated_at INTEGER DEFAULT null,
-        config_hash TEXT DEFAULT null,
-        user_hash TEXT DEFAULT null)""")
-
-    # Table for Cluster History
-    # usage_intervals: List[Tuple[int, int]]
-    #  Specifies start and end timestamps of cluster.
-    #  When the last end time is None, the cluster is still UP.
-    #  Example: [(start1, end1), (start2, end2), (start3, None)]
-
-    # requested_resources: Set[resource_lib.Resource]
-    #  Requested resources fetched from task that user specifies.
-
-    # launched_resources: Optional[resources_lib.Resources]
-    #  Actual launched resources fetched from handle for cluster.
-
-    # num_nodes: Optional[int] number of nodes launched.
-
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS cluster_history (
-        cluster_hash TEXT PRIMARY KEY,
-        name TEXT,
-        num_nodes int,
-        requested_resources BLOB,
-        launched_resources BLOB,
-        usage_intervals BLOB,
-        user_hash TEXT)""")
-    # Table for configs (e.g. enabled clouds)
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS config (
-        key TEXT PRIMARY KEY, value TEXT)""")
-    # Table for Storage
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS storage (
-        name TEXT PRIMARY KEY,
-        launched_at INTEGER,
-        handle BLOB,
-        last_use TEXT,
-        status TEXT)""")
-    # Table for User
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS users (
-        id TEXT PRIMARY KEY,
-        name TEXT)""")
+    # Create tables if they don't exist
+    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.
-    # Add autostop column to clusters table
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'autostop',
-                                 'INTEGER DEFAULT -1')
-
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'metadata',
-                                 'TEXT DEFAULT \'{}\'')
-
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'to_down',
-                                 'INTEGER DEFAULT 0')
-
-    # The cloud identity that created the cluster.
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'owner', 'TEXT')
-
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'cluster_hash',
-                                 'TEXT DEFAULT null')
-
-    db_utils.add_column_to_table(cursor, conn, 'clusters',
-                                 'storage_mounts_metadata', 'BLOB DEFAULT null')
-    db_utils.add_column_to_table(
-        cursor,
-        conn,
-        'clusters',
-        'cluster_ever_up',
-        'INTEGER DEFAULT 0',
-        # Set the value to 1 so that all the existing clusters before #2977
-        # are considered as ever up, i.e:
-        #   existing cluster's default (null) -> 1;
-        #   new cluster's default -> 0;
-        # This is conservative for the existing clusters: even if some INIT
-        # clusters were never really UP, setting it to 1 means they won't be
-        # auto-deleted during any failover.
-        value_to_replace_existing_entries=1)
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'status_updated_at',
-                                 'INTEGER DEFAULT null')
-    db_utils.add_column_to_table(
-        cursor,
-        conn,
-        'clusters',
-        'user_hash',
-        'TEXT DEFAULT null',
-        value_to_replace_existing_entries=common_utils.get_user_hash())
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'config_hash',
-                                 'TEXT DEFAULT null')
-
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'config_hash',
-                                 'TEXT DEFAULT null')
-
-    db_utils.add_column_to_table(cursor, conn, 'cluster_history', 'user_hash',
-                                 'TEXT DEFAULT null')
-    conn.commit()
-
-
-_DB = db_utils.SQLiteConn(_DB_PATH, create_table)
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        # Add autostop column to clusters table
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters', 'autostop',
+                                                'INTEGER DEFAULT -1')
+
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters', 'metadata',
+                                                'TEXT DEFAULT \'{}\'')
+
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters', 'to_down',
+                                                'INTEGER DEFAULT 0')
+
+        # The cloud identity that created the cluster.
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters', 'owner',
+                                                'TEXT')
+
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters',
+                                                'cluster_hash',
+                                                'TEXT DEFAULT null')
+
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters',
+                                                'storage_mounts_metadata',
+                                                'BLOB DEFAULT null')
+        db_utils.add_column_to_table_sqlalchemy(
+            session,
+            'clusters',
+            'cluster_ever_up',
+            'INTEGER DEFAULT 0',
+            # Set the value to 1 so that all the existing clusters before #2977
+            # are considered as ever up, i.e:
+            #   existing cluster's default (null) -> 1;
+            #   new cluster's default -> 0;
+            # This is conservative for the existing clusters: even if some INIT
+            # clusters were never really UP, setting it to 1 means they won't be
+            # auto-deleted during any failover.
+            value_to_replace_existing_entries=1)
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters',
+                                                'status_updated_at',
+                                                'INTEGER DEFAULT null')
+        db_utils.add_column_to_table_sqlalchemy(
+            session,
+            'clusters',
+            'user_hash',
+            'TEXT DEFAULT null',
+            value_to_replace_existing_entries=common_utils.get_user_hash())
+        db_utils.add_column_to_table_sqlalchemy(session, 'clusters',
+                                                'config_hash',
+                                                'TEXT DEFAULT null')
+
+        db_utils.add_column_to_table_sqlalchemy(session, 'cluster_history',
+                                                'user_hash',
+                                                'TEXT DEFAULT null')
+
+        db_utils.add_column_to_table_sqlalchemy(
+            session,
+            'clusters',
+            'workspace',
+            'TEXT DEFAULT \'default\'',
+            value_to_replace_existing_entries=constants.
+            SKYPILOT_DEFAULT_WORKSPACE)
+        session.commit()
+
+
+create_table()
 
 
 def add_or_update_user(user: models.User):
     """Store the mapping from user hash to user name for display purposes."""
     if user.name is None:
         return
-    _DB.cursor.execute('INSERT OR REPLACE INTO users (id, name) VALUES (?, ?)',
-                       (user.id, user.name))
-    _DB.conn.commit()
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_stmnt = sqlite.insert(user_table).values(id=user.id,
+                                                            name=user.name)
+            do_update_stmt = insert_stmnt.on_conflict_do_update(
+                index_elements=[user_table.c.id],
+                set_={user_table.c.name: user.name})
+            session.execute(do_update_stmt)
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        else:
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        session.commit()
 
 
 def get_user(user_id: str) -> models.User:
-    row = _DB.cursor.execute('SELECT id, name FROM users WHERE id=?',
-                             (user_id,)).fetchone()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(user_table).filter_by(id=user_id).first()
     if row is None:
         return models.User(id=user_id)
-    return models.User(id=row[0], name=row[1])
+    return models.User(id=row.id, name=row.name)
 
 
 def get_all_users() -> List[models.User]:
-    rows = _DB.cursor.execute('SELECT id, name FROM users').fetchall()
-    return [models.User(id=row[0], name=row[1]) for row in rows]
+    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]
 
 
 def add_or_update_cluster(cluster_name: str,
@@ -209,6 +271,9 @@ def add_or_update_cluster(cluster_name: str,
         is_launch: if the cluster is firstly launched. If True, the launched_at
             and last_use will be updated. Otherwise, use the old value.
     """
+    # TODO(zhwu): have to be imported here to avoid circular import.
+    from sky import skypilot_config  # pylint: disable=import-outside-toplevel
+
     # 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
@@ -242,142 +307,118 @@ def add_or_update_cluster(cluster_name: str,
         usage_intervals.append((cluster_launched_at, None))
 
     user_hash = common_utils.get_user_hash()
+    active_workspace = skypilot_config.get_active_workspace()
 
-    _DB.cursor.execute(
-        'INSERT or REPLACE INTO clusters'
-        # All the fields need to exist here, even if they don't need
-        # be changed, as the INSERT OR REPLACE statement will replace
-        # the field of the existing row with the default value if not
-        # specified.
-        '(name, launched_at, handle, last_use, status, '
-        'autostop, to_down, metadata, owner, cluster_hash, '
-        'storage_mounts_metadata, cluster_ever_up, status_updated_at, '
-        'config_hash, user_hash) '
-        'VALUES ('
-        # name
-        '?, '
-        # launched_at
-        'COALESCE('
-        '?, (SELECT launched_at FROM clusters WHERE name=?)), '
-        # handle
-        '?, '
-        # last_use
-        'COALESCE('
-        '?, (SELECT last_use FROM clusters WHERE name=?)), '
-        # status
-        '?, '
-        # autostop
-        # Keep the old autostop value if it exists, otherwise set it to
-        # default -1.
-        'COALESCE('
-        '(SELECT autostop FROM clusters WHERE name=? AND status!=?), -1), '
-        # Keep the old to_down value if it exists, otherwise set it to
-        # default 0.
-        'COALESCE('
-        '(SELECT to_down FROM clusters WHERE name=? AND status!=?), 0),'
-        # Keep the old metadata value if it exists, otherwise set it to
-        # default {}.
-        'COALESCE('
-        '(SELECT metadata FROM clusters WHERE name=?), \'{}\'),'
-        # Keep the old owner value if it exists, otherwise set it to
-        # default null.
-        'COALESCE('
-        '(SELECT owner FROM clusters WHERE name=?), null),'
-        # cluster_hash
-        '?,'
-        # storage_mounts_metadata
-        'COALESCE('
-        '(SELECT storage_mounts_metadata FROM clusters WHERE name=?), null), '
-        # cluster_ever_up
-        '((SELECT cluster_ever_up FROM clusters WHERE name=?) OR ?), '
-        # status_updated_at
-        '?,'
-        # config_hash
-        'COALESCE(?, (SELECT config_hash FROM clusters WHERE name=?)),'
-        # user_hash: keep original user_hash if it exists
-        'COALESCE('
-        '(SELECT user_hash FROM clusters WHERE name=?), ?)'
-        ')',
-        (
-            # name
-            cluster_name,
-            # launched_at
-            cluster_launched_at,
-            cluster_name,
-            # handle
-            handle,
-            # last_use
-            last_use,
-            cluster_name,
-            # status
-            status.value,
-            # autostop
-            cluster_name,
-            status_lib.ClusterStatus.STOPPED.value,
-            # to_down
-            cluster_name,
-            status_lib.ClusterStatus.STOPPED.value,
-            # metadata
-            cluster_name,
-            # owner
-            cluster_name,
-            # cluster_hash
-            cluster_hash,
-            # storage_mounts_metadata
-            cluster_name,
-            # cluster_ever_up
-            cluster_name,
-            int(ready),
-            # status_updated_at
-            status_updated_at,
-            # config_hash
-            config_hash,
-            cluster_name,
-            # user_hash
-            cluster_name,
-            user_hash,
-        ))
-
-    launched_nodes = getattr(cluster_handle, 'launched_nodes', None)
-    launched_resources = getattr(cluster_handle, 'launched_resources', None)
-    _DB.cursor.execute(
-        'INSERT or REPLACE INTO cluster_history'
-        '(cluster_hash, name, num_nodes, requested_resources, '
-        'launched_resources, usage_intervals, user_hash) '
-        'VALUES ('
-        # hash
-        '?, '
-        # name
-        '?, '
-        # requested resources
-        '?, '
-        # launched resources
-        '?, '
-        # number of nodes
-        '?, '
-        # usage intervals
-        '?, '
-        # user_hash
-        '?'
-        ')',
-        (
-            # hash
-            cluster_hash,
-            # name
-            cluster_name,
-            # number of nodes
-            launched_nodes,
-            # requested resources
-            pickle.dumps(requested_resources),
-            # launched resources
-            pickle.dumps(launched_resources),
-            # usage intervals
-            pickle.dumps(usage_intervals),
-            # user_hash
-            user_hash,
-        ))
-
-    _DB.conn.commit()
+    conditional_values = {}
+    if is_launch:
+        conditional_values.update({
+            'launched_at': cluster_launched_at,
+            'last_use': last_use
+        })
+
+    if int(ready) == 1:
+        conditional_values.update({
+            'cluster_ever_up': 1,
+        })
+
+    if config_hash is not None:
+        conditional_values.update({
+            'config_hash': config_hash,
+        })
+
+    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(
+            name=cluster_name).with_for_update().first()
+        if (not cluster_row or
+                cluster_row.status == status_lib.ClusterStatus.STOPPED.value):
+            conditional_values.update({
+                'autostop': -1,
+                'to_down': 0,
+            })
+        if not cluster_row or not cluster_row.user_hash:
+            conditional_values.update({
+                'user_hash': user_hash,
+            })
+        if not cluster_row or not cluster_row.workspace:
+            conditional_values.update({
+                'workspace': active_workspace,
+            })
+
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_stmnt = sqlite.insert(cluster_table).values(
+                name=cluster_name,
+                **conditional_values,
+                handle=handle,
+                status=status.value,
+                # set metadata to server default ('{}')
+                # set owner to server default (null)
+                cluster_hash=cluster_hash,
+                # set storage_mounts_metadata to server default (null)
+                status_updated_at=status_updated_at,
+            )
+            do_update_stmt = insert_stmnt.on_conflict_do_update(
+                index_elements=[cluster_table.c.name],
+                set_={
+                    **conditional_values,
+                    cluster_table.c.handle: handle,
+                    cluster_table.c.status: status.value,
+                    # do not update metadata value
+                    # do not update owner value
+                    cluster_table.c.cluster_hash: cluster_hash,
+                    # do not update storage_mounts_metadata
+                    cluster_table.c.status_updated_at: status_updated_at,
+                    # do not update user_hash
+                })
+            session.execute(do_update_stmt)
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        else:
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+
+        # Modify cluster history table
+        launched_nodes = getattr(cluster_handle, 'launched_nodes', None)
+        launched_resources = getattr(cluster_handle, 'launched_resources', None)
+
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_stmnt = sqlite.insert(cluster_history_table).values(
+                cluster_hash=cluster_hash,
+                name=cluster_name,
+                num_nodes=launched_nodes,
+                requested_resources=pickle.dumps(requested_resources),
+                launched_resources=pickle.dumps(launched_resources),
+                usage_intervals=pickle.dumps(usage_intervals),
+                user_hash=user_hash)
+            do_update_stmt = insert_stmnt.on_conflict_do_update(
+                index_elements=[cluster_history_table.c.cluster_hash],
+                set_={
+                    cluster_history_table.c.name: cluster_name,
+                    cluster_history_table.c.num_nodes: launched_nodes,
+                    cluster_history_table.c.requested_resources:
+                        pickle.dumps(requested_resources),
+                    cluster_history_table.c.launched_resources:
+                        pickle.dumps(launched_resources),
+                    cluster_history_table.c.usage_intervals:
+                        pickle.dumps(usage_intervals),
+                    cluster_history_table.c.user_hash: user_hash
+                })
+            session.execute(do_update_stmt)
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        else:
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        session.commit()
 
 
 def _get_user_hash_or_current_user(user_hash: Optional[str]) -> str:
@@ -395,16 +436,18 @@ def _get_user_hash_or_current_user(user_hash: Optional[str]) -> str:
 def update_cluster_handle(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle'):
     handle = pickle.dumps(cluster_handle)
-    _DB.cursor.execute('UPDATE clusters SET handle=(?) WHERE name=(?)',
-                       (handle, cluster_name))
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.query(cluster_table).filter_by(name=cluster_name).update(
+            {cluster_table.c.handle: handle})
+        session.commit()
 
 
 def update_last_use(cluster_name: str):
     """Updates the last used command for the cluster."""
-    _DB.cursor.execute('UPDATE clusters SET last_use=(?) WHERE name=(?)',
-                       (common_utils.get_current_command(), cluster_name))
-    _DB.conn.commit()
+    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()
 
 
 def remove_cluster(cluster_name: str, terminate: bool) -> None:
@@ -412,63 +455,73 @@ def remove_cluster(cluster_name: str, terminate: bool) -> None:
     cluster_hash = _get_hash_for_existing_cluster(cluster_name)
     usage_intervals = _get_cluster_usage_intervals(cluster_hash)
 
-    # usage_intervals is not None and not empty
-    if usage_intervals:
-        assert cluster_hash is not None, cluster_name
-        start_time = usage_intervals.pop()[0]
-        end_time = int(time.time())
-        usage_intervals.append((start_time, end_time))
-        _set_cluster_usage_intervals(cluster_hash, usage_intervals)
-
-    if terminate:
-        _DB.cursor.execute('DELETE FROM clusters WHERE name=(?)',
-                           (cluster_name,))
-    else:
-        handle = get_handle_from_cluster_name(cluster_name)
-        if handle is None:
-            return
-        # Must invalidate IP list to avoid directly trying to ssh into a
-        # stopped VM, which leads to timeout.
-        if hasattr(handle, 'stable_internal_external_ips'):
-            handle = typing.cast('backends.CloudVmRayResourceHandle', handle)
-            handle.stable_internal_external_ips = None
-        current_time = int(time.time())
-        _DB.cursor.execute(
-            'UPDATE clusters SET handle=(?), status=(?), '
-            'status_updated_at=(?) WHERE name=(?)', (
-                pickle.dumps(handle),
-                status_lib.ClusterStatus.STOPPED.value,
-                current_time,
-                cluster_name,
-            ))
-    _DB.conn.commit()
+    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
+            start_time = usage_intervals.pop()[0]
+            end_time = int(time.time())
+            usage_intervals.append((start_time, end_time))
+            _set_cluster_usage_intervals(cluster_hash, usage_intervals)
+
+        if terminate:
+            session.query(cluster_table).filter_by(name=cluster_name).delete()
+        else:
+            handle = get_handle_from_cluster_name(cluster_name)
+            if handle is None:
+                return
+            # Must invalidate IP list to avoid directly trying to ssh into a
+            # stopped VM, which leads to timeout.
+            if hasattr(handle, 'stable_internal_external_ips'):
+                handle = typing.cast('backends.CloudVmRayResourceHandle',
+                                     handle)
+                handle.stable_internal_external_ips = None
+            current_time = int(time.time())
+            session.query(cluster_table).filter_by(name=cluster_name).update({
+                cluster_table.c.handle: pickle.dumps(handle),
+                cluster_table.c.status: status_lib.ClusterStatus.STOPPED.value,
+                cluster_table.c.status_updated_at: current_time
+            })
+        session.commit()
 
 
 def get_handle_from_cluster_name(
         cluster_name: str) -> Optional['backends.ResourceHandle']:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    rows = _DB.cursor.execute('SELECT handle FROM clusters WHERE name=(?)',
-                              (cluster_name,))
-    for (handle,) in rows:
-        return pickle.loads(handle)
-    return 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
+    return pickle.loads(row.handle)
 
 
 def get_glob_cluster_names(cluster_name: str) -> List[str]:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    rows = _DB.cursor.execute('SELECT name FROM clusters WHERE name GLOB (?)',
-                              (cluster_name,))
-    return [row[0] for row in rows]
+    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 ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            # postgres does not support GLOB
+            raise ValueError('Unsupported database dialect')
+        else:
+            raise ValueError('Unsupported database dialect')
+    return [row.name for row in rows]
 
 
 def set_cluster_status(cluster_name: str,
                        status: status_lib.ClusterStatus) -> None:
     current_time = int(time.time())
-    _DB.cursor.execute(
-        'UPDATE clusters SET status=(?), status_updated_at=(?) WHERE name=(?)',
-        (status.value, current_time, cluster_name))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        count = session.query(cluster_table).filter_by(
+            name=cluster_name).update({
+                cluster_table.c.status: status.value,
+                cluster_table.c.status_updated_at: current_time
+            })
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
@@ -476,46 +529,40 @@ def set_cluster_status(cluster_name: str,
 
 def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
                                to_down: bool) -> None:
-    _DB.cursor.execute(
-        'UPDATE clusters SET autostop=(?), to_down=(?) WHERE name=(?)', (
-            idle_minutes,
-            int(to_down),
-            cluster_name,
-        ))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        count = session.query(cluster_table).filter_by(
+            name=cluster_name).update({
+                cluster_table.c.autostop: idle_minutes,
+                cluster_table.c.to_down: int(to_down)
+            })
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
 def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
-    rows = _DB.cursor.execute('SELECT launched_at FROM clusters WHERE name=(?)',
-                              (cluster_name,))
-    for (launch_time,) in rows:
-        if launch_time is None:
-            return None
-        return int(launch_time)
-    return 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)
 
 
 def get_cluster_info(cluster_name: str) -> Optional[Dict[str, Any]]:
-    rows = _DB.cursor.execute('SELECT metadata FROM clusters WHERE name=(?)',
-                              (cluster_name,))
-    for (metadata,) in rows:
-        if metadata is None:
-            return None
-        return json.loads(metadata)
-    return 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)
 
 
 def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
-    _DB.cursor.execute('UPDATE clusters SET metadata=(?) WHERE name=(?)', (
-        json.dumps(metadata),
-        cluster_name,
-    ))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    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)})
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
@@ -523,25 +570,22 @@ def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
 
 def get_cluster_storage_mounts_metadata(
         cluster_name: str) -> Optional[Dict[str, Any]]:
-    rows = _DB.cursor.execute(
-        'SELECT storage_mounts_metadata FROM clusters WHERE name=(?)',
-        (cluster_name,))
-    for (storage_mounts_metadata,) in rows:
-        if storage_mounts_metadata is None:
-            return None
-        return pickle.loads(storage_mounts_metadata)
-    return 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)
 
 
 def set_cluster_storage_mounts_metadata(
         cluster_name: str, storage_mounts_metadata: Dict[str, Any]) -> None:
-    _DB.cursor.execute(
-        'UPDATE clusters SET storage_mounts_metadata=(?) WHERE name=(?)', (
-            pickle.dumps(storage_mounts_metadata),
-            cluster_name,
-        ))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        count = session.query(cluster_table).filter_by(
+            name=cluster_name).update({
+                cluster_table.c.storage_mounts_metadata:
+                    pickle.dumps(storage_mounts_metadata)
+            })
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
@@ -552,14 +596,12 @@ def _get_cluster_usage_intervals(
 ) -> Optional[List[Tuple[int, Optional[int]]]]:
     if cluster_hash is None:
         return None
-    rows = _DB.cursor.execute(
-        'SELECT usage_intervals FROM cluster_history WHERE cluster_hash=(?)',
-        (cluster_hash,))
-    for (usage_intervals,) in rows:
-        if usage_intervals is None:
-            return None
-        return pickle.loads(usage_intervals)
-    return 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 or row.usage_intervals is None:
+        return None
+    return pickle.loads(row.usage_intervals)
 
 
 def _get_cluster_launch_time(cluster_hash: str) -> Optional[int]:
@@ -591,15 +633,13 @@ def _get_cluster_duration(cluster_hash: str) -> int:
 def _set_cluster_usage_intervals(
         cluster_hash: str, usage_intervals: List[Tuple[int,
                                                        Optional[int]]]) -> None:
-    _DB.cursor.execute(
-        'UPDATE cluster_history SET usage_intervals=(?) WHERE cluster_hash=(?)',
-        (
-            pickle.dumps(usage_intervals),
-            cluster_hash,
-        ))
-
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    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:
+                    pickle.dumps(usage_intervals)
+            })
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster hash {cluster_hash} not found.')
@@ -610,38 +650,38 @@ def set_owner_identity_for_cluster(cluster_name: str,
     if owner_identity is None:
         return
     owner_identity_str = json.dumps(owner_identity)
-    _DB.cursor.execute('UPDATE clusters SET owner=(?) WHERE name=(?)',
-                       (owner_identity_str, cluster_name))
-
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    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})
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Cluster {cluster_name} not found.')
 
 
 def _get_hash_for_existing_cluster(cluster_name: str) -> Optional[str]:
-    rows = _DB.cursor.execute(
-        'SELECT cluster_hash FROM clusters WHERE name=(?)', (cluster_name,))
-    for (cluster_hash,) in rows:
-        if cluster_hash is None:
-            return None
-        return cluster_hash
-    return 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
 
 
 def get_launched_resources_from_cluster_hash(
         cluster_hash: str) -> Optional[Tuple[int, Any]]:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(cluster_history_table).filter_by(
+            cluster_hash=cluster_hash).first()
+    if row is None:
+        return None
+    num_nodes = row.num_nodes
+    launched_resources = row.launched_resources
 
-    rows = _DB.cursor.execute(
-        'SELECT num_nodes, launched_resources '
-        'FROM cluster_history WHERE cluster_hash=(?)', (cluster_hash,))
-    for (num_nodes, launched_resources) in rows:
-        if num_nodes is None or launched_resources is None:
-            return None
-        launched_resources = pickle.loads(launched_resources)
-        return num_nodes, launched_resources
-    return None
+    if num_nodes is None or launched_resources is None:
+        return None
+    launched_resources = pickle.loads(launched_resources)
+    return num_nodes, launched_resources
 
 
 def _load_owner(record_owner: Optional[str]) -> Optional[List[str]]:
@@ -675,74 +715,62 @@ def _load_storage_mounts_metadata(
 @context_utils.cancellation_guard
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
-    rows = _DB.cursor.execute(
-        'SELECT name, launched_at, handle, last_use, status, autostop, '
-        'metadata, to_down, owner, cluster_hash, storage_mounts_metadata, '
-        'cluster_ever_up, status_updated_at, config_hash, user_hash '
-        'FROM clusters WHERE name=(?)', (cluster_name,)).fetchall()
-    for row in rows:
-        # Explicitly specify the number of fields to unpack, so that
-        # we can add new fields to the database in the future without
-        # breaking the previous code.
-        (name, launched_at, handle, last_use, status, autostop, metadata,
-         to_down, owner, cluster_hash, storage_mounts_metadata, cluster_ever_up,
-         status_updated_at, config_hash, user_hash) = row
-        user_hash = _get_user_hash_or_current_user(user_hash)
-        # TODO: use namedtuple instead of dict
-        record = {
-            'name': name,
-            'launched_at': launched_at,
-            'handle': pickle.loads(handle),
-            'last_use': last_use,
-            'status': status_lib.ClusterStatus[status],
-            'autostop': autostop,
-            'to_down': bool(to_down),
-            'owner': _load_owner(owner),
-            'metadata': json.loads(metadata),
-            'cluster_hash': cluster_hash,
-            'storage_mounts_metadata':
-                _load_storage_mounts_metadata(storage_mounts_metadata),
-            'cluster_ever_up': bool(cluster_ever_up),
-            'status_updated_at': status_updated_at,
-            'user_hash': user_hash,
-            'user_name': get_user(user_hash).name,
-            'config_hash': config_hash,
-        }
-        return record
-    return 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
+    user_hash = _get_user_hash_or_current_user(row.user_hash)
+    # TODO: use namedtuple instead of dict
+    record = {
+        'name': row.name,
+        'launched_at': row.launched_at,
+        'handle': pickle.loads(row.handle),
+        'last_use': row.last_use,
+        'status': status_lib.ClusterStatus[row.status],
+        'autostop': row.autostop,
+        'to_down': bool(row.to_down),
+        'owner': _load_owner(row.owner),
+        'metadata': json.loads(row.metadata),
+        'cluster_hash': row.cluster_hash,
+        'storage_mounts_metadata': _load_storage_mounts_metadata(
+            row.storage_mounts_metadata),
+        'cluster_ever_up': bool(row.cluster_ever_up),
+        'status_updated_at': row.status_updated_at,
+        'user_hash': user_hash,
+        'user_name': get_user(user_hash).name,
+        'config_hash': row.config_hash,
+        'workspace': row.workspace,
+    }
+    return record
 
 
 def get_clusters() -> List[Dict[str, Any]]:
-    rows = _DB.cursor.execute(
-        'select name, launched_at, handle, last_use, status, autostop, '
-        'metadata, to_down, owner, cluster_hash, storage_mounts_metadata, '
-        'cluster_ever_up, status_updated_at, config_hash, user_hash '
-        'from clusters order by launched_at desc').fetchall()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.query(cluster_table).order_by(
+            sqlalchemy.desc(cluster_table.c.launched_at)).all()
     records = []
     for row in rows:
-        (name, launched_at, handle, last_use, status, autostop, metadata,
-         to_down, owner, cluster_hash, storage_mounts_metadata, cluster_ever_up,
-         status_updated_at, config_hash, user_hash) = row
-        user_hash = _get_user_hash_or_current_user(user_hash)
+        user_hash = _get_user_hash_or_current_user(row.user_hash)
         # TODO: use namedtuple instead of dict
         record = {
-            'name': name,
-            'launched_at': launched_at,
-            'handle': pickle.loads(handle),
-            'last_use': last_use,
-            'status': status_lib.ClusterStatus[status],
-            'autostop': autostop,
-            'to_down': bool(to_down),
-            'owner': _load_owner(owner),
-            'metadata': json.loads(metadata),
-            'cluster_hash': cluster_hash,
-            'storage_mounts_metadata':
-                _load_storage_mounts_metadata(storage_mounts_metadata),
-            'cluster_ever_up': bool(cluster_ever_up),
-            'status_updated_at': status_updated_at,
+            'name': row.name,
+            'launched_at': row.launched_at,
+            'handle': pickle.loads(row.handle),
+            'last_use': row.last_use,
+            'status': status_lib.ClusterStatus[row.status],
+            'autostop': row.autostop,
+            'to_down': bool(row.to_down),
+            'owner': _load_owner(row.owner),
+            'metadata': json.loads(row.metadata),
+            'cluster_hash': row.cluster_hash,
+            'storage_mounts_metadata': _load_storage_mounts_metadata(
+                row.storage_mounts_metadata),
+            'cluster_ever_up': bool(row.cluster_ever_up),
+            'status_updated_at': row.status_updated_at,
             'user_hash': user_hash,
             'user_name': get_user(user_hash).name,
-            'config_hash': config_hash,
+            'config_hash': row.config_hash,
+            'workspace': row.workspace,
         }
 
         records.append(record)
@@ -750,43 +778,30 @@ def get_clusters() -> List[Dict[str, Any]]:
 
 
 def get_clusters_from_history() -> List[Dict[str, Any]]:
-    rows = _DB.cursor.execute(
-        'SELECT ch.cluster_hash, ch.name, ch.num_nodes, '
-        'ch.launched_resources, ch.usage_intervals, clusters.status, '
-        'ch.user_hash  '
-        'FROM cluster_history ch '
-        'LEFT OUTER JOIN clusters '
-        'ON ch.cluster_hash=clusters.cluster_hash ').fetchall()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.query(
+            cluster_history_table.join(cluster_table,
+                                       cluster_history_table.c.cluster_hash ==
+                                       cluster_table.c.cluster_hash,
+                                       isouter=True)).all()
 
     # '(cluster_hash, name, num_nodes, requested_resources, '
     #         'launched_resources, usage_intervals) '
     records = []
-
     for row in rows:
         # TODO: use namedtuple instead of dict
-
-        (
-            cluster_hash,
-            name,
-            num_nodes,
-            launched_resources,
-            usage_intervals,
-            status,
-            user_hash,
-        ) = row[:7]
-        user_hash = _get_user_hash_or_current_user(user_hash)
-
+        user_hash = _get_user_hash_or_current_user(row.user_hash)
+        status = row.status
         if status is not None:
             status = status_lib.ClusterStatus[status]
-
         record = {
-            'name': name,
-            'launched_at': _get_cluster_launch_time(cluster_hash),
-            'duration': _get_cluster_duration(cluster_hash),
-            'num_nodes': num_nodes,
-            'resources': pickle.loads(launched_resources),
-            'cluster_hash': cluster_hash,
-            'usage_intervals': pickle.loads(usage_intervals),
+            'name': row.name,
+            'launched_at': _get_cluster_launch_time(row.cluster_hash),
+            'duration': _get_cluster_duration(row.cluster_hash),
+            'num_nodes': row.num_nodes,
+            'resources': pickle.loads(row.launched_resources),
+            'cluster_hash': row.cluster_hash,
+            'usage_intervals': pickle.loads(row.usage_intervals),
             'status': status,
             'user_hash': user_hash,
         }
@@ -799,29 +814,29 @@ def get_clusters_from_history() -> List[Dict[str, Any]]:
 
 
 def get_cluster_names_start_with(starts_with: str) -> List[str]:
-    rows = _DB.cursor.execute('SELECT name FROM clusters WHERE name LIKE (?)',
-                              (f'{starts_with}%',))
-    return [row[0] for row in rows]
-
+    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]
 
-def get_cached_enabled_clouds(
-        cloud_capability: 'cloud.CloudCapability') -> List['clouds.Cloud']:
 
-    rows = _DB.cursor.execute('SELECT value FROM config WHERE key = ?',
-                              (_get_capability_key(cloud_capability),))
+def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
+                              workspace: str) -> List['clouds.Cloud']:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(config_table).filter_by(
+            key=_get_enabled_clouds_key(cloud_capability, workspace)).first()
     ret = []
-    for (value,) in rows:
-        ret = json.loads(value)
-        break
+    if row:
+        ret = json.loads(row.value)
     enabled_clouds: List['clouds.Cloud'] = []
     for c in ret:
         try:
             cloud = registry.CLOUD_REGISTRY.from_str(c)
         except ValueError:
-            # Handle the case for the clouds whose support has been removed from
-            # SkyPilot, e.g., 'local' was a cloud in the past and may be stored
-            # in the database for users before #3037. We should ignore removed
-            # clouds and continue.
+            # Handle the case for the clouds whose support has been
+            # removed from SkyPilot, e.g., 'local' was a cloud in the past
+            # and may be stored in the database for users before #3037.
+            # We should ignore removed clouds and continue.
             continue
         if cloud is not None:
             enabled_clouds.append(cloud)
@@ -829,15 +844,32 @@ def get_cached_enabled_clouds(
 
 
 def set_enabled_clouds(enabled_clouds: List[str],
-                       cloud_capability: 'cloud.CloudCapability') -> None:
-    _DB.cursor.execute(
-        'INSERT OR REPLACE INTO config VALUES (?, ?)',
-        (_get_capability_key(cloud_capability), json.dumps(enabled_clouds)))
-    _DB.conn.commit()
-
-
-def _get_capability_key(cloud_capability: 'cloud.CloudCapability') -> str:
-    return _ENABLED_CLOUDS_KEY_PREFIX + cloud_capability.value
+                       cloud_capability: 'cloud.CloudCapability',
+                       workspace: str) -> None:
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_stmnt = sqlite.insert(config_table).values(
+                key=_get_enabled_clouds_key(cloud_capability, workspace),
+                value=json.dumps(enabled_clouds))
+            do_update_stmt = insert_stmnt.on_conflict_do_update(
+                index_elements=[config_table.c.key],
+                set_={config_table.c.value: json.dumps(enabled_clouds)})
+            session.execute(do_update_stmt)
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        else:
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        session.commit()
+
+
+def _get_enabled_clouds_key(cloud_capability: 'cloud.CloudCapability',
+                            workspace: str) -> str:
+    return _ENABLED_CLOUDS_KEY_PREFIX + workspace + '_' + cloud_capability.value
 
 
 def add_or_update_storage(storage_name: str,
@@ -853,26 +885,48 @@ 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')
-    _DB.cursor.execute('INSERT OR REPLACE INTO storage VALUES (?, ?, ?, ?, ?)',
-                       (storage_name, storage_launched_at, handle, last_use,
-                        storage_status.value))
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_stmnt = sqlite.insert(storage_table).values(
+                name=storage_name,
+                handle=handle,
+                last_use=last_use,
+                launched_at=storage_launched_at,
+                status=storage_status.value)
+            do_update_stmt = insert_stmnt.on_conflict_do_update(
+                index_elements=[storage_table.c.name],
+                set_={
+                    storage_table.c.handle: handle,
+                    storage_table.c.last_use: last_use,
+                    storage_table.c.launched_at: storage_launched_at,
+                    storage_table.c.status: storage_status.value
+                })
+            session.execute(do_update_stmt)
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        else:
+            session.rollback()
+            raise ValueError('Unsupported database dialect')
+        session.commit()
 
 
 def remove_storage(storage_name: str):
     """Removes Storage from Database"""
-    _DB.cursor.execute('DELETE FROM storage WHERE name=(?)', (storage_name,))
-    _DB.conn.commit()
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.query(storage_table).filter_by(name=storage_name).delete()
+        session.commit()
 
 
 def set_storage_status(storage_name: str,
                        status: status_lib.StorageStatus) -> None:
-    _DB.cursor.execute('UPDATE storage SET status=(?) WHERE name=(?)', (
-        status.value,
-        storage_name,
-    ))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    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()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Storage {storage_name} not found.')
@@ -880,21 +934,20 @@ def set_storage_status(storage_name: str,
 
 def get_storage_status(storage_name: str) -> Optional[status_lib.StorageStatus]:
     assert storage_name is not None, 'storage_name cannot be None'
-    rows = _DB.cursor.execute('SELECT status FROM storage WHERE name=(?)',
-                              (storage_name,))
-    for (status,) in rows:
-        return status_lib.StorageStatus[status]
+    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
 
 
 def set_storage_handle(storage_name: str,
                        handle: 'Storage.StorageMetadata') -> None:
-    _DB.cursor.execute('UPDATE storage SET handle=(?) WHERE name=(?)', (
-        pickle.dumps(handle),
-        storage_name,
-    ))
-    count = _DB.cursor.rowcount
-    _DB.conn.commit()
+    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)})
+        session.commit()
     assert count <= 1, count
     if count == 0:
         raise ValueError(f'Storage{storage_name} not found.')
@@ -904,38 +957,48 @@ def get_handle_from_storage_name(
         storage_name: Optional[str]) -> Optional['Storage.StorageMetadata']:
     if storage_name is None:
         return None
-    rows = _DB.cursor.execute('SELECT handle FROM storage WHERE name=(?)',
-                              (storage_name,))
-    for (handle,) in rows:
-        if handle is None:
-            return None
-        return pickle.loads(handle)
+    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
 
 
 def get_glob_storage_name(storage_name: str) -> List[str]:
     assert storage_name is not None, 'storage_name cannot be None'
-    rows = _DB.cursor.execute('SELECT name FROM storage WHERE name GLOB (?)',
-                              (storage_name,))
-    return [row[0] for row in rows]
+    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 ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # TODO(syang) support postgres dialect
+            # postgres does not support GLOB
+            raise ValueError('Unsupported database dialect')
+        else:
+            raise ValueError('Unsupported database dialect')
+    return [row.name for row in rows]
 
 
 def get_storage_names_start_with(starts_with: str) -> List[str]:
-    rows = _DB.cursor.execute('SELECT name FROM storage WHERE name LIKE (?)',
-                              (f'{starts_with}%',))
-    return [row[0] for row in rows]
+    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]
 
 
 def get_storage() -> List[Dict[str, Any]]:
-    rows = _DB.cursor.execute('SELECT * FROM storage')
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.query(storage_table).all()
     records = []
-    for name, launched_at, handle, last_use, status in rows:
+    for row in rows:
         # TODO: use namedtuple instead of dict
         records.append({
-            'name': name,
-            'launched_at': launched_at,
-            'handle': pickle.loads(handle),
-            'last_use': last_use,
-            'status': status_lib.StorageStatus[status],
+            'name': row.name,
+            'launched_at': row.launched_at,
+            'handle': pickle.loads(row.handle),
+            'last_use': row.last_use,
+            'status': status_lib.StorageStatus[row.status],
         })
     return records