def __init__()

in evalbench/databases/alloydb.py [0:0]

• 28 lines of code
• 2 McCabe index (conditional complexity)

    def __init__(self, db_config):
        """
        Initializes the AlloyDB connection, overriding the PGDB's
        default Google Cloud SQL connection mechanism.
        """
        super().__init__(db_config)
        self.nl_config = db_config['nl_config']

        self.connector = AlloyDBConnector()

        def get_conn_alloydb():
            return self.connector.connect(
                self.db_path,
                "pg8000",
                user=self.username,
                password=self.password,
                db=self.db_name,
                enable_iam_auth=False,
                ip_type=AlloyDBIPTypes.PUBLIC,
            )

        def get_engine_args_alloydb():
            common_args = {
                "creator": get_conn_alloydb,
                "connect_args": {"command_timeout": 60},
            }
            if "is_tmp_db" in db_config:
                common_args["poolclass"] = NullPool
            else:
                common_args["pool_size"] = 50
                common_args["pool_recycle"] = 300
            return common_args

        self.engine = sqlalchemy.create_engine(
            "postgresql+pg8000://", **get_engine_args_alloydb()
        )