Python utils.get_pool函数代码示例

    def test_pool_with_fork(self):
        # Test that separate MongoClients have separate Pools, and that the
        # driver can create a new MongoClient after forking
        if sys.platform == "win32":
            raise SkipTest("Can't test forking on Windows")

        try:
            from multiprocessing import Process, Pipe
        except ImportError:
            raise SkipTest("No multiprocessing module")

        a = self.get_client(auto_start_request=False)
        a.pymongo_test.test.remove()
        a.pymongo_test.test.insert({'_id':1})
        a.pymongo_test.test.find_one()
        self.assertEqual(1, len(get_pool(a).sockets))
        a_sock = one(get_pool(a).sockets)

        def loop(pipe):
            c = self.get_client(auto_start_request=False)
            self.assertEqual(1,len(get_pool(c).sockets))
            c.pymongo_test.test.find_one()
            self.assertEqual(1,len(get_pool(c).sockets))
            pipe.send(one(get_pool(c).sockets).sock.getsockname())

        cp1, cc1 = Pipe()
        cp2, cc2 = Pipe()

        p1 = Process(target=loop, args=(cc1,))
        p2 = Process(target=loop, args=(cc2,))

        p1.start()
        p2.start()

        p1.join(1)
        p2.join(1)

        p1.terminate()
        p2.terminate()

        p1.join()
        p2.join()

        cc1.close()
        cc2.close()

        b_sock = cp1.recv()
        c_sock = cp2.recv()
        self.assertTrue(a_sock.sock.getsockname() != b_sock)
        self.assertTrue(a_sock.sock.getsockname() != c_sock)
        self.assertTrue(b_sock != c_sock)

        # a_sock, created by parent process, is still in the pool
        d_sock = get_pool(a).get_socket()
        self.assertEqual(a_sock, d_sock)
        d_sock.close()

 def assert_pool_size(self, pool_size):
     if pool_size == 0:
         self.assertTrue(
             self.c._MongoClient__member is None
             or not get_pool(self.c).sockets
         )
     else:
         self.assertEqual(
             pool_size, len(get_pool(self.c).sockets)
         )

    def test_socket_timeout_ms_validation(self):
        c = rs_or_single_client(socketTimeoutMS=10 * 1000)
        self.assertEqual(10, get_pool(c).opts.socket_timeout)

        c = connected(rs_or_single_client(socketTimeoutMS=None))
        self.assertEqual(None, get_pool(c).opts.socket_timeout)

        self.assertRaises(ValueError, rs_or_single_client, socketTimeoutMS=0)

        self.assertRaises(ValueError, rs_or_single_client, socketTimeoutMS=-1)

        self.assertRaises(ValueError, rs_or_single_client, socketTimeoutMS=1e10)

        self.assertRaises(ValueError, rs_or_single_client, socketTimeoutMS="foo")

    def test_exhaust_getmore_network_error(self):
        # When doing a getmore on an exhaust cursor, the socket stays checked
        # out on success but it's checked in on error to avoid semaphore leaks.
        client = rs_or_single_client(maxPoolSize=1)
        collection = client.pymongo_test.test
        collection.drop()
        collection.insert_many([{} for _ in range(200)])  # More than one batch.
        pool = get_pool(client)
        pool._check_interval_seconds = None  # Never check.

        cursor = collection.find(cursor_type=CursorType.EXHAUST)

        # Initial query succeeds.
        cursor.next()

        # Cause a network error.
        sock_info = cursor._Cursor__exhaust_mgr.sock
        sock_info.sock.close()

        # A getmore fails.
        self.assertRaises(ConnectionFailure, list, cursor)
        self.assertTrue(sock_info.closed)

        # The socket was closed and the semaphore was decremented.
        self.assertNotIn(sock_info, pool.sockets)
        self.assertTrue(pool._socket_semaphore.acquire(blocking=False))

    def test_nested_request(self):
        # auto_start_request is False
        client = get_client()
        pool = get_pool(client)
        self.assertFalse(client.in_request())

        # Start and end request
        client.start_request()
        self.assertInRequestAndSameSock(client, pool)
        client.end_request()
        self.assertNotInRequestAndDifferentSock(client, pool)

        # Double-nesting
        client.start_request()
        client.start_request()
        client.end_request()
        self.assertInRequestAndSameSock(client, pool)
        client.end_request()
        self.assertNotInRequestAndDifferentSock(client, pool)

        # Extra end_request calls have no effect - count stays at zero
        client.end_request()
        self.assertNotInRequestAndDifferentSock(client, pool)

        client.start_request()
        self.assertInRequestAndSameSock(client, pool)
        client.end_request()
        self.assertNotInRequestAndDifferentSock(client, pool)

    def test_auto_start_request(self):
        ctx = catch_warnings()
        try:
            warnings.simplefilter("ignore", DeprecationWarning)
            for bad_horrible_value in (None, 5, 'hi!'):
                self.assertRaises(
                    (TypeError, ConfigurationError),
                    lambda: get_client(auto_start_request=bad_horrible_value)
                )

            # auto_start_request should default to False
            client = get_client()
            self.assertFalse(client.auto_start_request)

            client = get_client(auto_start_request=True)
            self.assertTrue(client.auto_start_request)

            # Assure we acquire a request socket.
            client.pymongo_test.test.find_one()
            self.assertTrue(client.in_request())
            pool = get_pool(client)
            self.assertRequestSocket(pool)
            self.assertSameSock(pool)

            client.end_request()
            self.assertNoRequest(pool)
            self.assertDifferentSock(pool)

            # Trigger auto_start_request
            client.pymongo_test.test.find_one()
            self.assertRequestSocket(pool)
            self.assertSameSock(pool)
        finally:
            ctx.exit()

    def test_auth_network_error(self):
        # Make sure there's no semaphore leak if we get a network error
        # when authenticating a new socket with cached credentials.
        auth_client = self._get_client()
        if not server_started_with_auth(auth_client):
            raise SkipTest('Authentication is not enabled on server')

        auth_client.admin.add_user('admin', 'password')
        auth_client.admin.authenticate('admin', 'password')
        try:
            # Get a client with one socket so we detect if it's leaked.
            c = self._get_client(max_pool_size=1, waitQueueTimeoutMS=1)

            # Simulate an authenticate() call on a different socket.
            credentials = auth._build_credentials_tuple(
                'MONGODB-CR', 'admin',
                unicode('admin'), unicode('password'),
                {})

            c._cache_credentials('test', credentials, connect=False)

            # Cause a network error on the actual socket.
            pool = get_pool(c)
            socket_info = one(pool.sockets)
            socket_info.sock.close()

            # In __check_auth, the client authenticates its socket with the
            # new credential, but gets a socket.error. Should be reraised as
            # AutoReconnect.
            self.assertRaises(AutoReconnect, c.test.collection.find_one)

            # No semaphore leak, the pool is allowed to make a new socket.
            c.test.collection.find_one()
        finally:
            remove_all_users(auth_client.admin)

    def test_connection(self):
        c = Connection(host, port)
        self.assertTrue(c.auto_start_request)
        self.assertEqual(None, c.max_pool_size)
        self.assertFalse(c.slave_okay)
        self.assertFalse(c.safe)
        self.assertEqual({}, c.get_lasterror_options())

        # Connection's writes are unacknowledged by default
        doc = {"_id": ObjectId()}
        coll = c.pymongo_test.write_concern_test
        coll.drop()
        coll.insert(doc)
        coll.insert(doc)

        c = Connection("mongodb://%s:%s/?safe=true" % (host, port))
        self.assertTrue(c.safe)

        # To preserve legacy Connection's behavior, max_size should be None.
        # Pool should handle this without error.
        self.assertEqual(None, get_pool(c).max_size)
        c.end_request()

        # Connection's network_timeout argument is translated into
        # socketTimeoutMS
        self.assertEqual(123, Connection(
            host, port, network_timeout=123)._MongoClient__net_timeout)

        for network_timeout in 'foo', 0, -1:
            self.assertRaises(
                ConfigurationError,
                Connection, host, port, network_timeout=network_timeout)

    def test_auth_network_error(self):
        # Make sure there's no semaphore leak if we get a network error
        # when authenticating a new socket with cached credentials.

        # Get a client with one socket so we detect if it's leaked.
        c = connected(rs_or_single_client(maxPoolSize=1,
                                          waitQueueTimeoutMS=1))

        # Simulate an authenticate() call on a different socket.
        credentials = auth._build_credentials_tuple(
            'DEFAULT', 'admin', db_user, db_pwd, {})

        c._cache_credentials('test', credentials, connect=False)

        # Cause a network error on the actual socket.
        pool = get_pool(c)
        socket_info = one(pool.sockets)
        socket_info.sock.close()

        # SocketInfo.check_auth logs in with the new credential, but gets a
        # socket.error. Should be reraised as AutoReconnect.
        self.assertRaises(AutoReconnect, c.test.collection.find_one)

        # No semaphore leak, the pool is allowed to make a new socket.
        c.test.collection.find_one()

    def test_server_disconnect(self):
        # PYTHON-345, we need to make sure that threads' request sockets are
        # closed by disconnect().
        #
        # 1. Create a client with auto_start_request=True
        # 2. Start N threads and do a find() in each to get a request socket
        # 3. Pause all threads
        # 4. In the main thread close all sockets, including threads' request
        #       sockets
        # 5. In main thread, do a find(), which raises AutoReconnect and resets
        #       pool
        # 6. Resume all threads, do a find() in them
        #
        # If we've fixed PYTHON-345, then only one AutoReconnect is raised,
        # and all the threads get new request sockets.
        cx = get_client(auto_start_request=True)
        collection = cx.db.pymongo_test

        # acquire a request socket for the main thread
        collection.find_one()
        pool = get_pool(collection.database.connection)
        socket_info = pool._get_request_state()
        assert isinstance(socket_info, SocketInfo)
        request_sock = socket_info.sock

        state = FindPauseFind.create_shared_state(nthreads=40)

        threads = [FindPauseFind(collection, state) for _ in range(state.nthreads)]

        # Each thread does a find(), thus acquiring a request socket
        for t in threads:
            t.start()

        # Wait for the threads to reach the rendezvous
        FindPauseFind.wait_for_rendezvous(state)

        try:
            # Simulate an event that closes all sockets, e.g. primary stepdown
            for t in threads:
                t.request_sock.close()

            # Finally, ensure the main thread's socket's last_checkout is
            # updated:
            collection.find_one()

            # ... and close it:
            request_sock.close()

            # Doing an operation on the client raises an AutoReconnect and
            # resets the pool behind the scenes
            self.assertRaises(AutoReconnect, collection.find_one)

        finally:
            # Let threads do a second find()
            FindPauseFind.resume_after_rendezvous(state)

        joinall(threads)

        for t in threads:
            self.assertTrue(t.passed, "%s threw exception" % t)

    def run_mongo_thread(self):
        pool = get_pool(self.client)
        assert len(pool.sockets) == 1, "Expected 1 socket, found %d" % (
            len(pool.sockets)
        )

        sock_info = one(pool.sockets)

        self.client.start_request()

        # start_request() hasn't yet moved the socket from the general pool into
        # the request
        assert len(pool.sockets) == 1
        assert one(pool.sockets) == sock_info

        self.client[DB].test.find_one()

        # find_one() causes the socket to be used in the request, so now it's
        # bound to this thread
        assert len(pool.sockets) == 0
        assert pool._get_request_state() == sock_info
        self.client.end_request()

        # The socket is back in the pool
        assert len(pool.sockets) == 1
        assert one(pool.sockets) == sock_info

    def test_auto_start_request(self):
        for bad_horrible_value in (None, 5, 'hi!'):
            self.assertRaises(
                (TypeError, ConfigurationError),
                lambda: get_client(auto_start_request=bad_horrible_value)
            )

        # auto_start_request should default to False
        client = get_client()
        self.assertFalse(client.auto_start_request)

        client = get_client(auto_start_request=True)
        self.assertTrue(client.auto_start_request)

        # Assure we acquire a request socket.
        client.pymongo_test.test.find_one()
        self.assertTrue(client.in_request())
        pool = get_pool(client)
        self.assertRequestSocket(pool)
        self.assertSameSock(pool)

        client.end_request()
        self.assertNoRequest(pool)
        self.assertDifferentSock(pool)

        # Trigger auto_start_request
        client.pymongo_test.test.find_one()
        self.assertRequestSocket(pool)
        self.assertSameSock(pool)

    def test_max_pool_size_none(self):
        c = rs_or_single_client(maxPoolSize=None)
        collection = c[DB].test

        # Need one document.
        collection.drop()
        collection.insert_one({})

        cx_pool = get_pool(c)
        nthreads = 10
        threads = []
        lock = threading.Lock()
        self.n_passed = 0

        def f():
            for _ in range(5):
                collection.find_one({'$where': delay(0.1)})

            with lock:
                self.n_passed += 1

        for i in range(nthreads):
            t = threading.Thread(target=f)
            threads.append(t)
            t.start()

        joinall(threads)
        self.assertEqual(nthreads, self.n_passed)
        self.assertTrue(len(cx_pool.sockets) > 1)

    def test_max_pool_size(self):
        max_pool_size = 4
        c = rs_or_single_client(maxPoolSize=max_pool_size)
        collection = c[DB].test

        # Need one document.
        collection.drop()
        collection.insert_one({})

        # nthreads had better be much larger than max_pool_size to ensure that
        # max_pool_size sockets are actually required at some point in this
        # test's execution.
        cx_pool = get_pool(c)
        nthreads = 10
        threads = []
        lock = threading.Lock()
        self.n_passed = 0

        def f():
            for _ in range(5):
                collection.find_one({'$where': delay(0.1)})
                assert len(cx_pool.sockets) <= max_pool_size

            with lock:
                self.n_passed += 1

        for i in range(nthreads):
            t = threading.Thread(target=f)
            threads.append(t)
            t.start()

        joinall(threads)
        self.assertEqual(nthreads, self.n_passed)
        self.assertTrue(len(cx_pool.sockets) > 1)
        self.assertEqual(max_pool_size, cx_pool._socket_semaphore.counter)

    def test_contextlib(self):
        client = rs_or_single_client()
        client.pymongo_test.drop_collection("test")
        client.pymongo_test.test.insert_one({"foo": "bar"})

        # The socket used for the previous commands has been returned to the
        # pool
        self.assertEqual(1, len(get_pool(client).sockets))

        with contextlib.closing(client):
            self.assertEqual("bar", client.pymongo_test.test.find_one()["foo"])
            self.assertEqual(1, len(get_pool(client).sockets))
        self.assertEqual(0, len(get_pool(client).sockets))

        with client as client:
            self.assertEqual("bar", client.pymongo_test.test.find_one()["foo"])
        self.assertEqual(0, len(get_pool(client).sockets))