From f5704cd7c4f8fc7aa087eab61ad8019897d4f30b Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Tue, 24 Sep 2024 13:37:39 -0400
Subject: [PATCH] Add `KclMessageDrivenChannelAdapter.leaseTableName` property

---
 .../KclMessageDrivenChannelAdapter.java       | 20 +++++++++++++++++--
 .../KclMessageDrivenChannelAdapterTests.java  |  4 ++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
index b521860..22bdfba 100644
--- a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
+++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
@@ -23,8 +23,6 @@ import java.util.List;
 import java.util.UUID;
 import java.util.function.Consumer;
 
-import javax.annotation.Nullable;
-
 import com.amazonaws.services.schemaregistry.deserializers.GlueSchemaRegistryDeserializer;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -87,6 +85,7 @@ import org.springframework.integration.support.ErrorMessageStrategy;
 import org.springframework.integration.support.ErrorMessageUtils;
 import org.springframework.integration.support.management.IntegrationManagedResource;
 import org.springframework.jmx.export.annotation.ManagedResource;
+import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.util.Assert;
 
@@ -123,6 +122,9 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 
 	private String consumerGroup = "SpringIntegration";
 
+	@Nullable
+	private String leaseTableName;
+
 	private InboundMessageMapper<byte[]> embeddedHeadersMapper;
 
 	private ConfigsBuilder config;
@@ -211,6 +213,16 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 		return this.consumerGroup;
 	}
 
+	/**
+	 * Set a name of the DynamoDB table name for leases.
+	 * Defaults to {@link #consumerGroup}.
+	 * @param leaseTableName the DynamoDB table name for leases.
+	 * @since 3.0.8
+	 */
+	public void setLeaseTableName(String leaseTableName) {
+		this.leaseTableName = leaseTableName;
+	}
+
 	/**
 	 * Specify an {@link InboundMessageMapper} to extract message headers embedded into
 	 * the record data.
@@ -372,6 +384,10 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 						this.cloudWatchClient,
 						this.workerId,
 						this.recordProcessorFactory);
+
+		if (this.leaseTableName != null) {
+			this.config.tableName(this.leaseTableName);
+		}
 	}
 
 	private StreamTracker buildStreamTracker() {
diff --git a/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java b/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
index b4696a1..9b83387 100644
--- a/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
+++ b/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
@@ -121,6 +121,9 @@ public class KclMessageDrivenChannelAdapterTests implements LocalstackContainerT
 
 		// Because FanOut is false, there would be no Stream Consumers.
 		assertThat(streamConsumers).hasSize(0);
+
+		List<String> tableNames = DYNAMO_DB.listTables().join().tableNames();
+		assertThat(tableNames).containsOnly("test_table");
 	}
 
 	@Test
@@ -172,6 +175,7 @@ public class KclMessageDrivenChannelAdapterTests implements LocalstackContainerT
 					InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON));
 			adapter.setConverter(String::new);
 			adapter.setConsumerGroup("single_stream_group");
+			adapter.setLeaseTableName("test_table");
 			adapter.setFanOut(false);
 			adapter.setMetricsLevel(MetricsLevel.NONE);
 			adapter.setLeaseManagementConfigCustomizer(leaseManagementConfig ->