Amazon Provider: Consistency in boto3-based Hooks

[Taragolis]

Hello Airflow folks, I hope you are doing well and Airflow pipelines work as expected.

I have a proposal/idea about AWS Hooks which interact with AWS API by boto3 (and botocore as Éminence grise)

Goals

Lets starts with what I expected to have in Amazon Provider Hooks:

• Truly generic thin wrapper around boto3.session.Session, boto3.client (low-level client which usual refers to specific botocore Client generated from json data-definitions) and boto3.resource (high-level client) without define method specific parameters in hook constructor, some sample of hooks which break this idea:
  • DynamoDBHook: table names and table keys
  • BatchClientHook class attributes as hook configurations, class instance attributes which not relevant for all methods
  • S3Hook: Defined as client but also required resource-specific objects
  • Ec2Hook: could create boto3.client or boto3.resource, depend on hook specific constructor attribute api_type
• Type Annotations for boto3 clients and resources: More type hints for Amazon provider #11297

Demo

boto3-hooks.mp4

Suggested Implementation

This is only idea if I had solution I would create PR rather than discussion. I have some thoughts about how it should looks like
All new naming just for reference, you could find PoC implementation in Gist

Current Inheritance model

Proposal Inheritance model

Changes in attributes

New attributes (cached):

• service_name will replace client_type and resource_type attributes

New properties (cached):

• session return cached session associated with Hook
• client either Boto3ResourceHook and Boto3ClientHook will provide boto3.client
• resource only Boto3ResourceHook will provide boto3.resource

Deprecate:

• conn this property as well as get_conn depend on client_type and resource_type and does not required anymore if we have session, client and resource.

Backward compatibility

Without backward compatibility it could turned into

Idea is add legacy_mode attribute and set it implicit to True for all existed hooks.
Property conn and get_conn will provide compatibility with previous version of Hooks for a while

Client vs Resource

Service / Resources / Hooks statistic

