PYTHON-1723 Support zstd wire compression

Author: behackett

.evergreen/config.yml

@@ -1026,6 +1026,10 @@ axes:
  display_name: zlib compression
  variables:
  COMPRESSORS: "zlib"
+ - id: zstd
+ display_name: zstd compression
+ variables:
+ COMPRESSORS: "zstd"
  - id: python-version
  display_name: "Python"
  values:
@@ -1283,15 +1287,22 @@ buildvariants:
  compression: "*"
  - python-version: ["jython2.7"]
  c-extensions: "*"
- compression: "snappy"
+ compression: ["snappy", "zstd"]
  batchtime: 10080 # 7 days
  display_name: "${compression} ${c-extensions} ${python-version} (x86_64)"
  # Ubuntu 16.04 images have libsnappy-dev installed
  run_on: ubuntu1604-test
  tasks:
- - "test-latest-standalone"
- - "test-4.0-standalone"
- - "test-3.6-standalone"
+ - "test-latest-standalone"
+ rules:
+ - if:
+ python-version: "*"
+ c-extensions: "*"
+ compression: ["snappy", "zlib"]
+ then:
+ add_tasks:
+ - "test-4.0-standalone"
+ - "test-3.6-standalone"
 
 - matrix_name: "tests-python-version-py37-plus-compression"
  matrix_spec: {"python-version-requires-openssl-102-plus": "*", "c-extensions": "*", "compression": "*"}
@@ -1301,9 +1312,16 @@ buildvariants:
  # Ubuntu 16.04 images have libsnappy-dev installed, and provides OpenSSL 1.0.2
  run_on: ubuntu1604-test
  tasks:
- - "test-latest-standalone"
- - "test-4.0-standalone"
- - "test-3.6-standalone"
+ - "test-latest-standalone"
+ rules:
+ - if:
+ python-version: "*"
+ c-extensions: "*"
+ compression: ["snappy", "zlib"]
+ then:
+ add_tasks:
+ - "test-4.0-standalone"
+ - "test-3.6-standalone"
 
 - matrix_name: "tests-python-version-green-framework-rhel62"
  matrix_spec: {"python-version": "*", "green-framework": "*", auth-ssl: "*"}

.evergreen/run-tests.sh

@@ -56,6 +56,12 @@ elif [ $COMPRESSORS = "snappy" ]; then
  # 0.5.2 has issues in pypy3(.5)
  pip install python-snappy==0.5.1
  PYTHON=python
+elif [ $COMPRESSORS = "zstd" ]; then
+ $PYTHON_BINARY -m virtualenv --system-site-packages --never-download zstdtest
+ . zstdtest/bin/activate
+ trap "deactivate; rm -rf zstdtest" EXIT HUP
+ pip install zstandard
+ PYTHON=python
 else
  PYTHON="$PYTHON_BINARY"
 fi

README.rst

@@ -118,10 +118,15 @@ Wire protocol compression with snappy requires `python-snappy
 
  $ python -m pip install pymongo[snappy]
 
+Wire protocol compression with zstandard requires `zstandard
+<https://pypi.org/project/zstandard>`_::
+
+ $ python -m pip install pymongo[zstd]
+
 You can install all dependencies automatically with the following
 command::
 
- $ python -m pip install pymongo[snappy,gssapi,srv,tls]
+ $ python -m pip install pymongo[snappy,gssapi,srv,tls,zstd]
 
 Other optional packages:
 

doc/changelog.rst

@@ -36,6 +36,7 @@ Version 3.9 adds support for MongoDB 4.2. Highlights include:
 - Support for retryable reads and the ``retryReads`` URI option which is
  enabled by default. See the :class:`~pymongo.mongo_client.MongoClient`
  documentation for details.
+- Support zstandard for wire protocol compression.
 
  Now that supported operations are retried automatically and transparently,
  users should consider adjusting any custom retry logic to prevent

doc/installation.rst

@@ -75,10 +75,15 @@ Wire protocol compression with snappy requires `python-snappy
 
  $ python -m pip install pymongo[snappy]
 
+Wire protocol compression with zstandard requires `zstandard
+<https://pypi.org/project/zstandard>`_::
+
+ $ python -m pip install pymongo[zstd]
+
 You can install all dependencies automatically with the following
 command::
 
- $ python -m pip install pymongo[snappy,gssapi,srv,tls]
+ $ python -m pip install pymongo[snappy,gssapi,srv,tls,zstd]
 
 Other optional packages:
 

pymongo/compression_support.py

@@ -28,9 +28,15 @@
  # Python built without zlib support.
  _HAVE_ZLIB = False
 
+try:
+ from zstandard import ZstdCompressor, ZstdDecompressor
+ _HAVE_ZSTD = True
+except ImportError:
+ _HAVE_ZSTD = False
+
 from pymongo.monitoring import _SENSITIVE_COMMANDS
 
-_SUPPORTED_COMPRESSORS = set(["snappy", "zlib"])
+_SUPPORTED_COMPRESSORS = set(["snappy", "zlib", "zstd"])
 _NO_COMPRESSION = set(['ismaster'])
 _NO_COMPRESSION.update(_SENSITIVE_COMMANDS)
 
