9
9
import software .amazon .smithy .aws .traits .ServiceTrait ;
10
10
import software .amazon .smithy .model .Model ;
11
11
import software .amazon .smithy .model .node .Node ;
12
+ import software .amazon .smithy .model .shapes .ServiceShape ;
12
13
import software .amazon .smithy .rulesengine .traits .EndpointRuleSetTrait ;
13
14
import software .amazon .smithy .typescript .codegen .TypeScriptSettings ;
14
15
import software .amazon .smithy .typescript .codegen .endpointsV2 .AddDefaultEndpointRuleSet ;
@@ -34,20 +35,21 @@ public List<String> runBefore() {
34
35
public Model preprocessModel (Model model , TypeScriptSettings settings ) {
35
36
Model .Builder modelBuilder = model .toBuilder ();
36
37
37
- model . getServiceShapes (). forEach ( serviceShape -> {
38
- if (!serviceShape .hasTrait (EndpointRuleSetTrait .class )
39
- && AwsTraitsUtils .isAwsService (settings , model )) {
40
- // this branch is for models that identify as AWS services
41
- // but do not include an endpoint ruleset.
38
+ ServiceShape serviceShape = settings . getService ( model );
39
+ if (!serviceShape .hasTrait (EndpointRuleSetTrait .class )
40
+ && AwsTraitsUtils .isAwsService (serviceShape )) {
41
+ // this branch is for models that identify as AWS services
42
+ // but do not include an endpoint ruleset.
42
43
43
- modelBuilder .removeShape (serviceShape .toShapeId ());
44
- modelBuilder .addShape (serviceShape .toBuilder ()
44
+ modelBuilder .removeShape (serviceShape .toShapeId ());
45
+ modelBuilder .addShape (
46
+ serviceShape .toBuilder ()
45
47
.addTrait (getDefaultRegionalEndpointRuleSet (
46
- serviceShape .expectTrait (ServiceTrait .class ).getEndpointPrefix ())
47
- )
48
- .build ());
49
- }
50
- });
48
+ serviceShape .expectTrait (ServiceTrait .class ).getEndpointPrefix ()
49
+ ))
50
+ .build ()
51
+ );
52
+ }
51
53
52
54
return modelBuilder .build ();
53
55
}
0 commit comments