From af305e58ec3c6550b588ffa110fc862abe26d837 Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Mon, 4 Apr 2022 11:50:30 -0500 Subject: [PATCH 1/7] feat: support dynamodb client and resource with endpoint_url Signed-off-by: Miguel Trejo --- .../feast/infra/online_stores/dynamodb.py | 63 ++++++++++++------- 1 file changed, 42 insertions(+), 21 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index c161a5b9554..f20168bc324 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -50,17 +50,20 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): type: Literal["dynamodb"] = "dynamodb" """Online store type selector""" + batch_size: int = 40 + """Number of items to retrieve in a DynamoDB BatchGetItem call.""" + + endpoint_url: str + """DynamoDB local development endpoint Url, i.e. http://localhost:8000""" + region: StrictStr """AWS Region Name""" - table_name_template: StrictStr = "{project}.{table_name}" - """DynamoDB table name template""" - sort_response: bool = True """Whether or not to sort BatchGetItem response.""" - batch_size: int = 40 - """Number of items to retrieve in a DynamoDB BatchGetItem call.""" + table_name_template: StrictStr = "{project}.{table_name}" + """DynamoDB table name template""" class DynamoDBOnlineStore(OnlineStore): @@ -95,8 +98,12 @@ def update( """ online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) - dynamodb_client = self._get_dynamodb_client(online_config.region) - dynamodb_resource = self._get_dynamodb_resource(online_config.region) + dynamodb_client = self._get_dynamodb_client( + online_config.region, online_config.endpoint_url + ) + dynamodb_resource = self._get_dynamodb_resource( + online_config.region, online_config.endpoint_url + ) for table_instance in tables_to_keep: try: @@ -141,7 +148,9 @@ def teardown( """ online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) - dynamodb_resource = self._get_dynamodb_resource(online_config.region) + dynamodb_resource = self._get_dynamodb_resource( + online_config.region, online_config.endpoint_url + ) for table in tables: _delete_table_idempotent( @@ -175,7 +184,9 @@ def online_write_batch( """ online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) - dynamodb_resource = self._get_dynamodb_resource(online_config.region) + dynamodb_resource = self._get_dynamodb_resource( + online_config.region, online_config.endpoint_url + ) table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) @@ -217,7 +228,9 @@ def online_read( """ online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) - dynamodb_resource = self._get_dynamodb_resource(online_config.region) + dynamodb_resource = self._get_dynamodb_resource( + online_config.region, online_config.endpoint_url + ) table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) ) @@ -260,14 +273,16 @@ def online_read( result.extend(batch_size_nones) return result - def _get_dynamodb_client(self, region: str): + def _get_dynamodb_client(self, region: str, endpoint_url: str): if self._dynamodb_client is None: - self._dynamodb_client = _initialize_dynamodb_client(region) + self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url) return self._dynamodb_client - def _get_dynamodb_resource(self, region: str): + def _get_dynamodb_resource(self, region: str, endpoint_url: str): if self._dynamodb_resource is None: - self._dynamodb_resource = _initialize_dynamodb_resource(region) + self._dynamodb_resource = _initialize_dynamodb_resource( + region, endpoint_url + ) return self._dynamodb_resource def _sort_dynamodb_response(self, responses: list, order: list): @@ -285,12 +300,12 @@ def _sort_dynamodb_response(self, responses: list, order: list): return table_responses_ordered -def _initialize_dynamodb_client(region: str): - return boto3.client("dynamodb", region_name=region) +def _initialize_dynamodb_client(region: str, endpoint_url: str = None): + return boto3.client("dynamodb", region_name=region, endpoint_url=endpoint_url) -def _initialize_dynamodb_resource(region: str): - return boto3.resource("dynamodb", region_name=region) +def _initialize_dynamodb_resource(region: str, endpoint_url: str = None): + return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) # TODO(achals): This form of user-facing templating is experimental. @@ -362,8 +377,12 @@ def from_proto(dynamodb_table_proto: DynamoDBTableProto) -> Any: ) def update(self): - dynamodb_client = _initialize_dynamodb_client(region=self.region) - dynamodb_resource = _initialize_dynamodb_resource(region=self.region) + dynamodb_client = _initialize_dynamodb_client( + region=self.region, endpoint_url=None + ) + dynamodb_resource = _initialize_dynamodb_resource( + region=self.region, endpoint_url=None + ) try: dynamodb_resource.create_table( @@ -384,5 +403,7 @@ def update(self): dynamodb_client.get_waiter("table_exists").wait(TableName=f"{self.name}") def teardown(self): - dynamodb_resource = _initialize_dynamodb_resource(region=self.region) + dynamodb_resource = _initialize_dynamodb_resource( + region=self.region, endpoint_url=None + ) _delete_table_idempotent(dynamodb_resource, self.name) From 2180e2cf6daa8644d2a754a96bceb7641635cd35 Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Mon, 4 Apr 2022 12:08:56 -0500 Subject: [PATCH 2/7] fix: overwrite by partitionkey batchwrite operation Signed-off-by: Miguel Trejo --- sdk/python/feast/infra/online_stores/dynamodb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index f20168bc324..d307f53c4b1 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -191,7 +191,7 @@ def online_write_batch( table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) ) - with table_instance.batch_writer() as batch: + with table_instance.batch_writer(overwrite_by_pkeys=["entity_id"]) as batch: for entity_key, features, timestamp, created_ts in data: entity_id = compute_entity_id(entity_key) batch.put_item( From 02f8071be5f39af03bdec08670792e0b89cf0004 Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Mon, 4 Apr 2022 12:23:36 -0500 Subject: [PATCH 3/7] docs: how to configure local dynamodb Signed-off-by: Miguel Trejo --- CONTRIBUTING.md | 3 ++- .../tests/integration/feature_repos/repo_configuration.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ea1dd8ad13d..b0ef72b65fb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -154,7 +154,8 @@ To test across clouds, on top of setting up Redis, you also need GCP / AWS / Sno **AWS** 1. TODO(adchia): flesh out setting up AWS login (or create helper script) -2. Modify `RedshiftDataSourceCreator` to use your credentials +2. To avoid AWS fees `DynamoDBOnlineStore` can be tested locally if you deploy DynamoDB Locally on your Computer, for this setup a [Local DynamoDB on your Computer](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.DownloadingAndRunning.html) +3. Modify `RedshiftDataSourceCreator` to use your credentials **Snowflake** - See https://signup.snowflake.com/ diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py index 3029b2e7740..17fdd2e5b3c 100644 --- a/sdk/python/tests/integration/feature_repos/repo_configuration.py +++ b/sdk/python/tests/integration/feature_repos/repo_configuration.py @@ -5,6 +5,7 @@ import re import tempfile import uuid +from copy import deepcopy from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path @@ -69,9 +70,13 @@ IntegrationTestRepoConfig(python_feature_server=True), ] if os.getenv("FEAST_IS_LOCAL_TEST", "False") != "True": + # Local integration tests use local Dynamodb + LOCAL_DYNAMO_CONFIG = deepcopy(DYNAMO_CONFIG) + LOCAL_DYNAMO_CONFIG["endpoint_url"] = "http://localhost:8000" DEFAULT_FULL_REPO_CONFIGS.extend( [ IntegrationTestRepoConfig(online_store=REDIS_CONFIG), + IntegrationTestRepoConfig(online_store=LOCAL_DYNAMO_CONFIG), # GCP configurations IntegrationTestRepoConfig( provider="gcp", @@ -87,7 +92,7 @@ IntegrationTestRepoConfig( provider="aws", offline_store_creator=RedshiftDataSourceCreator, - online_store=DYNAMO_CONFIG, + online_store=LOCAL_DYNAMO_CONFIG, python_feature_server=True, ), IntegrationTestRepoConfig( From 76efcea13b712a98bf58b54fe1711d577f8b3e0f Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Mon, 4 Apr 2022 12:31:54 -0500 Subject: [PATCH 4/7] fix: DynamoDBonlineStore endpoint_url defaults to None Signed-off-by: Miguel Trejo --- sdk/python/feast/infra/online_stores/dynamodb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index d307f53c4b1..7c150455a60 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -53,7 +53,7 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): batch_size: int = 40 """Number of items to retrieve in a DynamoDB BatchGetItem call.""" - endpoint_url: str + endpoint_url: str = None """DynamoDB local development endpoint Url, i.e. http://localhost:8000""" region: StrictStr From 0440598922d111af932d2488ae28e4e37bc6a75f Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Mon, 4 Apr 2022 13:18:34 -0500 Subject: [PATCH 5/7] docs: setup dummy aws credentials Signed-off-by: Miguel Trejo --- CONTRIBUTING.md | 12 +++++++++++- sdk/python/feast/infra/online_stores/dynamodb.py | 4 ++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b0ef72b65fb..c6c0256f6b6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -154,7 +154,17 @@ To test across clouds, on top of setting up Redis, you also need GCP / AWS / Sno **AWS** 1. TODO(adchia): flesh out setting up AWS login (or create helper script) -2. To avoid AWS fees `DynamoDBOnlineStore` can be tested locally if you deploy DynamoDB Locally on your Computer, for this setup a [Local DynamoDB on your Computer](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.DownloadingAndRunning.html) +2. To avoid AWS fees `DynamoDBOnlineStore` can be tested locally if you deploy DynamoDB Locally on your Computer, for this + + A. Setup a [Local DynamoDB on your Computer](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.DownloadingAndRunning.html). + + B. Setup dummy AWS Credentials and `us-west-2` as region. (Optional: Only if testing the `DynamoDBOnlineStore`) + + ```txt + export AWS_ACCESS_KEY_ID: 'DUMMYIDEXAMPLE' + export AWS_SECRET_ACCESS_KEY: 'DUMMYEXAMPLEKEY' + export AWS_DEFAULT_REGION: 'us-west-2' + ``` 3. Modify `RedshiftDataSourceCreator` to use your credentials **Snowflake** diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 7c150455a60..551d554e1fe 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -17,7 +17,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from pydantic import StrictStr -from pydantic.typing import Literal +from pydantic.typing import Literal, Optional from feast import Entity, FeatureView, utils from feast.infra.infra_object import DYNAMODB_INFRA_OBJECT_CLASS_TYPE, InfraObject @@ -53,7 +53,7 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): batch_size: int = 40 """Number of items to retrieve in a DynamoDB BatchGetItem call.""" - endpoint_url: str = None + endpoint_url: Optional[str] = None """DynamoDB local development endpoint Url, i.e. http://localhost:8000""" region: StrictStr From 7d23380752e766af950e4678fd3acf3a7b7e0be4 Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Tue, 5 Apr 2022 22:34:41 -0500 Subject: [PATCH 6/7] test: DynamoDBOnlineStoreConfig endpoint_url configuration Signed-off-by: Miguel Trejo --- CONTRIBUTING.md | 13 +--- .../feast/infra/online_stores/dynamodb.py | 14 ++-- .../feature_repos/repo_configuration.py | 7 +- .../test_dynamodb_online_store.py | 64 +++++++++++++++++++ 4 files changed, 73 insertions(+), 25 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c6c0256f6b6..ea1dd8ad13d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -154,18 +154,7 @@ To test across clouds, on top of setting up Redis, you also need GCP / AWS / Sno **AWS** 1. TODO(adchia): flesh out setting up AWS login (or create helper script) -2. To avoid AWS fees `DynamoDBOnlineStore` can be tested locally if you deploy DynamoDB Locally on your Computer, for this - - A. Setup a [Local DynamoDB on your Computer](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.DownloadingAndRunning.html). - - B. Setup dummy AWS Credentials and `us-west-2` as region. (Optional: Only if testing the `DynamoDBOnlineStore`) - - ```txt - export AWS_ACCESS_KEY_ID: 'DUMMYIDEXAMPLE' - export AWS_SECRET_ACCESS_KEY: 'DUMMYEXAMPLEKEY' - export AWS_DEFAULT_REGION: 'us-west-2' - ``` -3. Modify `RedshiftDataSourceCreator` to use your credentials +2. Modify `RedshiftDataSourceCreator` to use your credentials **Snowflake** - See https://signup.snowflake.com/ diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 551d554e1fe..058b7c29d53 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -17,7 +17,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from pydantic import StrictStr -from pydantic.typing import Literal, Optional +from pydantic.typing import Literal, Union from feast import Entity, FeatureView, utils from feast.infra.infra_object import DYNAMODB_INFRA_OBJECT_CLASS_TYPE, InfraObject @@ -53,7 +53,7 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): batch_size: int = 40 """Number of items to retrieve in a DynamoDB BatchGetItem call.""" - endpoint_url: Optional[str] = None + endpoint_url: Union[str, None] = None """DynamoDB local development endpoint Url, i.e. http://localhost:8000""" region: StrictStr @@ -191,7 +191,7 @@ def online_write_batch( table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) ) - with table_instance.batch_writer(overwrite_by_pkeys=["entity_id"]) as batch: + with table_instance.batch_writer() as batch: for entity_key, features, timestamp, created_ts in data: entity_id = compute_entity_id(entity_key) batch.put_item( @@ -273,12 +273,12 @@ def online_read( result.extend(batch_size_nones) return result - def _get_dynamodb_client(self, region: str, endpoint_url: str): + def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None): if self._dynamodb_client is None: self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url) return self._dynamodb_client - def _get_dynamodb_resource(self, region: str, endpoint_url: str): + def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None): if self._dynamodb_resource is None: self._dynamodb_resource = _initialize_dynamodb_resource( region, endpoint_url @@ -300,11 +300,11 @@ def _sort_dynamodb_response(self, responses: list, order: list): return table_responses_ordered -def _initialize_dynamodb_client(region: str, endpoint_url: str = None): +def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None): return boto3.client("dynamodb", region_name=region, endpoint_url=endpoint_url) -def _initialize_dynamodb_resource(region: str, endpoint_url: str = None): +def _initialize_dynamodb_resource(region: str, endpoint_url: Optional[str] = None): return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py index 17fdd2e5b3c..3029b2e7740 100644 --- a/sdk/python/tests/integration/feature_repos/repo_configuration.py +++ b/sdk/python/tests/integration/feature_repos/repo_configuration.py @@ -5,7 +5,6 @@ import re import tempfile import uuid -from copy import deepcopy from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path @@ -70,13 +69,9 @@ IntegrationTestRepoConfig(python_feature_server=True), ] if os.getenv("FEAST_IS_LOCAL_TEST", "False") != "True": - # Local integration tests use local Dynamodb - LOCAL_DYNAMO_CONFIG = deepcopy(DYNAMO_CONFIG) - LOCAL_DYNAMO_CONFIG["endpoint_url"] = "http://localhost:8000" DEFAULT_FULL_REPO_CONFIGS.extend( [ IntegrationTestRepoConfig(online_store=REDIS_CONFIG), - IntegrationTestRepoConfig(online_store=LOCAL_DYNAMO_CONFIG), # GCP configurations IntegrationTestRepoConfig( provider="gcp", @@ -92,7 +87,7 @@ IntegrationTestRepoConfig( provider="aws", offline_store_creator=RedshiftDataSourceCreator, - online_store=LOCAL_DYNAMO_CONFIG, + online_store=DYNAMO_CONFIG, python_feature_server=True, ), IntegrationTestRepoConfig( diff --git a/sdk/python/tests/unit/online_store/test_dynamodb_online_store.py b/sdk/python/tests/unit/online_store/test_dynamodb_online_store.py index 0f42230ef53..b15edd325be 100644 --- a/sdk/python/tests/unit/online_store/test_dynamodb_online_store.py +++ b/sdk/python/tests/unit/online_store/test_dynamodb_online_store.py @@ -38,6 +38,70 @@ def repo_config(): ) +def test_online_store_config_default(): + """Test DynamoDBOnlineStoreConfig default parameters.""" + aws_region = "us-west-2" + dynamodb_store_config = DynamoDBOnlineStoreConfig(region=aws_region) + assert dynamodb_store_config.type == "dynamodb" + assert dynamodb_store_config.batch_size == 40 + assert dynamodb_store_config.endpoint_url is None + assert dynamodb_store_config.region == aws_region + assert dynamodb_store_config.sort_response is True + assert dynamodb_store_config.table_name_template == "{project}.{table_name}" + + +def test_online_store_config_custom_params(): + """Test DynamoDBOnlineStoreConfig custom parameters.""" + aws_region = "us-west-2" + batch_size = 20 + endpoint_url = "http://localhost:8000" + sort_response = False + table_name_template = "feast_test.dynamodb_table" + dynamodb_store_config = DynamoDBOnlineStoreConfig( + region=aws_region, + batch_size=batch_size, + endpoint_url=endpoint_url, + sort_response=sort_response, + table_name_template=table_name_template, + ) + assert dynamodb_store_config.type == "dynamodb" + assert dynamodb_store_config.batch_size == batch_size + assert dynamodb_store_config.endpoint_url == endpoint_url + assert dynamodb_store_config.region == aws_region + assert dynamodb_store_config.sort_response == sort_response + assert dynamodb_store_config.table_name_template == table_name_template + + +def test_online_store_config_dynamodb_client(): + """Test DynamoDBOnlineStoreConfig configure DynamoDB client with endpoint_url.""" + aws_region = "us-west-2" + endpoint_url = "http://localhost:8000" + dynamodb_store = DynamoDBOnlineStore() + dynamodb_store_config = DynamoDBOnlineStoreConfig( + region=aws_region, endpoint_url=endpoint_url + ) + dynamodb_client = dynamodb_store._get_dynamodb_client( + dynamodb_store_config.region, dynamodb_store_config.endpoint_url + ) + assert dynamodb_client.meta.region_name == aws_region + assert dynamodb_client.meta.endpoint_url == endpoint_url + + +def test_online_store_config_dynamodb_resource(): + """Test DynamoDBOnlineStoreConfig configure DynamoDB Resource with endpoint_url.""" + aws_region = "us-west-2" + endpoint_url = "http://localhost:8000" + dynamodb_store = DynamoDBOnlineStore() + dynamodb_store_config = DynamoDBOnlineStoreConfig( + region=aws_region, endpoint_url=endpoint_url + ) + dynamodb_resource = dynamodb_store._get_dynamodb_resource( + dynamodb_store_config.region, dynamodb_store_config.endpoint_url + ) + assert dynamodb_resource.meta.client.meta.region_name == aws_region + assert dynamodb_resource.meta.client.meta.endpoint_url == endpoint_url + + @mock_dynamodb2 @pytest.mark.parametrize("n_samples", [5, 50, 100]) def test_online_read(repo_config, n_samples): From 77c9f134d3f39c24930c5623584422c575d7df9c Mon Sep 17 00:00:00 2001 From: Miguel Trejo Date: Thu, 7 Apr 2022 20:32:59 -0500 Subject: [PATCH 7/7] feat: DynamoDBTable support endpoint_url Signed-off-by: Miguel Trejo --- .../feast/infra/online_stores/dynamodb.py | 33 ++++++++---- .../test_dynamodb_online_store.py | 52 +++++++++++++++++++ 2 files changed, 75 insertions(+), 10 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 058b7c29d53..61334be1a92 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -342,13 +342,20 @@ class DynamoDBTable(InfraObject): Attributes: name: The name of the table. region: The region of the table. + endpoint_url: Local DynamoDB Endpoint Url. + _dynamodb_client: Boto3 DynamoDB client. + _dynamodb_resource: Boto3 DynamoDB resource. """ region: str + endpoint_url = None + _dynamodb_client = None + _dynamodb_resource = None - def __init__(self, name: str, region: str): + def __init__(self, name: str, region: str, endpoint_url: Optional[str] = None): super().__init__(name) self.region = region + self.endpoint_url = endpoint_url def to_infra_object_proto(self) -> InfraObjectProto: dynamodb_table_proto = self.to_proto() @@ -377,12 +384,8 @@ def from_proto(dynamodb_table_proto: DynamoDBTableProto) -> Any: ) def update(self): - dynamodb_client = _initialize_dynamodb_client( - region=self.region, endpoint_url=None - ) - dynamodb_resource = _initialize_dynamodb_resource( - region=self.region, endpoint_url=None - ) + dynamodb_client = self._get_dynamodb_client(self.region, self.endpoint_url) + dynamodb_resource = self._get_dynamodb_resource(self.region, self.endpoint_url) try: dynamodb_resource.create_table( @@ -403,7 +406,17 @@ def update(self): dynamodb_client.get_waiter("table_exists").wait(TableName=f"{self.name}") def teardown(self): - dynamodb_resource = _initialize_dynamodb_resource( - region=self.region, endpoint_url=None - ) + dynamodb_resource = self._get_dynamodb_resource(self.region, self.endpoint_url) _delete_table_idempotent(dynamodb_resource, self.name) + + def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None): + if self._dynamodb_client is None: + self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url) + return self._dynamodb_client + + def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None): + if self._dynamodb_resource is None: + self._dynamodb_resource = _initialize_dynamodb_resource( + region, endpoint_url + ) + return self._dynamodb_resource diff --git a/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py b/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py index b15edd325be..7b0c5a4a619 100644 --- a/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py +++ b/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py @@ -7,6 +7,7 @@ from feast.infra.online_stores.dynamodb import ( DynamoDBOnlineStore, DynamoDBOnlineStoreConfig, + DynamoDBTable, ) from feast.repo_config import RepoConfig from tests.utils.online_store_utils import ( @@ -50,6 +51,18 @@ def test_online_store_config_default(): assert dynamodb_store_config.table_name_template == "{project}.{table_name}" +def test_dynamodb_table_default_params(): + """Test DynamoDBTable default parameters.""" + tbl_name = "dynamodb-test" + aws_region = "us-west-2" + dynamodb_table = DynamoDBTable(tbl_name, aws_region) + assert dynamodb_table.name == tbl_name + assert dynamodb_table.region == aws_region + assert dynamodb_table.endpoint_url is None + assert dynamodb_table._dynamodb_client is None + assert dynamodb_table._dynamodb_resource is None + + def test_online_store_config_custom_params(): """Test DynamoDBOnlineStoreConfig custom parameters.""" aws_region = "us-west-2" @@ -72,6 +85,19 @@ def test_online_store_config_custom_params(): assert dynamodb_store_config.table_name_template == table_name_template +def test_dynamodb_table_custom_params(): + """Test DynamoDBTable custom parameters.""" + tbl_name = "dynamodb-test" + aws_region = "us-west-2" + endpoint_url = "http://localhost:8000" + dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url) + assert dynamodb_table.name == tbl_name + assert dynamodb_table.region == aws_region + assert dynamodb_table.endpoint_url == endpoint_url + assert dynamodb_table._dynamodb_client is None + assert dynamodb_table._dynamodb_resource is None + + def test_online_store_config_dynamodb_client(): """Test DynamoDBOnlineStoreConfig configure DynamoDB client with endpoint_url.""" aws_region = "us-west-2" @@ -87,6 +113,19 @@ def test_online_store_config_dynamodb_client(): assert dynamodb_client.meta.endpoint_url == endpoint_url +def test_dynamodb_table_dynamodb_client(): + """Test DynamoDBTable configure DynamoDB client with endpoint_url.""" + tbl_name = "dynamodb-test" + aws_region = "us-west-2" + endpoint_url = "http://localhost:8000" + dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url) + dynamodb_client = dynamodb_table._get_dynamodb_client( + dynamodb_table.region, dynamodb_table.endpoint_url + ) + assert dynamodb_client.meta.region_name == aws_region + assert dynamodb_client.meta.endpoint_url == endpoint_url + + def test_online_store_config_dynamodb_resource(): """Test DynamoDBOnlineStoreConfig configure DynamoDB Resource with endpoint_url.""" aws_region = "us-west-2" @@ -102,6 +141,19 @@ def test_online_store_config_dynamodb_resource(): assert dynamodb_resource.meta.client.meta.endpoint_url == endpoint_url +def test_dynamodb_table_dynamodb_resource(): + """Test DynamoDBTable configure DynamoDB resource with endpoint_url.""" + tbl_name = "dynamodb-test" + aws_region = "us-west-2" + endpoint_url = "http://localhost:8000" + dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url) + dynamodb_resource = dynamodb_table._get_dynamodb_resource( + dynamodb_table.region, dynamodb_table.endpoint_url + ) + assert dynamodb_resource.meta.client.meta.region_name == aws_region + assert dynamodb_resource.meta.client.meta.endpoint_url == endpoint_url + + @mock_dynamodb2 @pytest.mark.parametrize("n_samples", [5, 50, 100]) def test_online_read(repo_config, n_samples):