Address the rest of the comments

Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai>
feast-dev · feast-ci-bot · Jul 3, 2021 · Apr 27, 2021 · Apr 27, 2021 · May 4, 2021
commit 57a607cd5c52a3ddc1e799eb24d508c921d361df
@@ -48,12 +48,6 @@ class DynamoDBOnlineStore(OnlineStore):
     Online feature store for AWS DynamoDB.
     """
 
-    def _initialize_dynamodb(self, online_config: DynamoDBOnlineStoreConfig):
-        return (
-            boto3.client("dynamodb", region_name=online_config.region),
-            boto3.resource("dynamodb", region_name=online_config.region),
-        )
-
     def update(
         self,
         config: RepoConfig,
@@ -65,11 +59,11 @@ def update(
     ):
         online_config = config.online_store
         assert isinstance(online_config, DynamoDBOnlineStoreConfig)
-        client, dynamodb = self._initialize_dynamodb(online_config)
+        dynamodb_client, dynamodb_resource = self._initialize_dynamodb(online_config)
 
         for table_instance in tables_to_keep:
             try:
-                table = dynamodb.create_table(
+                dynamodb_resource.create_table(
                     TableName=f"{config.project}.{table_instance.name}",
                     KeySchema=[{"AttributeName": "entity_id", "KeyType": "HASH"}],
                     AttributeDefinitions=[
@@ -78,18 +72,18 @@ def update(
                     BillingMode="PAY_PER_REQUEST",
                 )
             except ClientError as ce:
-                print(ce)
-                if ce.response["Error"]["Code"] == "ResourceNotFoundException":
-                    table = dynamodb.Table(f"{config.project}.{table_instance.name}")
+                # If the table creation fails with ResourceInUseException,
+                # it means the table already exists or is being created.
+                # Otherwise, re-raise the exception
+                if ce.response["Error"]["Code"] != "ResourceInUseException":
+                    raise
 
         for table_instance in tables_to_keep:
-            client.get_waiter("table_exists").wait(
+            dynamodb_client.get_waiter("table_exists").wait(
                 TableName=f"{config.project}.{table_instance.name}"
             )
 
-        for table_instance in tables_to_delete:
-            table = dynamodb.Table(f"{config.project}.{table_instance.name}")
-            table.delete()
+        self._delete_tables_idempotent(dynamodb_resource, config, tables_to_delete)
 
     def teardown(
         self,
@@ -99,11 +93,9 @@ def teardown(
     ):
         online_config = config.online_store
         assert isinstance(online_config, DynamoDBOnlineStoreConfig)
-        _, dynamodb = self._initialize_dynamodb(online_config)
+        _, dynamodb_resource = self._initialize_dynamodb(online_config)
 
-        for table_instance in tables:
-            table = dynamodb.Table(f"{config.project}.{table_instance.name}")
-            table.delete()
+        self._delete_tables_idempotent(dynamodb_resource, config, tables)
 
     def online_write_batch(
         self,
@@ -116,9 +108,9 @@ def online_write_batch(
     ) -> None:
         online_config = config.online_store
         assert isinstance(online_config, DynamoDBOnlineStoreConfig)
-        _, dynamodb = self._initialize_dynamodb(online_config)
+        _, dynamodb_resource = self._initialize_dynamodb(online_config)
 
-        table_instance = dynamodb.Table(f"{config.project}.{table.name}")
+        table_instance = dynamodb_resource.Table(f"{config.project}.{table.name}")
         with table_instance.batch_writer() as batch:
             for entity_key, features, timestamp, created_ts in data:
                 entity_id = compute_entity_id(entity_key)
@@ -144,12 +136,12 @@ def online_read(
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
         online_config = config.online_store
         assert isinstance(online_config, DynamoDBOnlineStoreConfig)
-        _, dynamodb = self._initialize_dynamodb(online_config)
+        _, dynamodb_resource = self._initialize_dynamodb(online_config)
 
         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
         for entity_key in entity_keys:
-            table_instance = dynamodb.Table(f"{config.project}.{table.name}")
-            entity_id = compute_entity_id(entity_key)  # TODO check id
+            table_instance = dynamodb_resource.Table(f"{config.project}.{table.name}")
+            entity_id = compute_entity_id(entity_key)
             response = table_instance.get_item(Key={"entity_id": entity_id})
             value = response.get("Item")
 
@@ -163,3 +155,28 @@ def online_read(
             else:
                 result.append((None, None))
         return result
+
+    def _initialize_dynamodb(self, online_config: DynamoDBOnlineStoreConfig):
+        return (
+            boto3.client("dynamodb", region_name=online_config.region),
+            boto3.resource("dynamodb", region_name=online_config.region),
+        )
+
+    def _delete_tables_idempotent(
+        self,
+        dynamodb_resource,
+        config: RepoConfig,
+        tables: Sequence[Union[FeatureTable, FeatureView]],
+    ):
+        for table_instance in tables:
+            try:
+                table = dynamodb_resource.Table(
+                    f"{config.project}.{table_instance.name}"
+                )
+                table.delete()
+            except ClientError as ce:
+                # If the table deletion fails with ResourceNotFoundException,
+                # it means the table has already been deleted.
+                # Otherwise, re-raise the exception
+                if ce.response["Error"]["Code"] != "ResourceNotFoundException":
+                    raise
@@ -564,33 +564,33 @@ def get_registry_proto(self):
         file_obj = TemporaryFile()
         registry_proto = RegistryProto()
         try:
-            import botocore
+            from botocore.exceptions import ClientError
         except ImportError as e:
             from feast.errors import FeastExtrasDependencyImportError
 
             raise FeastExtrasDependencyImportError("aws", str(e))
         try:
             bucket = self.s3_client.Bucket(self._bucket)
             self.s3_client.meta.client.head_bucket(Bucket=bucket.name)
-        except botocore.client.ClientError as e:
+        except ClientError as e:
             # If a client error is thrown, then check that it was a 404 error.
             # If it was a 404 error, then the bucket does not exist.
             error_code = int(e.response["Error"]["Code"])
             if error_code == 404:
                 raise S3RegistryBucketNotExist(self._bucket)
             else:
-                raise S3RegistryBucketForbiddenAccess(self._bucket)
+                raise S3RegistryBucketForbiddenAccess(self._bucket) from e
 
         try:
             obj = bucket.Object(self._key)
             obj.download_fileobj(file_obj)
             file_obj.seek(0)
             registry_proto.ParseFromString(file_obj.read())
             return registry_proto
-        except botocore.exceptions.ClientError as e:
+        except ClientError as e:
             raise FileNotFoundError(
-                f'Error while trying to locate Registry at path "{self._uri.geturl()}"with [original error]: {e.response}'
-            )
+                f"Error while trying to locate Registry at path {self._uri.geturl()}"
+            ) from e
 
     def update_registry_proto(
         self, updater: Optional[Callable[[RegistryProto], RegistryProto]] = None

@@ -159,7 +159,7 @@ def _validate_offline_store_config(cls, values):
             elif values["provider"] == "gcp":
                 values["offline_store"]["type"] = "bigquery"
             elif values["provider"] == "aws":
-                values["offline_store"]["type"] = "redshift"
+                values["offline_store"]["type"] = "file"
 
         offline_store_type = values["offline_store"]["type"]
 

@@ -71,6 +71,10 @@
     "redis-py-cluster==2.1.2",
 ]
 
+AWS_REQUIRED = [
+    "boto3==1.17.*",
+]
+
 CI_REQUIRED = [
     "cryptography==3.3.2",
     "flake8",
@@ -200,6 +204,7 @@ def run(self):
         "dev": ["mypy-protobuf==1.*", "grpcio-testing==1.*"],
         "ci": CI_REQUIRED,
         "gcp": GCP_REQUIRED,
+        "aws": AWS_REQUIRED,
         "redis": REDIS_REQUIRED,
     },
     include_package_data=True,

@@ -33,8 +33,6 @@ def test_basic() -> None:
             online_store:
                 type: dynamodb
                 region: us-west-2
-            offline_store:
-                type: file
             """
             )
         )

