From 01497126ab9d42482d3403179d7d066082453af5 Mon Sep 17 00:00:00 2001 From: asingh9530 Date: Sat, 5 Oct 2024 18:28:30 +0530 Subject: [PATCH 1/6] added session auth for dynamodb online store Signed-off-by: asingh9530 --- .../feast/infra/online_stores/dynamodb.py | 53 ++++++++++++------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index b2488543b02..289347e52ee 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -70,6 +70,8 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): tags: Union[Dict[str, str], None] = None """AWS resource tags added to each table""" + session_based_auth: bool = False + """AWS session based client authentication""" class DynamoDBOnlineStore(OnlineStore): """ @@ -104,10 +106,10 @@ def update( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_client = self._get_dynamodb_client( - online_config.region, online_config.endpoint_url + online_config.region, online_config.endpoint_url, online_config.session_based_auth ) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url + online_config.region, online_config.endpoint_url, online_config.session_based_auth ) # Add Tags attribute to creation request only if configured to prevent # TagResource permission issues, even with an empty Tags array. @@ -166,7 +168,7 @@ def teardown( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url + online_config.region, online_config.endpoint_url, online_config.session_based_auth ) for table in tables: @@ -201,7 +203,7 @@ def online_write_batch( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url + online_config.region, online_config.endpoint_url, online_config.session_based_auth ) table_instance = dynamodb_resource.Table( @@ -228,7 +230,7 @@ def online_read( assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url + online_config.region, online_config.endpoint_url, online_config.session_based_auth ) table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) @@ -323,15 +325,15 @@ def _get_aioboto_session(self): def _get_aiodynamodb_client(self, region: str): return self._get_aioboto_session().create_client("dynamodb", region_name=region) - def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None): + def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): if self._dynamodb_client is None: - self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url) + self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url, session_based_auth) return self._dynamodb_client - def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None): + def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): if self._dynamodb_resource is None: self._dynamodb_resource = _initialize_dynamodb_resource( - region, endpoint_url + region, endpoint_url, session_based_auth ) return self._dynamodb_resource @@ -443,17 +445,32 @@ def _to_client_batch_get_payload(online_config, table_name, batch): } -def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None): - return boto3.client( - "dynamodb", - region_name=region, - endpoint_url=endpoint_url, - config=Config(user_agent=get_user_agent()), - ) +def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): + + if session_based_auth is True: + return boto3.Session().client( + "dynamodb", + region_name=region, + endpoint_url=endpoint_url, + config=Config(user_agent=get_user_agent()), + ) + else: + return boto3.client( + "dynamodb", + region_name=region, + endpoint_url=endpoint_url, + config=Config(user_agent=get_user_agent()) + + ) + + +def _initialize_dynamodb_resource(region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): + if session_based_auth is True: + return boto3.Session().resource("dynamodb", region_name=region, endpoint_url=endpoint_url) + else: + return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) -def _initialize_dynamodb_resource(region: str, endpoint_url: Optional[str] = None): - return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) # TODO(achals): This form of user-facing templating is experimental. From c307da710441e4bdd333d556cc6bcd126001300a Mon Sep 17 00:00:00 2001 From: asingh9530 Date: Sat, 5 Oct 2024 18:31:25 +0530 Subject: [PATCH 2/6] lint Signed-off-by: asingh9530 --- .../feast/infra/online_stores/dynamodb.py | 61 ++++++++++++++----- 1 file changed, 45 insertions(+), 16 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index 289347e52ee..7f70dcf925f 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -73,6 +73,7 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel): session_based_auth: bool = False """AWS session based client authentication""" + class DynamoDBOnlineStore(OnlineStore): """ AWS DynamoDB implementation of the online store interface. @@ -106,10 +107,14 @@ def update( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_client = self._get_dynamodb_client( - online_config.region, online_config.endpoint_url, online_config.session_based_auth + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, ) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url, online_config.session_based_auth + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, ) # Add Tags attribute to creation request only if configured to prevent # TagResource permission issues, even with an empty Tags array. @@ -168,7 +173,9 @@ def teardown( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url, online_config.session_based_auth + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, ) for table in tables: @@ -203,7 +210,9 @@ def online_write_batch( online_config = config.online_store assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url, online_config.session_based_auth + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, ) table_instance = dynamodb_resource.Table( @@ -230,7 +239,9 @@ def online_read( assert isinstance(online_config, DynamoDBOnlineStoreConfig) dynamodb_resource = self._get_dynamodb_resource( - online_config.region, online_config.endpoint_url, online_config.session_based_auth + online_config.region, + online_config.endpoint_url, + online_config.session_based_auth, ) table_instance = dynamodb_resource.Table( _get_table_name(online_config, config, table) @@ -325,12 +336,24 @@ def _get_aioboto_session(self): def _get_aiodynamodb_client(self, region: str): return self._get_aioboto_session().create_client("dynamodb", region_name=region) - def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): + def _get_dynamodb_client( + self, + region: str, + endpoint_url: Optional[str] = None, + session_based_auth: Optional[bool] = False, + ): if self._dynamodb_client is None: - self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url, session_based_auth) + self._dynamodb_client = _initialize_dynamodb_client( + region, endpoint_url, session_based_auth + ) return self._dynamodb_client - def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): + def _get_dynamodb_resource( + self, + region: str, + endpoint_url: Optional[str] = None, + session_based_auth: Optional[bool] = False, + ): if self._dynamodb_resource is None: self._dynamodb_resource = _initialize_dynamodb_resource( region, endpoint_url, session_based_auth @@ -445,8 +468,11 @@ def _to_client_batch_get_payload(online_config, table_name, batch): } -def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): - +def _initialize_dynamodb_client( + region: str, + endpoint_url: Optional[str] = None, + session_based_auth: Optional[bool] = False, +): if session_based_auth is True: return boto3.Session().client( "dynamodb", @@ -459,20 +485,23 @@ def _initialize_dynamodb_client(region: str, endpoint_url: Optional[str] = None, "dynamodb", region_name=region, endpoint_url=endpoint_url, - config=Config(user_agent=get_user_agent()) - + config=Config(user_agent=get_user_agent()), ) -def _initialize_dynamodb_resource(region: str, endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False): +def _initialize_dynamodb_resource( + region: str, + endpoint_url: Optional[str] = None, + session_based_auth: Optional[bool] = False, +): if session_based_auth is True: - return boto3.Session().resource("dynamodb", region_name=region, endpoint_url=endpoint_url) + return boto3.Session().resource( + "dynamodb", region_name=region, endpoint_url=endpoint_url + ) else: return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) - - # TODO(achals): This form of user-facing templating is experimental. # Please refer to https://github.com/feast-dev/feast/issues/2438 before building on top of it, def _get_table_name( From 9a29b10687913ff24b86846d9380fecc61f5f787 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sat, 5 Oct 2024 18:59:33 +0530 Subject: [PATCH 3/6] Update sdk/python/feast/infra/online_stores/dynamodb.py Co-authored-by: Francisco Arceo Signed-off-by: asingh9530 --- 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 7f70dcf925f..6be77e75574 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -473,7 +473,7 @@ def _initialize_dynamodb_client( endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False, ): - if session_based_auth is True: + if session_based_auth: return boto3.Session().client( "dynamodb", region_name=region, From 8c7f0d6ada18ae5d5365a2a22e5a60cda4c51511 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sat, 5 Oct 2024 18:59:39 +0530 Subject: [PATCH 4/6] Update sdk/python/feast/infra/online_stores/dynamodb.py Co-authored-by: Francisco Arceo Signed-off-by: asingh9530 --- 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 6be77e75574..c0494272b34 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -494,7 +494,7 @@ def _initialize_dynamodb_resource( endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False, ): - if session_based_auth is True: + if session_based_auth: return boto3.Session().resource( "dynamodb", region_name=region, endpoint_url=endpoint_url ) From 3a4597d5b2b979186cfb29d1de3814e7ffd9db26 Mon Sep 17 00:00:00 2001 From: asingh9530 Date: Sun, 6 Oct 2024 16:22:18 +0530 Subject: [PATCH 5/6] dummy commit Signed-off-by: asingh9530 --- sdk/python/feast/infra/online_stores/dynamodb.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index c0494272b34..fd2a9716a27 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -473,6 +473,7 @@ def _initialize_dynamodb_client( endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False, ): + if session_based_auth: return boto3.Session().client( "dynamodb", From 054c1a6bd45932f7dcbfa172de6da92e3cb73cac Mon Sep 17 00:00:00 2001 From: asingh9530 Date: Sun, 6 Oct 2024 16:22:52 +0530 Subject: [PATCH 6/6] dummy commit Signed-off-by: asingh9530 --- sdk/python/feast/infra/online_stores/dynamodb.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index fd2a9716a27..c0494272b34 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -473,7 +473,6 @@ def _initialize_dynamodb_client( endpoint_url: Optional[str] = None, session_based_auth: Optional[bool] = False, ): - if session_based_auth: return boto3.Session().client( "dynamodb",