cherry pick otf1725

Author: Jennifer Baldwin

aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java

@@ -82,8 +82,7 @@ public S3Client s3() {
           .applyMutation(s3FileIOProperties()::applyEndpointConfigurations)
           .applyMutation(s3FileIOProperties()::applyServiceConfigurations)
           .applyMutation(s3FileIOProperties()::applyRetryConfigurations)
-          .credentialsProvider(
-              new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn()))
+          .credentialsProvider(lakeFormationCredentialsProvider())
           .region(Region.of(region()))
           .build();
     } else {
@@ -97,16 +96,15 @@ public KmsClient kms() {
       return KmsClient.builder()
           .applyMutation(httpClientProperties()::applyHttpClientConfigurations)
           .applyMutation(awsClientProperties()::applyRetryConfigurations)
-          .credentialsProvider(
-              new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn()))
+          .credentialsProvider(lakeFormationCredentialsProvider())
           .region(Region.of(region()))
           .build();
     } else {
       return super.kms();
     }
   }
 
-  private boolean isTableRegisteredWithLakeFormation() {
+  protected boolean isTableRegisteredWithLakeFormation() {
     Preconditions.checkArgument(
         dbName != null && !dbName.isEmpty(), "Database name can not be empty");
     Preconditions.checkArgument(
@@ -120,10 +118,17 @@ private boolean isTableRegisteredWithLakeFormation() {
                     .databaseName(dbName)
                     .name(tableName)
                     .build());
-    return response.table().isRegisteredWithLakeFormation();
+    boolean isRegisteredWithLakeFormation = response.table().isRegisteredWithLakeFormation();
+    if (isRegisteredWithLakeFormation) {
+      // set the Glue account id, if not set
+      if (this.glueAccountId == null) {
+        this.glueAccountId = response.table().catalogId();
+      }
+    }
+    return isRegisteredWithLakeFormation;
   }
 
-  private String buildTableArn() {
+  protected String buildTableArn() {
     Preconditions.checkArgument(
         glueAccountId != null && !glueAccountId.isEmpty(),
         "%s can not be empty",
@@ -133,13 +138,17 @@ private String buildTableArn() {
         "arn:%s:glue:%s:%s:table/%s/%s", partitionName, region(), glueAccountId, dbName, tableName);
   }
 
-  private LakeFormationClient lakeFormation() {
+  protected LakeFormationClient lakeFormation() {
     return LakeFormationClient.builder()
         .applyMutation(this::applyAssumeRoleConfigurations)
         .applyMutation(httpClientProperties()::applyHttpClientConfigurations)
         .build();
   }
 
+  protected AwsCredentialsProvider lakeFormationCredentialsProvider() {
+    return new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn());
+  }
+
   static class LakeFormationCredentialsProvider implements AwsCredentialsProvider {
     private final LakeFormationClient client;
     private final String tableArn;