chore: Support consistent reads in dynamodb (feast-dev#3119)

achals · web-flow · commit eaaa872fb323 · 2022-08-22T11:08:57.000-07:00
* chore: Support consistent reads in dynamodb

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* fix

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/infra/materialization/lambda/lambda_engine.py b/sdk/python/feast/infra/materialization/lambda/lambda_engine.py
@@ -227,7 +227,7 @@ def _materialize_one(
 
             logger.info(
                 f"Ingested task; request id {response['ResponseMetadata']['RequestId']}, "
-                f"rows written: {output['written_rows']}"
+                f"Output: {output}"
             )
 
         for f in not_done:
diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py
@@ -16,7 +16,7 @@
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
-from pydantic import StrictStr
+from pydantic import StrictBool, StrictStr
 from pydantic.typing import Literal, Union
 
 from feast import Entity, FeatureView, utils
@@ -63,6 +63,9 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
     table_name_template: StrictStr = "{project}.{table_name}"
     """DynamoDB table name template"""
 
+    consistent_reads: StrictBool = False
+    """Whether to read from Dynamodb by forcing consistent reads"""
+
 
 class DynamoDBOnlineStore(OnlineStore):
     """
@@ -237,12 +240,12 @@ def online_read(
             batch_entity_ids = {
                 table_instance.name: {
                     "Keys": [{"entity_id": entity_id} for entity_id in batch],
-                    "ConsistentRead": True,
+                    "ConsistentRead": online_config.consistent_reads,
                 }
             }
             with tracing_span(name="remote_call"):
                 response = dynamodb_resource.batch_get_item(
-                    RequestItems=batch_entity_ids
+                    RequestItems=batch_entity_ids,
                 )
             response = response.get("Responses")
             table_responses = response.get(table_instance.name)
diff --git a/sdk/python/tests/integration/materialization/test_lambda.py b/sdk/python/tests/integration/materialization/test_lambda.py
@@ -24,11 +24,15 @@
 def test_lambda_materialization_consistency():
     lambda_config = IntegrationTestRepoConfig(
         provider="aws",
-        online_store={"type": "dynamodb", "region": "us-west-2"},
+        online_store={
+            "type": "dynamodb",
+            "region": "us-west-2",
+            "consistent_reads": True,
+        },
         offline_store_creator=RedshiftDataSourceCreator,
         batch_engine={
             "type": "lambda",
-            "materialization_image": "402087665549.dkr.ecr.us-west-2.amazonaws.com/feast-lambda-consumer:v1",
+            "materialization_image": "402087665549.dkr.ecr.us-west-2.amazonaws.com/feast-lambda-consumer:v2",
             "lambda_role": "arn:aws:iam::402087665549:role/lambda_execution_role",
         },
         registry_location=RegistryLocation.S3,

Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,7 @@ def _materialize_one(`
`227`	`227`
`228`	`228`	`logger.info(`
`229`	`229`	`f"Ingested task; request id {response['ResponseMetadata']['RequestId']}, "`
`230`		`- f"rows written: {output['written_rows']}"`
	`230`	`+ f"Output: {output}"`
`231`	`231`	`)`
`232`	`232`
`233`	`233`	`for f in not_done:`