Skip to content

feat: Add boto3 session based auth for dynamodb online store for cross account access #4606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tokoko merged 6 commits into from
Oct 10, 2024
Merged

feat: Add boto3 session based auth for dynamodb online store for cross account access #4606

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0149712
added session auth for dynamodb online store
asingh9530 Oct 5, 2024
c307da7
lint
asingh9530 Oct 5, 2024
9a29b10
Update sdk/python/feast/infra/online_stores/dynamodb.py
asingh9530 Oct 5, 2024
8c7f0d6
Update sdk/python/feast/infra/online_stores/dynamodb.py
asingh9530 Oct 5, 2024
3a4597d
dummy commit
asingh9530 Oct 6, 2024
054c1a6
dummy commit
asingh9530 Oct 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
lint 
Signed-off-by: asingh9530 <abhinav199530singh@gmail.com>
  • Loading branch information
@asingh9530
asingh9530 committed Oct 5, 2024
commit c307da710441e4bdd333d556cc6bcd126001300a
61 changes: 45 additions & 16 deletions sdk/python/feast/infra/online_stores/dynamodb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
session_based_auth: bool = False
Copy link
Copy Markdown
Collaborator

@tokoko tokoko Oct 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is too much bother, I think. We should simply switch to session based, there's no real upside in keeping old behavior, unless I'm missing something.

Copy link
Copy Markdown
Member

@franciscojavierarceo franciscojavierarceo Oct 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Breaking changes create friction to using later versions. I'm fine for getting rid of it in a subsequent release if we give users notice through a warning.

Copy link
Copy Markdown
Collaborator

@tokoko tokoko Oct 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What makes this a breaking change? We're changing an internal detail of how a client is instantiated, it's not a change in an API. can't really imagine how anyone would be depending on the previous behavior of us using a default session.

Copy link
Copy Markdown
Member

@franciscojavierarceo franciscojavierarceo Oct 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I see, this isn't configured through the feature_store.yaml, my bad.

Copy link
Copy Markdown
Contributor Author

@asingh9530 asingh9530 Oct 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tokoko I would still argue we should use this setup in upcoming release and completely switch to session in next just to make sure we don't break anything internally or for feast existing customers although I can't really think if this would break anything but just as a caution.

Copy link
Copy Markdown
Contributor Author

@asingh9530 asingh9530 Oct 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think ?

Copy link
Copy Markdown
Contributor Author

@asingh9530 asingh9530 Oct 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tokoko any update ?

Copy link
Copy Markdown
Collaborator

@tokoko tokoko Oct 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@asingh9530 sorry. I'm fine with proceeding as is, although I think we're being over-cautious. Ironically when we make changes later on to remove this config, the removal will technically be a breaking change 😆 anyway, that's fine as well, let's just not forget to remove it though.

Copy link
Copy Markdown
Contributor Author

@asingh9530 asingh9530 Oct 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lol let me add a warning though if we use default client.

"""AWS session based client authentication"""


class DynamoDBOnlineStore(OnlineStore):
"""
AWS DynamoDB implementation of the online store interface.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Comment thread
asingh9530 marked this conversation as resolved.
Outdated
return boto3.Session().client(
"dynamodb",
Expand All @@ -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:
Comment thread
asingh9530 marked this conversation as resolved.
Outdated
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(
Expand Down