Python 3.9.10 (main, Feb 25 2022, 16:54:01) 
[Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import botocore
>>> import boto3

>>> botocore.__version__
'1.29.35'
>>> boto3.__version__
'1.26.35'

>>> session = boto3.session.Session(region_name="us-east-1")

>>> len(session.get_available_services())
336

>>> session.get_available_resources()
['cloudformation', 'cloudwatch', 'dynamodb', 'ec2', 'glacier', 'iam', 'opsworks', 's3', 'sns', 'sqs']


>>> from airflow.providers_manager import ProvidersManager

>>> aws_provider = ProvidersManager().providers["apache-airflow-providers-amazon"]
>>>len(aws_provider.data["hooks"])
31

According to internals: AWS have 336 different services,10 of them have high-level client and in total we have 31 hooks in Amazon Provider.

How to decide is new hook should based on Boto3ResourceHook or Boto3ClientHook ?

boto3.resource is high level wrapper around boto3.client that mean if boto3 supports resource for specific service we should use Boto3ResourceHook otherwise Boto3ClientHook.

Also I suggest migrate existed hooks to Resource-based

• cloudformation: Current implementation airflow.providers.amazon.aws.hooks.cloud_formation.CloudFormationHook(client_type="cloudformation")
• cloudwatch: Not existed (not Cloudwatch Logs provided by logs)
• dynamodb: Current implementation airflow.providers.amazon.aws.hooks.dynamodb.DynamoDBHook(resource_type="dynamodb")
• ec2: Current implementation airflow.providers.amazon.aws.hooks.ec2.EC2Hook(api_type="resource_type" or api_type="client_type")
  glacier: Current implementation airflow.providers.amazon.aws.hooks.glacier.GlacierHook(client_type="glacier")
  iam: Not existed
  opsworks: Not existed
  s3: Current implementation airflow.providers.amazon.aws.hooks.s3.S3Hook(client_type="s3") and also it manually create resource for some methods
  sns: Current implementation airflow.providers.amazon.aws.hooks.sns.SnsHook(client_type="sns")
  sqs: Current implementation airflow.providers.amazon.aws.hooks.sqs.SqsHook(client_type="sqs")

---

Conclusion

I'd like to hear any feedback/ideas/objections before create Issues/Tasks/PRs. Thanks for your time!

cc: @o-nikolas @ferruzzi @eladkal @shubham22 @potiuk

[ferruzzi]

Just got back from a vacation, sorry for the delay. I have read through this and want to take a minute to digest it. My initial thought is that this is the time to make a change I've been wanting to do. This doesn't necessarily need to be done at the same time, but if we are designing a new abstraction layer/interface then it might be a good time to think about how we could implement this. If we are making breaking changes, then some of the existing hooks can be converted to what we've been calling "thin hooks".

There are hooks which just take values and pass them through to boto. As one example, see here:

    def start_replication_task(
        self,
        replication_task_arn: str,
        start_replication_task_type: str,
        **kwargs,
    ):
        """
        Starts replication task.

        :param replication_task_arn: Replication task ARN
        :param start_replication_task_type: Replication task start type (default='start-replication')
            ('start-replication'|'resume-processing'|'reload-target')
        """
        dms_client = self.get_conn()
        dms_client.start_replication_task(
            ReplicationTaskArn=replication_task_arn,
            StartReplicationTaskType=start_replication_task_type,
            **kwargs,
        )

In these cases the method basically just converts python's snake_case to boto's CamelCase and adds IDE completion, but I always wonder if there is a better way to get those benefits using some interface like you are suggesting.

The best I have come up with so far is to replace those with a pass method which would inherit the boto call. We lose the python_case parameter names this way, but see the existing BatchProtocol for an example. Batch support right now is not handled like any other hook and should eventually be brought into line with other hooks, but it is an example of the pass method in action.

def describe_jobs(self, jobs: list[str]) -> dict:
    """
    Get job descriptions from AWS Batch
    :param jobs: a list of JobId to describe
    :return: an API response to describe jobs
    """
    ...

For example, without any changes to the existing Athena hook right now you can do something like AthenaHook().list_engine_versions() without having to add it to the hook.

IFF we go through with a change like that then another thing that would pair really well with it in the new interface would be a way to add a "wait for completion" block (new decorator/wrapper??) to any arbitrary boto call. We have been going through and adding wait_for_completion to hooks and operators left and right and I think if there was a convenient way to optionally wrap any hook in a wait block like that, it would allow us to clean out a ton of boilerplate code in the hooks. I don't have a proposal on what that would look like, but thought I'd throw it out there in case hat turns out to be something easy to implement in the new layer.

[Taragolis]

In these cases the method basically just converts python's snake_case to boto's CamelCase and adds IDE completion, but I always wonder if there is a better way to get those benefits using some interface like you are suggesting.

I don't know how to do this in runtime for IDE (without performance degradation) however all information stored into data files of botocore

Service definition for StartReplicationTask API call:
https://github.com/boto/botocore/blob/1d058bac6d8667db0606e036f4100e76357663dc/botocore/data/dms/2016-01-01/service-2.json#L858-L872

Input Parameters
https://github.com/boto/botocore/blob/1d058bac6d8667db0606e036f4100e76357663dc/botocore/data/dms/2016-01-01/service-2.json#L6176-L6205

Some of this info exists into client.meta, you could try to find this info by debugger

meta = boto3.client("dms", region_name="eu-west-1").meta

BTW, boto3-stubs packages use same info for generate boto3 client annotations

[ferruzzi]

yeah, I can get the list of all possible input parameters for a given service this way:

In [1]: from botocore.data import eks
     ...: import json

In [2]: with open(eks.__path__[0] + '/2017-11-01/service-2.json') as file:
     ...:     parsed_file = json.load(file)
     ...:

In [3]: parsed_file['shapes']['CreateNodegroupRequest']['members'].keys()
Out[3]: dict_keys(['clusterName', 'nodegroupName', 'scalingConfig', 'diskSize', 'subnets', 'instanceTypes', 'amiType', 'remoteAccess', 'nodeRole', 'labels', 'tags', 'clientRequestToken', 'launchTemplate', 'version', 'releaseVersion'])

Anyway, I didn't mean to derail your discussion, just a couple of ideas for the future that might be worth a thought. Sometimes decisions we make in the initial design can make future work easier or harder, so I thought I'd throw out a couple of ideas I've been thinking about. I'm not expecting you to take them on or to solve the problems.

[Taragolis]

I've just check, there is public method also exists in Meta object so we do not need read this manually

boto3_method = "create_nodegroup"
session = boto3.session.Session(region_name="eu-west-1")
client = session.client("eks")
meta = client.meta
api_method = meta.method_to_api_mapping[boto3_method]
service_model = meta.service_model
om = service_model.operation_model(api_method)
print(om.input_shape.members)
# OrderedDict([('clusterName', <StringShape(String)>), ('nodegroupName', <StringShape(String)>), ('scalingConfig', <StructureShape(NodegroupScalingConfig)>), ('diskSize', <Shape(BoxedInteger)>), ('subnets', <ListShape(StringList)>), ('instanceTypes', <ListShape(StringList)>), ('amiType', <StringShape(AMITypes)>), ('remoteAccess', <StructureShape(RemoteAccessConfig)>), ('nodeRole', <StringShape(String)>), ('labels', <MapShape(labelsMap)>), ('taints', <ListShape(taintsList)>), ('tags', <MapShape(TagMap)>), ('clientRequestToken', <StringShape(String)>), ('launchTemplate', <StructureShape(LaunchTemplateSpecification)>), ('updateConfig', <StructureShape(NodegroupUpdateConfig)>), ('capacityType', <StringShape(CapacityTypes)>), ('version', <StringShape(String)>), ('releaseVersion', <StringShape(String)>)])

print(om.input_shape.metadata)
# {'required': ['clusterName', 'nodegroupName', 'subnets', 'nodeRole']}

Unfortunetly I can't find any info about Meta into boto3 or botocore documentations. Just a few mentions like this https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#creating-clients

[ferruzzi]

I can't find any info about Meta into boto3 or botocore documentations

I couldn't either, which is why I had initially come up with that hack I shared. :( I don't think I knew about the meta.method_to_api_mapping[], that could be handy. Going through meta directly instead of the boto config file is much nicer overall.

[ferruzzi]

Also, I just wanted to take a sec to say thanks. You have been putting a ton of effort into improving the AWS provider package and it's absolutely appreciated. 👍

[o-nikolas]

Looks really cool @Taragolis, love the idea!

Idea is add legacy_mode attribute and set it implicit to True for all existed hooks.
Property conn and get_conn will provide compatibility with previous version of Hooks for a while

Do we really need the legacy_mode attribute? We could just add conn and get_conn to the new classes so that it's seamless (include deprecation warnings along with it of course) no need for two modes.

According to internals: AWS have 336 different services,10 of them have high-level client and in total we have 31 hooks in Amazon Provider.

Is part of this proposal creating hooks for all services? I'd love to see this and we've had some internal discussions about it as well. Having at least a (autogenerated) hook for all services would be great. And then any custom methods can be added to the hook as operators are written for each service.

[Taragolis]

Do we really need the legacy_mode attribute? We could just add conn and get_conn to the new classes so that it's seamless (include deprecation warnings along with it of course) no need for two modes.

It is possible. The idea with legacy_mode add additional time to migrate to new classes for end users, community and third parties providers with less braking changes.

Is part of this proposal creating hooks for all services?

Oh, no it just add this for some statistic with comparison how many of current hook could (potentially) use high-level clients (resources)

Having at least a (autogenerated) hook for all services would be great. And then any custom methods can be added to the hook as operators are written for each service.

That is a brilliant idea! We could create some script and Jinja templates which check is this hook presented into airflow/providers/amazon/aws/hooks/{service_name_in_snake_case}.py and generate if not and run it manually after new provider release.

[vincbeck]

Love the video! Thanks for the effort. Very minor technical question on my side, should not Boto3ResourceHook inherits from Boto3ClientHook?

[Taragolis]

Initially I create Boto3ResourceHook based on Boto3ClientHook and result was all methods defined in Boto3ClientHook overwrites in Boto3ResourceHook

[vincbeck]

Do you have any example of methods in Boto3ResourceHook and Boto3ClientHook. As far as I understand there are the properties client and resource in these classes but I am not sure what kind of methods you want to expose there

[vincbeck]

But it's definitely minor and implementation details so we can continue this conversation in the future PR ;)

[Taragolis]

Yeah, i mean properties. I don't think any specific methods should be presented into the this classes.

https://gist.github.com/Taragolis/d4ade059c1966a773a4ee7a1d29800d9

[vincbeck]

Oh I see! Very clear now :) Thanks for the explanations!