Some SQLAlchemy session cleanup

base_plugin.py

@@ -3,11 +3,11 @@ class BasePlugin(object):
     Defines an interface for all plugins to inherit from. Note that the __init__
     method should generally not be overrode; all setup work should be done in
     activate() if possible. If you do override __init__, remember to super()!
-    
+
     Note that only one instance of each plugin will be instantiated for *all*
     connected clients. self.protocol will be changed by the plugin manager to
     the current protocol.
-    
+
     You may access the factory if necessary via self.factory.protocols
     to access other clients, but this "Is Not A Very Good Idea" (tm)
 
@@ -357,8 +357,7 @@ def activate(self):
             for alias in alias_list:
                 self.plugins['command_dispatcher'].register(alias, command)
 
-
     def deactivate(self):
         super(SimpleCommandPlugin, self).deactivate()
         for command in self.commands:
-            self.plugins['command_dispatcher'].unregister(command)
+            self.plugins['command_dispatcher'].unregister(command)

plugin_manager.py

@@ -119,7 +119,6 @@ def load_plugins(self, plugin_dir):
             self.logger.critical(str(e))
         self.activate_plugins()
 
-
     def reload_plugins(self):
         self.logger.warning("Reloading plugins.")
         for x in self.plugins:
@@ -132,7 +131,6 @@ def reload_plugins(self):
             self.logger.exception("Couldn't reload plugins!")
             raise
 
-
     def activate_plugins(self):
         for plugin in [self.plugins[x] for x in self.load_order]:
             if self.config.config['plugin_config'][plugin.name]['auto_activate']:
@@ -217,4 +215,4 @@ def print_this_defered_failure(f):
 
 
 class FatalPluginError(Exception):
-    pass
+    pass

plugins/core/admin_commands_plugin/admin_command_plugin.py

@@ -107,41 +107,21 @@ def promote(self, data):
     def make_guest(self, player):
         self.logger.trace("Setting %s to GUEST", player.name)
         player.access_level = UserLevels.GUEST
-        try:
-            self.player_manager.session.commit()
-        except:
-            self.player_manager.session.rollback()
-            raise
 
     @permissions(UserLevels.MODERATOR)
     def make_registered(self, player):
         self.logger.trace("Setting %s to REGISTERED", player.name)
         player.access_level = UserLevels.REGISTERED
-        try:
-            self.player_manager.session.commit()
-        except:
-            self.player_manager.session.rollback()
-            raise
 
     @permissions(UserLevels.ADMIN)
     def make_mod(self, player):
         player.access_level = UserLevels.MODERATOR
         self.logger.trace("Setting %s to MODERATOR", player.name)
-        try:
-            self.player_manager.session.commit()
-        except:
-            self.player_manager.session.rollback()
-            raise
 
     @permissions(UserLevels.OWNER)
     def make_admin(self, player):
         self.logger.trace("Setting %s to ADMIN", player.name)
         player.access_level = UserLevels.ADMIN
-        try:
-            self.player_manager.session.commit()
-        except:
-            self.player_manager.session.rollback()
-            raise
 
     @permissions(UserLevels.MODERATOR)
     def kick(self, data):
@@ -198,7 +178,7 @@ def unban(self, data):
         ip = data[0]
         for ban in self.player_manager.bans:
             if ban.ip == ip:
-                self.player_manager.session.delete(ban)
+                self.player_manager.remove_ban(ban)
                 self.protocol.send_chat_message("Unbanned IP: %s" % ip)
                 break
         else:

plugins/core/colored_names/colored_names.py

@@ -27,4 +27,3 @@ def on_chat_received(self, data):
             self.logger.warning("Received AttributeError in colored_name. %s", str(e))
             self.protocol.transport.write(data.original_data)
         return False
-

plugins/core/player_manager/manager.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import datetime
 from functools import wraps
 import inspect
@@ -6,13 +7,29 @@
 
 from enum import Enum
 from sqlalchemy.ext.mutable import Mutable
