99import software .amazon .smithy .aws .traits .ServiceTrait ;
1010import software .amazon .smithy .model .Model ;
1111import software .amazon .smithy .model .node .Node ;
12+ import software .amazon .smithy .model .shapes .ServiceShape ;
1213import software .amazon .smithy .rulesengine .traits .EndpointRuleSetTrait ;
1314import software .amazon .smithy .typescript .codegen .TypeScriptSettings ;
1415import software .amazon .smithy .typescript .codegen .endpointsV2 .AddDefaultEndpointRuleSet ;
@@ -34,20 +35,21 @@ public List<String> runBefore() {
3435 public Model preprocessModel (Model model , TypeScriptSettings settings ) {
3536 Model .Builder modelBuilder = model .toBuilder ();
3637
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.
4243
43- modelBuilder .removeShape (serviceShape .toShapeId ());
44- modelBuilder .addShape (serviceShape .toBuilder ()
44+ modelBuilder .removeShape (serviceShape .toShapeId ());
45+ modelBuilder .addShape (
46+ serviceShape .toBuilder ()
4547 .addTrait (getDefaultRegionalEndpointRuleSet (
46- serviceShape .expectTrait (ServiceTrait .class ).getEndpointPrefix ())
47- )
48- .build ());
49- }
50- });
48+ serviceShape .expectTrait (ServiceTrait .class ).getEndpointPrefix ()
49+ ))
50+ .build ()
51+ );
52+ }
5153
5254 return modelBuilder .build ();
5355 }
0 commit comments