MemoryOS 2.0.3__py3-none-any.whl

memos/mem_user/mysql_user_manager.py

@@ -0,0 +1,502 @@
+"""User management system for MemOS.
+
+This module provides user authentication, authorization, and cube management
+functionality using SQLAlchemy and MySQL.
+"""
+
+import uuid
+
+from datetime import datetime
+from enum import Enum
+
+from sqlalchemy import (
+    Boolean,
+    Column,
+    DateTime,
+    ForeignKey,
+    String,
+    Table,
+    create_engine,
+)
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Session, declarative_base, relationship, sessionmaker
+
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+Base = declarative_base()
+
+
+class UserRole(Enum):
+    """User roles enumeration."""
+
+    ROOT = "ROOT"
+    ADMIN = "ADMIN"
+    USER = "USER"
+    GUEST = "GUEST"
+
+
+# Association table for many-to-many relationship between users and cubes
+user_cube_association = Table(
+    "user_cube_association",
+    Base.metadata,
+    Column("user_id", String(255), ForeignKey("users.user_id"), primary_key=True),
+    Column("cube_id", String(255), ForeignKey("cubes.cube_id"), primary_key=True),
+    Column("created_at", DateTime, default=datetime.now),
+)
+
+
+class User(Base):
+    """User model for the database."""
+
+    __tablename__ = "users"
+
+    user_id = Column(String(255), primary_key=True, default=lambda: str(uuid.uuid4()))
+    user_name = Column(String(255), unique=True, nullable=False)
+    role = Column(
+        String(20), default=UserRole.USER.value, nullable=False
+    )  # for sqlite backend this is SQLEnum
+    created_at = Column(DateTime, default=datetime.now, nullable=False)
+    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now, nullable=False)
+    is_active = Column(Boolean, default=True, nullable=False)
+
+    # Relationship with cubes
+    cubes = relationship("Cube", secondary=user_cube_association, back_populates="users")
+    owned_cubes = relationship("Cube", back_populates="owner", cascade="all, delete-orphan")
+
+    def __repr__(self):
+        return f"<User(user_id='{self.user_id}', user_name='{self.user_name}', role='{self.role}')>"
+
+
+class Cube(Base):
+    """Cube model for the database."""
+
+    __tablename__ = "cubes"
+
+    cube_id = Column(String(255), primary_key=True, default=lambda: str(uuid.uuid4()))
+    cube_name = Column(String(255), nullable=False)
+    cube_path = Column(String(500), nullable=True)  # Local path or remote repo
+    owner_id = Column(String(255), ForeignKey("users.user_id"), nullable=False)
+    created_at = Column(DateTime, default=datetime.now, nullable=False)
+    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now, nullable=False)
+    is_active = Column(Boolean, default=True, nullable=False)
+
+    # Relationships
+    owner = relationship("User", back_populates="owned_cubes")
+    users = relationship("User", secondary=user_cube_association, back_populates="cubes")
+
+    def __repr__(self):
+        return f"<Cube(cube_id='{self.cube_id}', cube_name='{self.cube_name}', owner_id='{self.owner_id}')>"
+
+
+class MySQLUserManager:
+    """User management system for MemOS using MySQL."""
+
+    def __init__(
+        self,
+        user_id: str = "root",
+        host: str = "localhost",
+        port: int = 3306,
+        username: str = "root",
+        password: str = "",
+        database: str = "memos_users",
+        charset: str = "utf8mb4",
+    ):
+        """Initialize the user manager with MySQL database connection.
+
+        Args:
+            user_id (str, optional): User ID. If None, uses default user ID.
+            host (str): MySQL server host. Defaults to "localhost".
+            port (int): MySQL server port. Defaults to 3306.
+            username (str): MySQL username. Defaults to "root".
+            password (str): MySQL password. Defaults to "".
+            database (str): MySQL database name. Defaults to "memos_users".
+            charset (str): MySQL charset. Defaults to "utf8mb4".
+        """
+        # Build MySQL connection URL
+        if password:
+            connection_url = (
+                f"mysql+pymysql://{username}:{password}@{host}:{port}/{database}?charset={charset}"
+            )
+        else:
+            connection_url = (
+                f"mysql+pymysql://{username}@{host}:{port}/{database}?charset={charset}"
+            )
+
+        self.connection_url = connection_url
+        self.engine = create_engine(connection_url, echo=False, pool_pre_ping=True)
+        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
+
+        # Create tables
+        Base.metadata.create_all(bind=self.engine)
+
+        # Initialize with root user if no users exist
+        self._init_root_user(user_id)
+
+        logger.info(f"MySQLUserManager initialized with database at {host}:{port}/{database}")
+
+    def _get_session(self) -> Session:
+        """Get a database session."""
+        return self.SessionLocal()
+
+    def _init_root_user(self, user_id: str) -> None:
+        """Initialize the root user if no users exist."""
+        session = self._get_session()
+        try:
+            # Check if any users exist
+            user_count = session.query(User).count()
+            if user_count == 0:
+                root_user = User(user_id=user_id, user_name=user_id, role=UserRole.ROOT)
+                session.add(root_user)
+                session.commit()
+                logger.info("Root user created successfully")
+            else:
+                self.create_user(user_name=user_id, user_id=user_id, role=UserRole.ROOT)
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Failed to create {user_id} user: {e}")
+        finally:
+            session.close()
+
+    def create_user(
+        self, user_name: str, role: UserRole = UserRole.USER, user_id: str | None = None
+    ) -> str:
+        """Create a new user.
+
+        Args:
+            user_name (str): Name of the user.
+            role (UserRole): Role of the user.
+            user_id (str, optional): Custom user ID. If None, generates UUID.
+
+        Returns:
+            str: The created user ID.
+
+        Raises:
+            ValueError: If user_name already exists.
+        """
+        session = self._get_session()
+        try:
+            # Check if user_name already exists
+            existing_user = session.query(User).filter(User.user_name == user_name).first()
+            if existing_user:
+                logger.info(f"User with name '{user_name}' already exists")
+                return existing_user.user_id
+            user = User(user_name=user_name, role=role.value, user_id=user_id or str(uuid.uuid4()))
+            session.add(user)
+            session.commit()
+            logger.info(f"User '{user_name}' created with ID: {user.user_id}")
+            return user.user_id
+        except IntegrityError:
+            session.rollback()
+            logger.info(f"Failed to create user with name '{user_name}' already exists")
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error creating user: {e}")
+            raise
+        finally:
+            session.close()
+
+    def get_user(self, user_id: str) -> User | None:
+        """Get user by ID.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            User: The user object or None if not found.
+        """
+        session = self._get_session()
+        try:
+            return session.query(User).filter(User.user_id == user_id).first()
+        finally:
+            session.close()
+
+    def get_user_by_name(self, user_name: str) -> User | None:
+        """Get user by name.
+
+        Args:
+            user_name (str): The user name.
+
+        Returns:
+            User: The user object or None if not found.
+        """
+        session = self._get_session()
+        try:
+            return session.query(User).filter(User.user_name == user_name).first()
+        finally:
+            session.close()
+
+    def validate_user(self, user_id: str) -> bool:
+        """Validate if a user exists and is active.
+
+        Args:
+            user_id (str): The user ID to validate.
+
+        Returns:
+            bool: True if user exists and is active, False otherwise.
+        """
+        user = self.get_user(user_id)
+        return user is not None and user.is_active
+
+    def list_users(self) -> list[User]:
+        """List all active users.
+
+        Returns:
+            list[User]: List of all active users.
+        """
+        session = self._get_session()
+        try:
+            return session.query(User).filter(User.is_active).all()
+        finally:
+            session.close()
+
+    def create_cube(
+        self,
+        cube_name: str,
+        owner_id: str,
+        cube_path: str | None = None,
+        cube_id: str | None = None,
+    ) -> str:
+        """Create a new cube.
+
+        Args:
+            cube_name (str): Name of the cube.
+            owner_id (str): ID of the cube owner.
+            cube_path (str, optional): Path to the cube.
+            cube_id (str, optional): Custom cube ID. If None, generates UUID.
+
+        Returns:
+            str: The created cube ID.
+
+        Raises:
+            ValueError: If owner doesn't exist.
+        """
+        session = self._get_session()
+        try:
+            # Validate owner exists
+            owner = session.query(User).filter(User.user_id == owner_id).first()
+            if not owner:
+                raise ValueError(f"User with ID '{owner_id}' does not exist")
+
+            cube = Cube(
+                cube_name=cube_name,
+                owner_id=owner_id,
+                cube_path=cube_path,
+                cube_id=cube_id or str(uuid.uuid4()),
+            )
+            session.add(cube)
+
+            # Add owner to cube users
+            cube.users.append(owner)
+
+            session.commit()
+            logger.info(f"Cube '{cube_name}' created with ID: {cube.cube_id}")
+            return cube.cube_id
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error creating cube: {e}")
+            raise
+        finally:
+            session.close()
+
+    def get_cube(self, cube_id: str) -> Cube | None:
+        """Get cube by ID.
+
+        Args:
+            cube_id (str): The cube ID.
+
+        Returns:
+            Cube: The cube object or None if not found.
+        """
+        session = self._get_session()
+        try:
+            return session.query(Cube).filter(Cube.cube_id == cube_id).first()
+        finally:
+            session.close()
+
+    def validate_user_cube_access(self, user_id: str, cube_id: str) -> bool:
+        """Validate if a user has access to a cube.
+
+        Args:
+            user_id (str): The user ID.
+            cube_id (str): The cube ID.
+
+        Returns:
+            bool: True if user has access to cube, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            # Check if user exists and is active
+            user = session.query(User).filter(User.user_id == user_id, User.is_active).first()
+            if not user:
+                return False
+
+            # Check if cube exists and is active
+            cube = session.query(Cube).filter(Cube.cube_id == cube_id, Cube.is_active).first()
+            if not cube:
+                return False
+
+            # Check if user has access to cube (owner or in users list)
+            if cube.owner_id == user_id:
+                return True
+
+            # Check many-to-many relationship
+            return user in cube.users
+        finally:
+            session.close()
+
+    def get_user_cubes(self, user_id: str) -> list[Cube]:
+        """Get all cubes accessible by a user.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            list[Cube]: List of cubes accessible by the user.
+        """
+        session = self._get_session()
+        try:
+            user = session.query(User).filter(User.user_id == user_id).first()
+            if not user:
+                return []
+
+            active_cubes = [cube for cube in user.cubes if cube.is_active]
+            return sorted(active_cubes, key=lambda cube: cube.created_at, reverse=True)
+        finally:
+            session.close()
+
+    def add_user_to_cube(self, user_id: str, cube_id: str) -> bool:
+        """Add a user to a cube's access list.
+
+        Args:
+            user_id (str): The user ID.
+            cube_id (str): The cube ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user = session.query(User).filter(User.user_id == user_id).first()
+            cube = session.query(Cube).filter(Cube.cube_id == cube_id).first()
+
+            if not user or not cube:
+                return False
+
+            if user not in cube.users:
+                cube.users.append(user)
+                session.commit()
+                logger.info(f"User '{user_id}' added to cube '{cube_id}'")
+
+            return True
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error adding user to cube: {e}")
+            return False
+        finally:
+            session.close()
+
+    def remove_user_from_cube(self, user_id: str, cube_id: str) -> bool:
+        """Remove a user from a cube's access list.
+
+        Args:
+            user_id (str): The user ID.
+            cube_id (str): The cube ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user = session.query(User).filter(User.user_id == user_id).first()
+            cube = session.query(Cube).filter(Cube.cube_id == cube_id).first()
+
+            if not user or not cube:
+                return False
+
+            # Don't remove owner
+            if cube.owner_id == user_id:
+                logger.warning(f"Cannot remove owner '{user_id}' from cube '{cube_id}'")
+                return False
+
+            if user in cube.users:
+                cube.users.remove(user)
+                session.commit()
+                logger.info(f"User '{user_id}' removed from cube '{cube_id}'")
+
+            return True
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error removing user from cube: {e}")
+            return False
+        finally:
+            session.close()
+
+    def delete_user(self, user_id: str) -> bool:
+        """Soft delete a user (set is_active to False).
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user = session.query(User).filter(User.user_id == user_id).first()
+            if not user:
+                return False
+
+            # Don't delete root user
+            if user.role == UserRole.ROOT:
+                logger.warning("Cannot delete root user")
+                return False
+
+            user.is_active = False
+            session.commit()
+            logger.info(f"User '{user_id}' deactivated")
+            return True
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error deleting user: {e}")
+            return False
+        finally:
+            session.close()
+
+    def delete_cube(self, cube_id: str) -> bool:
+        """Soft delete a cube (set is_active to False).
+
+        Args:
+            cube_id (str): The cube ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            cube = session.query(Cube).filter(Cube.cube_id == cube_id).first()
+            if not cube:
+                return False
+
+            cube.is_active = False
+            session.commit()
+            logger.info(f"Cube '{cube_id}' deactivated")
+            return True
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error deleting cube: {e}")
+            return False
+        finally:
+            session.close()
+
+    def close(self) -> None:
+        """Close the database engine and dispose of all connections.
+
+        This method should be called when the MySQLUserManager is no longer needed
+        to ensure proper cleanup of database connections.
+        """
+        if hasattr(self, "engine"):
+            self.engine.dispose()
+            logger.info("MySQLUserManager database connections closed")

memos/mem_user/persistent_factory.py

@@ -0,0 +1,98 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_user import UserManagerConfigFactory
+from memos.mem_user.mysql_persistent_user_manager import MySQLPersistentUserManager
+from memos.mem_user.persistent_user_manager import PersistentUserManager
+from memos.mem_user.redis_persistent_user_manager import RedisPersistentUserManager
+
+
+class PersistentUserManagerFactory:
+    """Factory class for creating persistent user manager instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "sqlite": PersistentUserManager,
+        "mysql": MySQLPersistentUserManager,
+        "redis": RedisPersistentUserManager,
+    }
+
+    @classmethod
+    def from_config(
+        cls, config_factory: UserManagerConfigFactory
+    ) -> PersistentUserManager | MySQLPersistentUserManager:
+        """Create a persistent user manager instance from configuration.
+
+        Args:
+            config_factory: Configuration factory containing backend and config
+
+        Returns:
+            Persistent user manager instance
+
+        Raises:
+            ValueError: If backend is not supported
+        """
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid persistent user manager backend: {backend}")
+
+        user_manager_class = cls.backend_to_class[backend]
+        config = config_factory.config
+
+        # Use model_dump() to convert Pydantic model to dict and unpack as kwargs
+        return user_manager_class(**config.model_dump())
+
+    @classmethod
+    def create_sqlite(
+        cls, db_path: str | None = None, user_id: str = "root"
+    ) -> PersistentUserManager:
+        """Create SQLite persistent user manager with default configuration.
+
+        Args:
+            db_path: Path to SQLite database file
+            user_id: Default user ID for initialization
+
+        Returns:
+            SQLite persistent user manager instance
+        """
+        config_factory = UserManagerConfigFactory(
+            backend="sqlite", config={"db_path": db_path, "user_id": user_id}
+        )
+        return cls.from_config(config_factory)
+
+    @classmethod
+    def create_mysql(
+        cls,
+        user_id: str = "root",
+        host: str = "localhost",
+        port: int = 3306,
+        username: str = "root",
+        password: str = "",
+        database: str = "memos_users",
+        charset: str = "utf8mb4",
+    ) -> MySQLPersistentUserManager:
+        """Create MySQL persistent user manager with specified configuration.
+
+        Args:
+            user_id: Default user ID for initialization
+            host: MySQL server host
+            port: MySQL server port
+            username: MySQL username
+            password: MySQL password
+            database: MySQL database name
+            charset: MySQL charset
+
+        Returns:
+            MySQL persistent user manager instance
+        """
+        config_factory = UserManagerConfigFactory(
+            backend="mysql",
+            config={
+                "user_id": user_id,
+                "host": host,
+                "port": port,
+                "username": username,
+                "password": password,
+                "database": database,
+                "charset": charset,
+            },
+        )
+        return cls.from_config(config_factory)