diff --git a/socs/agents/suprsync/agent.py b/socs/agents/suprsync/agent.py index 9e0540808..964de7615 100644 --- a/socs/agents/suprsync/agent.py +++ b/socs/agents/suprsync/agent.py @@ -46,7 +46,7 @@ class SupRsync: Time (sec) after which a copy command will timeout """ - def __init__(self, agent, args): + def __init__(self, agent: ocs_agent.OCSAgent, args: argparse.Namespace) -> None: self.agent = agent self.instance_id = args.instance_id self.log = txaio.make_logger() @@ -65,6 +65,9 @@ def __init__(self, agent, args): self.compression = args.compression self.bwlimit = args.bwlimit self.suprsync_file_root = args.suprsync_file_root + self.db_echo: bool = args.db_echo + self.db_pool_size: int = args.db_pool_size + self.db_pool_max_overflow: int = args.db_pool_max_overflow # Feed for counting transfer errors, loop iterations. self.agent.register_feed('transfer_stats', @@ -112,7 +115,10 @@ def run(self, session, params=None): } """ - srfm = SupRsyncFilesManager(self.db_path, create_all=True) + srfm = SupRsyncFilesManager( + self.db_path, create_all=True, echo=self.db_echo, + pool_size=self.db_pool_size, max_overflow=self.db_pool_max_overflow + ) handler = SupRsyncFileHandler( srfm, self.archive_name, self.remote_basedir, ssh_host=self.ssh_host, @@ -251,6 +257,16 @@ def make_parser(parser=None): help="Bandwidth limit arg (passed through to rsync)") pgroup.add_argument('--suprsync-file-root', type=str, required=True, help="Local path where agent will write suprsync files") + pgroup.add_argument('--db-echo', action='store_true', help="Echos db queries") + pgroup.add_argument( + '--db-pool-size', type=int, default=5, + help="Number of connections to the suprsync db to keep open inside the " + "connection pool" + ) + pgroup.add_argument( + '--db-pool-max-overflow', type=int, default=10, + help="Number of connections to allow in the overflow pool." + ) return parser diff --git a/socs/db/suprsync.py b/socs/db/suprsync.py index e887528ee..3e40be162 100644 --- a/socs/db/suprsync.py +++ b/socs/db/suprsync.py @@ -201,12 +201,18 @@ class SupRsyncFilesManager: If true, writes sql statements to stdout """ - def __init__(self, db_path, create_all=True, echo=False): + def __init__( + self, db_path: str, create_all: bool =True, echo: bool=False, + pool_size: int=5, max_overflow: int=10 + ) -> None: db_path = os.path.abspath(db_path) if not os.path.exists(os.path.dirname(db_path)): os.makedirs(os.path.dirname(db_path)) - self._engine = create_engine(f'sqlite:///{db_path}', echo=echo) + self._engine = create_engine( + f'sqlite:///{db_path}', echo=echo, + pool_size=pool_size, max_overflow=max_overflow, + ) self.Session = sessionmaker(bind=self._engine) if create_all: