Skip to content
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

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

Merged
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 all commits
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
82 changes: 64 additions & 18 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 @@ -70,6 +70,9 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
tags: Union[Dict[str, str], None] = None
"""AWS resource tags added to each table"""

session_based_auth: bool = False
Copy link
Collaborator

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
Member

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
Collaborator

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
Member

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
Contributor Author

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
Contributor Author

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
Contributor Author

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
Collaborator

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
Contributor Author

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):
"""
Expand Down Expand Up @@ -104,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.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.
Expand Down Expand Up @@ -166,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.region,
online_config.endpoint_url,
online_config.session_based_auth,
)

for table in tables:
Expand Down Expand Up @@ -201,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.region,
online_config.endpoint_url,
online_config.session_based_auth,
)

table_instance = dynamodb_resource.Table(
Expand All @@ -228,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.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 @@ -323,15 +336,27 @@ 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

Expand Down Expand Up @@ -443,17 +468,38 @@ 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:
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):
return boto3.resource("dynamodb", region_name=region, endpoint_url=endpoint_url)
def _initialize_dynamodb_resource(
region: str,
endpoint_url: Optional[str] = None,
session_based_auth: Optional[bool] = False,
):
if session_based_auth:
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.
Expand Down
Loading