@@ -76,16 +76,9 @@ def feature_store_with_gcs_registry():
 
 @pytest.fixture
 def feature_store_with_s3_registry():
-    import boto3
-
-    s3 = boto3.resource("s3")
-    bucket_name = "feast-integration-tests"
-    bucket = s3.Bucket(bucket_name)
-    s3.meta.client.head_bucket(Bucket=bucket.name)
-
     return FeatureStore(
         config=RepoConfig(
-            registry=f"s3://{bucket_name}/registries/{int(time.time() * 1000)}/registry.db",
+            registry=f"s3://feast-integration-tests/registries/{int(time.time() * 1000)}/registry.db",
             project="default",
             provider="aws",
             online_store=DynamoDBOnlineStoreConfig(region="us-west-2"),

@@ -168,7 +168,7 @@ def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
             join_key="driver_id",
             value_type=ValueType.INT32,
         )
-        with tempfile.TemporaryDirectory() as repo_dir_name, tempfile.TemporaryDirectory():
+        with tempfile.TemporaryDirectory() as repo_dir_name:
             config = RepoConfig(
                 registry=str(Path(repo_dir_name) / "registry.db"),
                 project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
@@ -206,7 +206,7 @@ def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
             join_key="driver_id",
             value_type=ValueType.INT32,
         )
-        with tempfile.TemporaryDirectory() as repo_dir_name, tempfile.TemporaryDirectory():
+        with tempfile.TemporaryDirectory() as repo_dir_name:
             config = RepoConfig(
                 registry=str(Path(repo_dir_name) / "registry.db"),
                 project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",