-from sqlalchemy.orm import Session, relationship, backref
+from sqlalchemy.orm import sessionmaker, relationship, backref
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey, Boolean, func
 from sqlalchemy.ext.declarative import declarative_base as sqla_declarative_base
 from twisted.words.ewords import AlreadyLoggedIn
 from sqlalchemy.types import TypeDecorator, VARCHAR
 from utility_functions import path
 
+
+@contextmanager
+def _autoclosing_session(sm):
+    session = sm()
+
+    try:
+        yield session
+
+    except:
+        session.rollback()
+        raise
+
+    finally:
+        session.close()
+
+
 class JSONEncodedDict(TypeDecorator):
     impl = VARCHAR
 
@@ -91,6 +108,32 @@ def __delitem__(self, key):
 MutableDict.associate_with(JSONEncodedDict)
 
 
+class RecordWithAttachedSession(object):
+    def __init__(self, record, sessionmaker):
+        self.__dict__['record'] = record
+        self.__dict__['sessionmaker'] = sessionmaker
+
+    def __getattr__(self, name):
+        with _autoclosing_session(self.sessionmaker) as session:
+            if sessionmaker.object_session(self.record) != session:
+                session.add(self.record)
+            session.refresh(self.record)
+            val = getattr(self.record, name)
+
+        return val
+
+    def __setattr__(self, name, val):
+        with _autoclosing_session(self.sessionmaker) as session:
+            if sessionmaker.object_session(self.record) != session:
+                session.add(self.record)
+            session.refresh(self.record)
+            setattr(self.record, name, val)
+            session.merge(self.record)
+            session.commit()
+
+        return val
+
+
 class Player(Base):
     __tablename__ = 'players'
 
@@ -121,7 +164,7 @@ def colored_name(self, colors):
 
     @property
     def storage(self):
-        caller = inspect.stack()[1][0].f_locals["self"].__class__.name
+        caller = inspect.stack()[2][0].f_locals["self"].__class__.name
         if self.plugin_storage is None:
             self.plugin_storage = {}
         try:
@@ -132,7 +175,7 @@ def storage(self):
 
     @storage.setter
     def storage(self, store):
-        caller = inspect.stack()[1][0].f_locals["self"].__class__.name
+        caller = inspect.stack()[2][0].f_locals["self"].__class__.name
         self.plugin_storage[caller] = store
 
     def as_dict(self):
@@ -160,75 +203,144 @@ class PlayerManager(object):
     def __init__(self, config):
         self.config = config
         self.engine = create_engine('sqlite:///%s' % path.preauthChild(self.config.player_db).path)
-        self.session = Session(self.engine)
         Base.metadata.create_all(self.engine)
-        for player in self.session.query(Player).all():
-            player.logged_in = False
-            player.protocol = None
+        self.sessionmaker = sessionmaker(bind=self.engine, autoflush=True)
+        with _autoclosing_session(self.sessionmaker) as session:
+            for player in session.query(Player).all():
+                player.logged_in = False
+                player.protocol = None
+                session.commit()
+
+    def _cache_and_return_from_session(self, session, record, collection=False):
+        to_return = record
+
+        if not isinstance(record, Base):
+            return to_return
+
+        if collection:
+            to_return = []
+            for r in record:
+                to_return.append(RecordWithAttachedSession(r, self.sessionmaker))
+        else:
+            to_return = RecordWithAttachedSession(record, self.sessionmaker)
+
+        return to_return
 
     def fetch_or_create(self, uuid, name, ip, protocol=None):
