MemoryOS 2.0.3__py3-none-any.whl

memos/mem_user/mysql_persistent_user_manager.py

@@ -0,0 +1,271 @@
+"""Persistent user management system for MemOS with configuration storage.
+
+This module extends the MySQL UserManager to provide persistent storage
+for user configurations and MOS instances.
+"""
+
+import json
+
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import Column, String, Text
+
+from memos.configs.mem_os import MOSConfig
+from memos.log import get_logger
+from memos.mem_user.mysql_user_manager import Base, MySQLUserManager
+
+
+logger = get_logger(__name__)
+
+
+class UserConfig(Base):
+    """User configuration model for the database."""
+
+    __tablename__ = "user_configs"
+
+    user_id = Column(String(255), primary_key=True)
+    config_data = Column(Text, nullable=False)  # JSON string of MOSConfig
+    created_at = Column(String(50), nullable=False)  # ISO format timestamp
+    updated_at = Column(String(50), nullable=False)  # ISO format timestamp
+
+    def __repr__(self):
+        return f"<UserConfig(user_id='{self.user_id}')>"
+
+
+class MySQLPersistentUserManager(MySQLUserManager):
+    """Extended MySQLUserManager with configuration persistence."""
+
+    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 persistent user manager.
+
+        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".
+        """
+        super().__init__(user_id, host, port, username, password, database, charset)
+
+        # Create user_configs table
+        Base.metadata.create_all(bind=self.engine)
+        logger.info("MySQLPersistentUserManager initialized with configuration storage")
+
+    def _convert_datetime_strings(self, obj: Any) -> Any:
+        """Recursively convert datetime strings back to datetime objects in config dict.
+
+        Args:
+            obj: The object to process (dict, list, or primitive type)
+
+        Returns:
+            The object with datetime strings converted to datetime objects
+        """
+        if isinstance(obj, dict):
+            result = {}
+            for key, value in obj.items():
+                if key == "created_at" and isinstance(value, str):
+                    try:
+                        result[key] = datetime.fromisoformat(value)
+                    except ValueError:
+                        # If parsing fails, keep the original string
+                        result[key] = value
+                else:
+                    result[key] = self._convert_datetime_strings(value)
+            return result
+        elif isinstance(obj, list):
+            return [self._convert_datetime_strings(item) for item in obj]
+        else:
+            return obj
+
+    def save_user_config(self, user_id: str, config: MOSConfig) -> bool:
+        """Save user configuration to database.
+
+        Args:
+            user_id (str): The user ID.
+            config (MOSConfig): The user's MOS configuration.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            # Convert config to JSON string with proper datetime handling
+            config_dict = config.model_dump(mode="json")
+            config_json = json.dumps(config_dict, indent=2)
+
+            now = datetime.now().isoformat()
+
+            # Check if config already exists
+            existing_config = (
+                session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+            )
+
+            if existing_config:
+                # Update existing config
+                existing_config.config_data = config_json
+                existing_config.updated_at = now
+                logger.info(f"Updated configuration for user {user_id}")
+            else:
+                # Create new config
+                user_config = UserConfig(
+                    user_id=user_id, config_data=config_json, created_at=now, updated_at=now
+                )
+                session.add(user_config)
+                logger.info(f"Saved new configuration for user {user_id}")
+
+            session.commit()
+            return True
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error saving user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def get_user_config(self, user_id: str) -> MOSConfig | None:
+        """Get user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            MOSConfig | None: The user's configuration or None if not found.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                config_dict = json.loads(user_config.config_data)
+                # Convert datetime strings back to datetime objects
+                config_dict = self._convert_datetime_strings(config_dict)
+                return MOSConfig(**config_dict)
+            return None
+
+        except Exception as e:
+            logger.error(f"Error loading user config for {user_id}: {e}")
+            return None
+        finally:
+            session.close()
+
+    def delete_user_config(self, user_id: str) -> bool:
+        """Delete user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                session.delete(user_config)
+                session.commit()
+                logger.info(f"Deleted configuration for user {user_id}")
+                return True
+            return False
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error deleting user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def list_user_configs(self, limit: int = 1) -> dict[str, MOSConfig]:
+        """List all user configurations.
+
+        Returns:
+            Dict[str, MOSConfig]: Dictionary mapping user_id to MOSConfig.
+        """
+        session = self._get_session()
+        try:
+            user_configs = session.query(UserConfig).limit(limit).all()
+            result = {}
+
+            for user_config in user_configs:
+                try:
+                    config_dict = json.loads(user_config.config_data)
+                    # Convert datetime strings back to datetime objects
+                    config_dict = self._convert_datetime_strings(config_dict)
+                    result[user_config.user_id] = MOSConfig(**config_dict)
+                except Exception as e:
+                    logger.error(f"Error parsing config for user {user_config.user_id}: {e}")
+                    continue
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error listing user configs: {e}")
+            return {}
+        finally:
+            session.close()
+
+    def create_user_with_config(
+        self, user_name: str, config: MOSConfig, role=None, user_id: str | None = None
+    ) -> str:
+        """Create a new user with configuration.
+
+        Args:
+            user_name (str): Name of the user.
+            config (MOSConfig): The user's configuration.
+            role: User role (optional, uses default from UserManager).
+            user_id (str, optional): Custom user ID.
+
+        Returns:
+            str: The created user ID.
+
+        Raises:
+            ValueError: If user_name already exists.
+        """
+        # Create user using parent method
+        created_user_id = self.create_user(user_name, role, user_id)
+
+        # Save configuration
+        if not self.save_user_config(created_user_id, config):
+            logger.error(f"Failed to save configuration for user {created_user_id}")
+
+        return created_user_id
+
+    def delete_user(self, user_id: str) -> bool:
+        """Delete a user and their configuration.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        # Delete configuration first
+        self.delete_user_config(user_id)
+
+        # Delete user using parent method
+        return super().delete_user(user_id)
+
+    def get_user_cube_access(self, user_id: str) -> list[str]:
+        """Get list of cube IDs that a user has access to.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            list[str]: List of cube IDs the user can access.
+        """
+        cubes = self.get_user_cubes(user_id)
+        return [cube.cube_id for cube in cubes]