Skip to content

Commit 24a117a

Browse files
xiaoxuandevxiaoxuandev
committed
AWS: Add LakeFormation credential support for GlueCatalog
1 parent 718ff6a commit 24a117a

7 files changed

Lines changed: 401 additions & 4 deletions

File tree

‎aws/src/integration/java/org/apache/iceberg/aws/lakeformation/TestLakeFormationAwsClientFactory.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iceberg.aws.lakeformation;
21+
22+
import java.util.Map;
23+
import java.util.UUID;
24+
import org.apache.iceberg.aws.AwsIntegTestUtil;
25+
import org.apache.iceberg.aws.AwsProperties;
26+
import org.apache.iceberg.aws.glue.GlueCatalog;
27+
import org.apache.iceberg.catalog.Namespace;
28+
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
29+
import org.junit.After;
30+
import org.junit.Assert;
31+
import org.junit.Before;
32+
import org.junit.Test;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
35+
import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
36+
import software.amazon.awssdk.regions.Region;
37+
import software.amazon.awssdk.services.glue.model.AccessDeniedException;
38+
import software.amazon.awssdk.services.glue.model.GlueException;
39+
import software.amazon.awssdk.services.iam.IamClient;
40+
import software.amazon.awssdk.services.iam.model.CreateRoleRequest;
41+
import software.amazon.awssdk.services.iam.model.CreateRoleResponse;
42+
import software.amazon.awssdk.services.iam.model.DeleteRolePolicyRequest;
43+
import software.amazon.awssdk.services.iam.model.DeleteRoleRequest;
44+
import software.amazon.awssdk.services.iam.model.PutRolePolicyRequest;
45+
46+
public class TestLakeFormationAwsClientFactory {
47+
48+
private static final Logger LOG = LoggerFactory.getLogger(TestLakeFormationAwsClientFactory.class);
49+
private static final int IAM_PROPAGATION_DELAY = 10000;
50+
private static final int ASSUME_ROLE_SESSION_DURATION = 3600;
51+
52+
private IamClient iam;
53+
private String roleName;
54+
private Map<String, String> assumeRoleProperties;
55+
private String policyName;
56+
57+
@Before
58+
public void before() {
59+
roleName = UUID.randomUUID().toString();
60+
iam = IamClient.builder()
61+
.region(Region.AWS_GLOBAL)
62+
.httpClientBuilder(UrlConnectionHttpClient.builder())
63+
.build();
64+
CreateRoleResponse response = iam.createRole(CreateRoleRequest.builder()
65+
.roleName(roleName)
66+
.assumeRolePolicyDocument("{" +
67+
"\"Version\":\"2012-10-17\"," +
68+
"\"Statement\":[{" +
69+
"\"Effect\":\"Allow\"," +
70+
"\"Principal\":{" +
71+
"\"AWS\":\"arn:aws:iam::" + AwsIntegTestUtil.testAccountId() + ":root\"}," +
72+
"\"Action\": [\"sts:AssumeRole\"," +
73+
"\"sts:TagSession\"]}]}")
74+
.maxSessionDuration(ASSUME_ROLE_SESSION_DURATION)
75+
.build());
76+
assumeRoleProperties = Maps.newHashMap();
77+
assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_REGION, "us-east-1");
78+
assumeRoleProperties.put(AwsProperties.GLUE_LAKEFORMATION_ENABLED, "true");
79+
assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_ARN, response.role().arn());
80+
assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_TAGS_PREFIX +
81+
LakeFormationAwsClientFactory.LF_AUTHORIZED_CALLER, "emr");
82+
policyName = UUID.randomUUID().toString();
83+
}
84+
85+
@After
86+
public void after() {
87+
iam.deleteRolePolicy(DeleteRolePolicyRequest.builder().roleName(roleName).policyName(policyName).build());
88+
iam.deleteRole(DeleteRoleRequest.builder().roleName(roleName).build());
89+
}
90+
91+
@Test
92+
public void testLakeFormationEnabledGlueCatalog() throws Exception {
93+
String glueArnPrefix = "arn:aws:glue:*:" + AwsIntegTestUtil.testAccountId();
94+
iam.putRolePolicy(PutRolePolicyRequest.builder()
95+
.roleName(roleName)
96+
.policyName(policyName)
97+
.policyDocument("{" +
98+
"\"Version\":\"2012-10-17\"," +
99+
"\"Statement\":[{" +
100+
"\"Sid\":\"policy1\"," +
101+
"\"Effect\":\"Allow\"," +
102+
"\"Action\":[\"glue:CreateDatabase\",\"glue:DeleteDatabase\"," +
103+
"\"glue:Get*\",\"lakeformation:GetDataAccess\"]," +
104+
"\"Resource\":[\"" + glueArnPrefix + ":catalog\"," +
105+
"\"" + glueArnPrefix + ":database/allowed_*\"," +
106+
"\"" + glueArnPrefix + ":table/allowed_*/*\"," +
107+
"\"" + glueArnPrefix + ":userDefinedFunction/allowed_*/*\"]}]}")
108+
.build());
109+
waitForIamConsistency();
110+
111+
GlueCatalog glueCatalog = new GlueCatalog();
112+
assumeRoleProperties.put("warehouse", "s3://path");
113+
glueCatalog.initialize("test", assumeRoleProperties);
114+
Namespace deniedNamespace = Namespace.of("denied_" + UUID.randomUUID().toString().replace("-", ""));
115+
try {
116+
glueCatalog.createNamespace(deniedNamespace);
117+
Assert.fail("Access to Glue should be denied");
118+
} catch (GlueException e) {
119+
Assert.assertEquals(AccessDeniedException.class, e.getClass());
120+
} catch (AssertionError e) {
121+
glueCatalog.dropNamespace(deniedNamespace);
122+
throw e;
123+
}
124+
125+
Namespace allowedNamespace = Namespace.of("allowed_" + UUID.randomUUID().toString().replace("-", ""));
126+
try {
127+
glueCatalog.createNamespace(allowedNamespace);
128+
} catch (GlueException e) {
129+
LOG.error("fail to create Glue database", e);
130+
Assert.fail("create namespace should succeed");
131+
} finally {
132+
glueCatalog.dropNamespace(allowedNamespace);
133+
try {
134+
glueCatalog.close();
135+
} catch (Exception e) {
136+
// swallow exception during closing
137+
LOG.error("Error closing GlueCatalog", e);
138+
}
139+
}
140+
}
141+
142+
private void waitForIamConsistency() throws Exception {
143+
Thread.sleep(IAM_PROPAGATION_DELAY); // sleep to make sure IAM up to date
144+
}
145+
}

