From 187cdd5ac95d0632fe7a1470a05f9e1cf80b1990 Mon Sep 17 00:00:00 2001 From: Danny Boland Date: Thu, 3 Mar 2022 12:19:58 +0000 Subject: [PATCH] feat: templating for dynamodb table name Signed-off-by: Danny Boland --- .../feast/infra/online_stores/dynamodb.py | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 66c32c1fb8b..e7627131523 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -52,6 +52,9 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): region: StrictStr """ AWS Region Name """ + table_name_template: StrictStr = "{project}.{table_name}" + """ DynamoDB table name template """ + class DynamoDBOnlineStore(OnlineStore): """ @@ -91,7 +94,7 @@ def update( for table_instance in tables_to_keep: try: dynamodb_resource.create_table( - TableName=_get_table_name(config, table_instance), + TableName=_get_table_name(online_config, config, table_instance), KeySchema=[{"AttributeName": "entity_id", "KeyType": "HASH"}], AttributeDefinitions=[ {"AttributeName": "entity_id", "AttributeType": "S"} @@ -107,12 +110,13 @@ def update( for table_instance in tables_to_keep: dynamodb_client.get_waiter("table_exists").wait( - TableName=_get_table_name(config, table_instance) + TableName=_get_table_name(online_config, config, table_instance) ) for table_to_delete in tables_to_delete: _delete_table_idempotent( - dynamodb_resource, _get_table_name(config, table_to_delete) + dynamodb_resource, + _get_table_name(online_config, config, table_to_delete), ) def teardown( @@ -133,7 +137,9 @@ def teardown( dynamodb_resource = self._get_dynamodb_resource(online_config.region) for table in tables: - _delete_table_idempotent(dynamodb_resource, _get_table_name(config, table)) + _delete_table_idempotent( + dynamodb_resource, _get_table_name(online_config, config, table) + ) @log_exceptions_and_usage(online_store="dynamodb") def online_write_batch( @@ -164,7 +170,9 @@ def online_write_batch( assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource(online_config.region) - table_instance = dynamodb_resource.Table(_get_table_name(config, table)) + table_instance = dynamodb_resource.Table( + _get_table_name(online_config, config, table) + ) with table_instance.batch_writer() as batch: for entity_key, features, timestamp, created_ts in data: entity_id = compute_entity_id(entity_key) @@ -206,7 +214,9 @@ def online_read( result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = [] for entity_key in entity_keys: - table_instance = dynamodb_resource.Table(_get_table_name(config, table)) + table_instance = dynamodb_resource.Table( + _get_table_name(online_config, config, table) + ) entity_id = compute_entity_id(entity_key) with tracing_span(name="remote_call"): response = table_instance.get_item(Key={"entity_id": entity_id}) @@ -242,8 +252,12 @@ def _initialize_dynamodb_resource(region: str): return boto3.resource("dynamodb", region_name=region) -def _get_table_name(config: RepoConfig, table: FeatureView) -> str: - return f"{config.project}.{table.name}" +def _get_table_name( + online_config: DynamoDBOnlineStoreConfig, config: RepoConfig, table: FeatureView +) -> str: + return online_config.table_name_template.format( + project=config.project, table_name=table.name + ) def _delete_table_idempotent(