@@ -57,6 +63,11 @@ def validate_compressors(dummy, value):
  warnings.warn(
  "Wire protocol compression with zlib is not available. "
  "The zlib module is not available.")
+ elif compressor == "zstd" and not _HAVE_ZSTD:
+ compressors.remove(compressor)
+ warnings.warn(
+ "Wire protocol compression with zstandard is not available. "
+ "You must install the zstandard module for zstandard support.")
  return compressors
 
 
@@ -83,6 +94,8 @@ def get_compression_context(self, compressors):
  return SnappyContext()
  elif chosen == "zlib":
  return ZlibContext(self.zlib_compression_level)
+ elif chosen == "zstd":
+ return ZstdContext()
 
 
 def _zlib_no_compress(data):
@@ -113,6 +126,16 @@ def __init__(self, level):
  self.compress = lambda data: zlib.compress(data, level)
 
 
+class ZstdContext(object):
+ compressor_id = 3
+
+ @staticmethod
+ def compress(data):
+ # ZstdCompressor is not thread safe.
+ # TODO: Use a pool?
+ return ZstdCompressor().compress(data)
+
+
 def decompress(data, compressor_id):
  if compressor_id == SnappyContext.compressor_id:
  # python-snappy doesn't support the buffer interface.
@@ -126,5 +149,9 @@ def decompress(data, compressor_id):
  return snappy.uncompress(bytes(data))
  elif compressor_id == ZlibContext.compressor_id:
  return zlib.decompress(data)
+ elif compressor_id == ZstdContext.compressor_id:
+ # ZstdDecompressor is not thread safe.
+ # TODO: Use a pool?
+ return ZstdDecompressor().decompress(data)
  else:
  raise ValueError("Unknown compressorId %d" % (compressor_id,))

pymongo/mongo_client.py

@@ -312,13 +312,15 @@ def __init__(
  https://docs.mongodb.com/manual/faq/diagnostics/#does-tcp-keepalive-time-affect-mongodb-deployments",
  - `compressors`: Comma separated list of compressors for wire
  protocol compression. The list is used to negotiate a compressor
- with the server. Currently supported options are "snappy" and
- "zlib". Support for snappy requires the
+ with the server. Currently supported options are "snappy", "zlib"
+ and "zstd". Support for snappy requires the
  `python-snappy <https://pypi.org/project/python-snappy/>`_ package.
- zlib support requires the Python standard library zlib module.
- By default no compression is used. Compression support must also be
- enabled on the server. MongoDB 3.4+ supports snappy compression.
- MongoDB 3.6+ supports snappy and zlib.
+ zlib support requires the Python standard library zlib module. zstd
+ requires the `zstandard <https://pypi.org/project/zstandard/>`_
+ package. By default no compression is used. Compression support
+ must also be enabled on the server. MongoDB 3.4+ supports snappy
+ compression. MongoDB 3.6 adds support for zlib. MongoDB 4.2 adds
+ support for zstd.
  - `zlibCompressionLevel`: (int) The zlib compression level to use
  when zlib is used as the wire protocol compressor. Supported values
  are -1 through 9. -1 tells the zlib library to use its default

setup.py

@@ -317,7 +317,7 @@ def build_extension(self, ext):
  sources=['pymongo/_cmessagemodule.c',
  'bson/buffer.c'])]
 
-extras_require = {'snappy': ["python-snappy"]}
+extras_require = {'snappy': ["python-snappy"], 'zstd': ["zstandard"]}
 vi = sys.version_info
 if vi[0] == 2:
  extras_require.update(

test/test_client.py

@@ -36,7 +36,7 @@
 from pymongo import auth, message
 from pymongo.common import _UUID_REPRESENTATIONS
 from pymongo.command_cursor import CommandCursor
-from pymongo.compression_support import _HAVE_SNAPPY
+from pymongo.compression_support import _HAVE_SNAPPY, _HAVE_ZSTD
 from pymongo.cursor import Cursor, CursorType
 from pymongo.database import Database
 from pymongo.errors import (AutoReconnect,
@@ -1343,6 +1343,21 @@ def compression_settings(client):
  opts = compression_settings(client)
  self.assertEqual(opts.compressors, ['snappy', 'zlib'])
 
+ if not _HAVE_ZSTD:
+ uri = "mongodb://localhost:27017/?compressors=zstd"
+ client = MongoClient(uri, connect=False)
+ opts = compression_settings(client)
+ self.assertEqual(opts.compressors, [])
+ else:
+ uri = "mongodb://localhost:27017/?compressors=zstd"
+ client = MongoClient(uri, connect=False)
+ opts = compression_settings(client)
+ self.assertEqual(opts.compressors, ['zstd'])
+ uri = "mongodb://localhost:27017/?compressors=zstd,zlib"
+ client = MongoClient(uri, connect=False)
+ opts = compression_settings(client)
+ self.assertEqual(opts.compressors, ['zstd', 'zlib'])
+
  options = client_context.default_client_options
  if "compressors" in options and "zlib" in options["compressors"]:
  for level in range(-1, 10):