‎aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ public void initialize(Map<String, String> properties) {
8686
this.tags = toTags(properties);
8787
}
8888

89-
private <T extends AwsClientBuilder & AwsSyncClientBuilder> T configure(T clientBuilder) {
89+
protected <T extends AwsClientBuilder & AwsSyncClientBuilder> T configure(T clientBuilder) {
9090
AssumeRoleRequest request = AssumeRoleRequest.builder()
9191
.roleArn(roleArn)
9292
.roleSessionName(genSessionName())
@@ -107,6 +107,18 @@ private <T extends AwsClientBuilder & AwsSyncClientBuilder> T configure(T client
107107
return clientBuilder;
108108
}
109109

110+
protected Set<Tag> tags() {
111+
return tags;
112+
}
113+
114+
protected String region() {
115+
return region;
116+
}
117+
118+
protected String s3Endpoint() {
119+
return s3Endpoint;
120+
}
121+
110122
private String genSessionName() {
111123
return String.format("iceberg-aws-%s", UUID.randomUUID());
112124
}

‎aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ public void initialize(Map<String, String> properties) {
121121
}
122122
}
123123

124-
static <T extends SdkClientBuilder> void configureEndpoint(T builder, String endpoint) {
124+
public static <T extends SdkClientBuilder> void configureEndpoint(T builder, String endpoint) {
125125
if (endpoint != null) {
126126
builder.endpointOverride(URI.create(endpoint));
127127
}

‎aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.Set;
2525
import java.util.stream.Collectors;
2626
import org.apache.iceberg.aws.dynamodb.DynamoDbCatalog;
27+
import org.apache.iceberg.aws.lakeformation.LakeFormationAwsClientFactory;
2728
import org.apache.iceberg.aws.s3.S3FileIO;
2829
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
2930
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
@@ -87,6 +88,11 @@ public class AwsProperties implements Serializable {
8788
*/
8889
public static final String GLUE_CATALOG_ID = "glue.id";
8990

91+
/**
92+
* The account ID used in a Glue resource ARN, e.g. arn:aws:glue:us-east-1:1000000000000:table/db1/table1
93+
*/
94+
public static final String GLUE_ACCOUNT_ID = "glue.account-id";
95+
9096
/**
9197
* If Glue should skip archiving an old table version when creating a new version in a commit.
9298
* By default Glue archives all old table versions after an UpdateTable call,
@@ -96,6 +102,15 @@ public class AwsProperties implements Serializable {
96102
public static final String GLUE_CATALOG_SKIP_ARCHIVE = "glue.skip-archive";
97103
public static final boolean GLUE_CATALOG_SKIP_ARCHIVE_DEFAULT = false;
98104

105+
/**
106+
* If set, GlueCatalog will use Lake Formation for access control.
107+
* For more credential vending details, see: https://docs.aws.amazon.com/lake-formation/latest/dg/api-overview.html.
108+
* If enabled, the {@link AwsClientFactory} implementation must be {@link LakeFormationAwsClientFactory}
109+
* or any class that extends it.
110+
*/
111+
public static final String GLUE_LAKEFORMATION_ENABLED = "glue.lakeformation-enabled";
112+
public static final boolean GLUE_LAKEFORMATION_ENABLED_DEFAULT = false;
113+
99114
/**
100115
* Number of threads to use for uploading parts to S3 (shared pool across all output streams),
101116
* default to {@link Runtime#availableProcessors()}
@@ -260,6 +275,18 @@ public class AwsProperties implements Serializable {
260275
@Deprecated
261276
public static final boolean CLIENT_ENABLE_ETAG_CHECK_DEFAULT = false;
262277

278+
/**
279+
* Used by {@link LakeFormationAwsClientFactory}.
280+
* The table name used as part of lake formation credentials request.
281+
*/
282+
public static final String LAKE_FORMATION_TABLE_NAME = "lakeformation.table-name";
283+
284+
/**
285+
* Used by {@link LakeFormationAwsClientFactory}.
286+
* The database name used as part of lake formation credentials request.
287+
*/
288+
public static final String LAKE_FORMATION_DB_NAME = "lakeformation.db-name";
289+
263290
private String s3FileIoSseType;
264291
private String s3FileIoSseKey;
265292
private String s3FileIoSseMd5;
@@ -274,6 +301,7 @@ public class AwsProperties implements Serializable {
274301

275302
private String glueCatalogId;
276303
private boolean glueCatalogSkipArchive;
304+
private boolean glueLakeFormationEnabled;
277305

278306
private String dynamoDbTableName;
279307

@@ -293,6 +321,7 @@ public AwsProperties() {
293321

294322
this.glueCatalogId = null;
295323
this.glueCatalogSkipArchive = GLUE_CATALOG_SKIP_ARCHIVE_DEFAULT;
324+
this.glueLakeFormationEnabled = GLUE_LAKEFORMATION_ENABLED_DEFAULT;
296325

297326
this.dynamoDbTableName = DYNAMODB_TABLE_NAME_DEFAULT;
298327
}
@@ -310,6 +339,9 @@ public AwsProperties(Map<String, String> properties) {
310339
this.glueCatalogId = properties.get(GLUE_CATALOG_ID);
311340
this.glueCatalogSkipArchive = PropertyUtil.propertyAsBoolean(properties,
312341
AwsProperties.GLUE_CATALOG_SKIP_ARCHIVE, AwsProperties.GLUE_CATALOG_SKIP_ARCHIVE_DEFAULT);
342+
this.glueLakeFormationEnabled = PropertyUtil.propertyAsBoolean(properties,
343+
GLUE_LAKEFORMATION_ENABLED,
344+
GLUE_LAKEFORMATION_ENABLED_DEFAULT);
313345

314346
this.s3FileIoMultipartUploadThreads = PropertyUtil.propertyAsInt(properties, S3FILEIO_MULTIPART_UPLOAD_THREADS,
315347
Runtime.getRuntime().availableProcessors());
@@ -402,6 +434,14 @@ public void setGlueCatalogSkipArchive(boolean skipArchive) {
402434
this.glueCatalogSkipArchive = skipArchive;
403435
}
404436

437+
public boolean glueLakeFormationEnabled() {
438+
return glueLakeFormationEnabled;
439+
}
440+
441+
public void setGlueLakeFormationEnabled(boolean glueLakeFormationEnabled) {
442+
this.glueLakeFormationEnabled = glueLakeFormationEnabled;
443+
}
444+
405445
public int s3FileIoMultipartUploadThreads() {
406446
return s3FileIoMultipartUploadThreads;
407447
}

‎aws/src/main/java/org/apache/iceberg/aws/glue/GlueCatalog.java‎

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,9 @@
3434
import org.apache.iceberg.TableMetadata;
3535
import org.apache.iceberg.TableOperations;
3636
import org.apache.iceberg.aws.AwsClientFactories;
37+
import org.apache.iceberg.aws.AwsClientFactory;
3738
import org.apache.iceberg.aws.AwsProperties;
39+
import org.apache.iceberg.aws.lakeformation.LakeFormationAwsClientFactory;
3840
import org.apache.iceberg.aws.s3.S3FileIO;
3941
import org.apache.iceberg.catalog.Namespace;
4042
import org.apache.iceberg.catalog.SupportsNamespaces;
@@ -49,9 +51,11 @@
4951
import org.apache.iceberg.io.FileIO;
5052
import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
5153
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
54+
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
5255
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
5356
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
5457
import org.apache.iceberg.util.LockManagers;
58+
import org.apache.iceberg.util.PropertyUtil;
5559
import org.slf4j.Logger;
5660
import org.slf4j.LoggerFactory;
5761
import software.amazon.awssdk.services.glue.GlueClient;
@@ -87,6 +91,7 @@ public class GlueCatalog extends BaseMetastoreCatalog
8791
private FileIO fileIO;
8892
private LockManager lockManager;
8993
private CloseableGroup closeableGroup;
94+
private Map<String, String> catalogProperties;
9095

9196
// Attempt to set versionId if available on the path
9297
private static final DynMethods.UnboundMethod SET_VERSION_ID = DynMethods.builder("versionId")
@@ -104,13 +109,36 @@ public GlueCatalog() {
104109

105110
@Override
106111
public void initialize(String name, Map<String, String> properties) {
112+
AwsClientFactory awsClientFactory;
113+
FileIO catalogFileIO;
114+
if (PropertyUtil.propertyAsBoolean(
115+
properties,
116+
AwsProperties.GLUE_LAKEFORMATION_ENABLED,
117+
AwsProperties.GLUE_LAKEFORMATION_ENABLED_DEFAULT)) {
118+
String factoryImpl = PropertyUtil.propertyAsString(properties, AwsProperties.CLIENT_FACTORY, null);
119+
ImmutableMap.Builder<String, String> builder = ImmutableMap.<String, String>builder().putAll(properties);
120+
if (factoryImpl == null) {
121+
builder.put(AwsProperties.CLIENT_FACTORY, LakeFormationAwsClientFactory.class.getName());
122+
}
123+
124+
this.catalogProperties = builder.build();
125+
awsClientFactory = AwsClientFactories.from(catalogProperties);
126+
Preconditions.checkArgument(awsClientFactory instanceof LakeFormationAwsClientFactory,
127+
"Detected LakeFormation enabled for Glue catalog, should use a client factory that extends %s, but found %s",
128+
LakeFormationAwsClientFactory.class.getName(), factoryImpl);
129+
catalogFileIO = null;
130+
} else {
131+
awsClientFactory = AwsClientFactories.from(properties);
132+
catalogFileIO = initializeFileIO(properties);
133+
}
134+
107135
initialize(
108136
name,
109137
properties.get(CatalogProperties.WAREHOUSE_LOCATION),
110138
new AwsProperties(properties),
111-
AwsClientFactories.from(properties).glue(),
139+
awsClientFactory.glue(),
112140
initializeLockManager(properties),
113-
initializeFileIO(properties));
141+
catalogFileIO);
114142
}
115143

116144
private LockManager initializeLockManager(Map<String, String> properties) {
@@ -162,6 +190,19 @@ private String cleanWarehousePath(String path) {
162190

163191
@Override
164192
protected TableOperations newTableOps(TableIdentifier tableIdentifier) {
193+
if (catalogProperties != null) {
194+
Map<String, String> tableSpecificCatalogProperties = ImmutableMap.<String, String>builder()
195+
.putAll(catalogProperties)
196+
.put(AwsProperties.LAKE_FORMATION_DB_NAME,
197+
IcebergToGlueConverter.getDatabaseName(tableIdentifier))
198+
.put(AwsProperties.LAKE_FORMATION_TABLE_NAME,
199+
IcebergToGlueConverter.getTableName(tableIdentifier))
200+
.build();
201+
// FileIO initialization depends on tableSpecificCatalogProperties, so a new FileIO is initialized each time
202+
return new GlueTableOperations(glue, lockManager, catalogName, awsProperties,
203+
initializeFileIO(tableSpecificCatalogProperties), tableIdentifier);
204+
}
205+
165206
return new GlueTableOperations(glue, lockManager, catalogName, awsProperties, fileIO, tableIdentifier);
166207
}
167208

0 commit comments

Comments
 (0)