PYTHON-2392 Implicit sessions should always be discarded after connection errors (#498)

PYTHON-2075 Add more sessions tests with more read and write commands

Author: ShaneHarvey

pymongo/mongo_client.py

@@ -1827,17 +1827,20 @@ def _tmp_session(self, session, close=True):
             return
 
         s = self._ensure_session(session)
-        if s and close:
-            with s:
-                # Call end_session when we exit this scope.
-                yield s
-        elif s:
+        if s:
             try:
-                # Only call end_session on error.
                 yield s
-            except Exception:
+            except Exception as exc:
+                if isinstance(exc, ConnectionFailure):
+                    s._server_session.mark_dirty()
+
+                # Always call end_session on error.
                 s.end_session()
                 raise
+            finally:
+                # Call end_session when we exit this scope.
+                if close:
+                    s.end_session()
         else:
             yield None
 

test/sessions/dirty-session-errors.json

@@ -192,6 +192,149 @@
         }
       }
     },
+    {
+      "description": "Dirty explicit session is discarded (non-bulk write)",
+      "clientOptions": {
+        "retryWrites": true
+      },
+      "failPoint": {
+        "configureFailPoint": "failCommand",
+        "mode": {
+          "times": 1
+        },
+        "data": {
+          "failCommands": [
+            "findAndModify"
+          ],
+          "closeConnection": true
+        }
+      },
+      "operations": [
+        {
+          "name": "assertSessionNotDirty",
+          "object": "testRunner",
+          "arguments": {
+            "session": "session0"
+          }
+        },
+        {
+          "name": "findOneAndUpdate",
+          "object": "collection",
+          "arguments": {
+            "session": "session0",
+            "filter": {
+              "_id": 1
+            },
+            "update": {
+              "$inc": {
+                "x": 1
+              }
+            },
+            "returnDocument": "Before"
+          },
+          "result": {
+            "_id": 1
+          }
+        },
+        {
+          "name": "assertSessionDirty",
+          "object": "testRunner",
+          "arguments": {
+            "session": "session0"
+          }
+        },
+        {
+          "name": "endSession",
+          "object": "session0"
+        },
+        {
+          "name": "find",
+          "object": "collection",
+          "arguments": {
+            "filter": {
+              "_id": -1
+            }
+          },
+          "result": []
+        },
+        {
+          "name": "assertDifferentLsidOnLastTwoCommands",
+          "object": "testRunner"
+        }
+      ],
+      "expectations": [
+        {
+          "command_started_event": {
+            "command": {
+              "findAndModify": "test",
+              "query": {
+                "_id": 1
+              },
+              "update": {
+                "$inc": {
+                  "x": 1
+                }
+              },
+              "new": false,
+              "lsid": "session0",
+              "txnNumber": {
+                "$numberLong": "1"
+              },
+              "readConcern": null,
+              "writeConcern": null
+            },
+            "command_name": "findAndModify",
+            "database_name": "session-tests"
+          }
+        },
+        {
+          "command_started_event": {
+            "command": {
+              "findAndModify": "test",
+              "query": {
+                "_id": 1
+              },
+              "update": {
+                "$inc": {
+                  "x": 1
+                }
+              },
+              "new": false,
+              "lsid": "session0",
+              "txnNumber": {
+                "$numberLong": "1"
+              },
+              "readConcern": null,
+              "writeConcern": null
+            },
+            "command_name": "findAndModify",
+            "database_name": "session-tests"
+          }
+        },
+        {
+          "command_started_event": {
+            "command": {
+              "find": "test",
+              "filter": {
+                "_id": -1
+              }
+            },
+            "command_name": "find",
+            "database_name": "session-tests"
+          }
+        }
+      ],
+      "outcome": {
+        "collection": {
+          "data": [
+            {
+              "_id": 1,
+              "x": 1
+            }
+          ]
+        }
+      }
+    },
     {
       "description": "Dirty implicit session is discarded (write)",
       "clientOptions": {
@@ -300,6 +443,128 @@
         }
       }
     },
+    {
+      "description": "Dirty implicit session is discarded (non-bulk write)",
+      "clientOptions": {
+        "retryWrites": true
+      },
+      "failPoint": {
+        "configureFailPoint": "failCommand",
+        "mode": {
+          "times": 1
+        },
+        "data": {
+          "failCommands": [
+            "findAndModify"
+          ],
+          "closeConnection": true
+        }
+      },
+      "operations": [
+        {
+          "name": "findOneAndUpdate",
+          "object": "collection",
+          "arguments": {
+            "filter": {
+              "_id": 1
+            },
+            "update": {
+              "$inc": {
+                "x": 1
+              }
+            },
+            "returnDocument": "Before"
+          },
+          "result": {
+            "_id": 1
+          }
+        },
+        {
+          "name": "find",
+          "object": "collection",
+          "arguments": {
+            "filter": {
+              "_id": -1
+            }
+          },
+          "result": []
+        },
+        {
+          "name": "assertDifferentLsidOnLastTwoCommands",
+          "object": "testRunner"
+        }
+      ],
+      "expectations": [
+        {
+          "command_started_event": {
+            "command": {
+              "findAndModify": "test",
+              "query": {
+                "_id": 1
+              },
+              "update": {
+                "$inc": {
+                  "x": 1
+                }
+              },
+              "new": false,
+              "txnNumber": {
+                "$numberLong": "1"
+              },
+              "readConcern": null,
+              "writeConcern": null
+            },
+            "command_name": "findAndModify",
+            "database_name": "session-tests"
+          }
+        },
+        {
+          "command_started_event": {
+            "command": {
+              "findAndModify": "test",
+              "query": {
+                "_id": 1
+              },
+              "update": {
+                "$inc": {
+                  "x": 1
+                }
+              },
+              "new": false,
+              "txnNumber": {
+                "$numberLong": "1"
+              },
+              "readConcern": null,
+              "writeConcern": null
+            },
+            "command_name": "findAndModify",
+            "database_name": "session-tests"
+          }
+        },
+        {
+          "command_started_event": {
+            "command": {
+              "find": "test",
+              "filter": {
+                "_id": -1
+              }
+            },
+            "command_name": "find",
+            "database_name": "session-tests"
+          }
+        }
+      ],
+      "outcome": {
+        "collection": {
+          "data": [
+            {
+              "_id": 1,
+              "x": 1
+            }
+          ]
+        }
+      }
+    },
     {
       "description": "Dirty implicit session is discarded (read)",
       "failPoint": {
@@ -353,6 +618,54 @@
           ]
         }
       }
+    },
+    {
+      "description": "Dirty implicit session is discarded (non-cursor returning read)",
+      "failPoint": {
+        "configureFailPoint": "failCommand",
+        "mode": {
+          "times": 2
+        },
+        "data": {
+          "failCommands": [
+            "aggregate"
+          ],
+          "closeConnection": true
+        }
+      },
+      "operations": [
+        {
+          "name": "countDocuments",
+          "object": "collection",
+          "arguments": {
+            "filter": {}
+          },
+          "error": true
+        },
+        {
+          "name": "find",
+          "object": "collection",
+          "arguments": {
+            "filter": {
+              "_id": -1
+            }
+          },
+          "result": []
+        },
+        {
+          "name": "assertDifferentLsidOnLastTwoCommands",
+          "object": "testRunner"
+        }
+      ],
+      "outcome": {
+        "collection": {
+          "data": [
+            {
+              "_id": 1
+            }
+          ]
+        }
+      }
     }
   ]
-}
+}