-        if self.session.query(Player).filter_by(uuid=uuid, logged_in=True).first():
-            raise AlreadyLoggedIn
-        if self.check_bans(ip):
-            raise Banned
-        while self.whois(name):
-            logger.info("Got a duplicate player, affixing _ to name")
-            name += "_"
-        player = self.session.query(Player).filter_by(uuid=uuid).first()
-        if player:
-            if player.name != name:
-                logger.info("Detected username change.")
-                player.name = name
-            if ip not in player.ips:
-                player.ips.append(IPAddress(ip=ip))
-                player.ip = ip
-            player.protocol = protocol
-        else:
-            logger.info("Adding new player with name: %s" % name)
-            player = Player(uuid=uuid, name=name,
-                            last_seen=datetime.datetime.utcnow(),
-                            access_level=int(UserLevels.GUEST),
-                            logged_in=False,
-                            protocol=protocol,
-                            client_id=-1,
-                            ip=ip,
-                            planet="",
-                            on_ship=True)
-            player.ips = [IPAddress(ip=ip)]
-            self.session.add(player)
-        if uuid == self.config.owner_uuid:
-            player.access_level = int(UserLevels.OWNER)
-        try:
-            self.session.commit()
-        except:
-            self.session.rollback()
-            raise
-        return player
+        with _autoclosing_session(self.sessionmaker) as session:
+            if session.query(Player).filter_by(uuid=uuid, logged_in=True).first():
+                raise AlreadyLoggedIn
+            if self.check_bans(ip):
+                raise Banned
+            while self.whois(name):
+                logger.info("Got a duplicate player, affixing _ to name")
+                name += "_"
+            player = session.query(Player).filter_by(uuid=uuid).first()
+            if player:
+                if player.name != name:
+                    logger.info("Detected username change.")
+                    player.name = name
+                if ip not in player.ips:
+                    player.ips.append(IPAddress(ip=ip))
+                    player.ip = ip
+                player.protocol = protocol
+            else:
+                logger.info("Adding new player with name: %s" % name)
+                player = Player(uuid=uuid, name=name,
+                                last_seen=datetime.datetime.now(),
+                                access_level=int(UserLevels.GUEST),
+                                logged_in=False,
+                                protocol=protocol,
+                                client_id=-1,
+                                ip=ip,
+                                planet="",
+                                on_ship=True)
+                player.ips = [IPAddress(ip=ip)]
+                session.add(player)
+            if uuid == self.config.owner_uuid:
+                player.access_level = int(UserLevels.OWNER)
+
+            session.commit()
+
+            return self._cache_and_return_from_session(session, player)
+
+    def delete(self, player_cache):
+        with _autoclosing_session(self.sessionmaker) as session:
+            session.delete(player_cache.record)
+            session.commit()
 
     def who(self):
-        return self.session.query(Player).filter_by(logged_in=True).all()
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).filter_by(logged_in=True).all(),
+                collection=True,
+                )
+
+    def all(self):
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).all(),
+                collection=True,
+                )
+
+    def all_like(self, regex):
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).filter(Player.name.like(regex)).all(),
+                collection=True,
+                )
 
     def whois(self, name):
-        return self.session.query(Player).filter(Player.logged_in == True,
-                                                 func.lower(Player.name) == func.lower(name)).first()
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).filter(
+                    Player.logged_in is True,
+                    func.lower(Player.name) == func.lower(name),
+                    ).first(),
+                )
 
     def check_bans(self, ip):
-        return self.session.query(Ban).filter_by(ip=ip).first() is not None
+        with _autoclosing_session(self.sessionmaker) as session:
+            return session.query(Ban).filter_by(ip=ip).first() is not None
 
     def ban(self, ip):
-        self.session.add(Ban(ip=ip))
-        try:
-            self.session.commit()
-        except:
-            self.session.rollback()
-            raise
+        with _autoclosing_session(self.sessionmaker) as session:
+            session.add(Ban(ip=ip))
+            session.commit()
+
+    @property
+    def bans(self):
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Ban).all(),
+                )
+
+    def delete_ban(self, ban_cache):
+        with _autoclosing_session(self.sessionmaker) as session:
+            session.delete(ban_cache.record)
+            session.commit()
 
     def get_by_name(self, name):
-        return self.session.query(Player).filter(func.lower(Player.name) == func.lower(name)).first()
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).filter(func.lower(Player.name) == func.lower(name)).first(),
+                )
 
     def get_logged_in_by_name(self, name):
-        return self.session.query(Player).filter(Player.logged_in == True,
-                                                 func.lower(Player.name) == func.lower(name)).first()
+        with _autoclosing_session(self.sessionmaker) as session:
+            return self._cache_and_return_from_session(
+                session,
+                session.query(Player).filter(
+                    Player.logged_in,
+                    func.lower(Player.name) == func.lower(name),
+                    ).first(),
+                )
 
 
 def permissions(level=UserLevels.OWNER):