fix: set the consistency level in connection (#179)

hipagesgroup · Jun 13, 2022 · 9132fe3 · 9132fe3
1 parent 74f074c
commit 9132fe3
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/.gitignore b/.gitignore
@@ -21,6 +21,7 @@ pip-delete-this-directory.txt
 
 # Env
 venv
+.venv/*
 .env
 .envrc
 env/

diff --git a/README.md b/README.md
@@ -75,6 +75,7 @@ from cassandra.policies import DCAwareRoundRobinPolicy
 from cassandra.cqlengine import columns
 from cassandra.cqlengine.management import sync_table
 from cassandra.cqlengine.models import Model
+from cassandra import ConsistencyLevel
 
 load_balancing_policy = DCAwareRoundRobinPolicy(local_dc='AWS_VPC_AP_SOUTHEAST_2')
 
@@ -83,7 +84,8 @@ conn = CassandraConnectionManager(
         cluster_ips=["1.1.1.1", "2.2.2.2"],
         port=9042,
         load_balancing_policy=load_balancing_policy,
-    )
+    ),
+    consistency_level=ConsistencyLevel.LOCAL_QUORUM
 )
 
 conn = CassandraConnectionManager(
@@ -93,7 +95,8 @@ conn = CassandraConnectionManager(
         load_balancing_policy=load_balancing_policy,
         secrets_manager=CassandraSecretsManager(
         username_var="MY_CUSTOM_USERNAME_ENV_VAR"),
-    )
+    ),
+    consistency_level=ConsistencyLevel.LOCAL_ONE
 )
 
 # For running Cassandra model operations

diff --git a/hip_data_tools/apache/cassandra.py b/hip_data_tools/apache/cassandra.py
@@ -204,15 +204,15 @@ class CassandraConnectionManager:
         for connecting to a cluster
     """
 
-    def __init__(self, settings: CassandraConnectionSettings):
+    def __init__(self, settings: CassandraConnectionSettings, consistency_level: ConsistencyLevel = ConsistencyLevel.LOCAL_ONE):
         self._settings = settings
         self.cluster = None
         self.session = None
         self._auth = PlainTextAuthProvider(
             username=self._settings.secrets_manager.username,
             password=self._settings.secrets_manager.password,
         )
-        self.consistency_level = ConsistencyLevel.LOCAL_QUORUM
+        self.consistency_level = consistency_level
 
     def get_cluster(self) -> Cluster:
         """
@@ -236,7 +236,7 @@ def get_session(self, keyspace) -> Session:
         """
         if self.session is None:
             self.session = self.get_cluster().connect(keyspace)
-
+            self.session.default_consistency_level = self.consistency_level
         return self.session
 
     def setup_connection(self, default_keyspace) -> None:

diff --git a/hip_data_tools/etl/s3_to_cassandra.py b/hip_data_tools/etl/s3_to_cassandra.py
@@ -5,6 +5,7 @@
 
 from attr import dataclass
 from cassandra.datastax.graph import Result
+from cassandra import ConsistencyLevel
 
 from hip_data_tools.apache.cassandra import CassandraUtil, CassandraConnectionManager, \
     CassandraConnectionSettings
@@ -42,7 +43,7 @@ def _get_cassandra_util(self):
         return CassandraUtil(
             keyspace=self.__settings.destination_keyspace,
             conn=CassandraConnectionManager(
-                settings=self.__settings.destination_connection_settings),
+                settings=self.__settings.destination_connection_settings,consistency_level=ConsistencyLevel.LOCAL_QUORUM),
         )
 
     def _get_s3_util(self):
-Original file line number
+Diff line change
@@ Expand Up / @@ -21,6 +21,7 @@ pip-delete-this-directory.txt @@
     # Env
     venv
+    .venv/*
     .env
     .envrc
     env/
@@ Expand Down @@