|
29 | 29 | from itertools import groupby, count, chain |
30 | 30 | import json |
31 | 31 | import logging |
32 | | -from typing import Optional |
| 32 | +from typing import Optional, Union |
33 | 33 | from warnings import warn |
34 | 34 | from random import random |
35 | 35 | import re |
|
51 | 51 | from cassandra.connection import (ConnectionException, ConnectionShutdown, |
52 | 52 | ConnectionHeartbeat, ProtocolVersionUnsupported, |
53 | 53 | EndPoint, DefaultEndPoint, DefaultEndPointFactory, |
54 | | - SniEndPointFactory, ConnectionBusy) |
| 54 | + SniEndPointFactory, ConnectionBusy, locally_supported_compressions) |
55 | 55 | from cassandra.cqltypes import UserType |
56 | 56 | import cassandra.cqltypes as types |
57 | 57 | from cassandra.encoder import Encoder |
@@ -686,7 +686,7 @@ class Cluster(object): |
686 | 686 | Used for testing new protocol features incrementally before the new version is complete. |
687 | 687 | """ |
688 | 688 |
|
689 | | - compression = True |
| 689 | + compression: Union[bool, str] = True |
690 | 690 | """ |
691 | 691 | Controls compression for communications between the driver and Cassandra. |
692 | 692 | If left as the default of :const:`True`, either lz4 or snappy compression |
@@ -1173,7 +1173,7 @@ def token_metadata_enabled(self, enabled): |
1173 | 1173 | def __init__(self, |
1174 | 1174 | contact_points=_NOT_SET, |
1175 | 1175 | port=9042, |
1176 | | - compression=True, |
| 1176 | + compression: Union[bool, str] = True, |
1177 | 1177 | auth_provider=None, |
1178 | 1178 | load_balancing_policy=None, |
1179 | 1179 | reconnection_policy=None, |
@@ -1302,6 +1302,23 @@ def __init__(self, |
1302 | 1302 |
|
1303 | 1303 | self._resolve_hostnames() |
1304 | 1304 |
|
| 1305 | + if isinstance(compression, bool): |
| 1306 | + if compression and not locally_supported_compressions: |
| 1307 | + log.error( |
| 1308 | + "Compression is enabled, but no compression libraries are available. " |
| 1309 | + "Consider installing one of the Python packages: lz4 or python-snappy." |
| 1310 | + ) |
| 1311 | + elif isinstance(compression, str): |
| 1312 | + if not locally_supported_compressions.get(compression): |
| 1313 | + raise ValueError( |
| 1314 | + "Compression '%s' was requested, but it is not available. " |
| 1315 | + "Consider installing the corresponding Python package." % compression |
| 1316 | + ) |
| 1317 | + else: |
| 1318 | + raise TypeError( |
| 1319 | + "The 'compression' option must be either a string (e.g., 'lz4' or 'snappy') " |
| 1320 | + "or a boolean (True to enable any available compression, False to disable it)." |
| 1321 | + ) |
1305 | 1322 | self.compression = compression |
1306 | 1323 |
|
1307 | 1324 | if protocol_version is not _NOT_SET: |
|
0 commit comments