diff --git a/testing/legal-test-aws/pom.xml b/testing/legal-test-aws/pom.xml
index 8f199b4de457db4d57edd690c0a9fabc045a1e8c..d6b09e3e622de6c48cf210bc6e2b95fe2b0cfda4 100644
--- a/testing/legal-test-aws/pom.xml
+++ b/testing/legal-test-aws/pom.xml
@@ -45,7 +45,7 @@
         <dependency>
             <groupId>org.opengroup.osdu.core.aws</groupId>
             <artifactId>os-core-lib-aws</artifactId>
-            <version>0.3.16</version>
+            <version>0.9.0-SNAPSHOT</version>
         </dependency>
 
         <!-- AWS managed packages -->
diff --git a/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/AwsLegalTagUtils.java b/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/AwsLegalTagUtils.java
index 329990130bcc766195d136c860373afb077578b0..c2fc4846e77286ff9a75365342c058f43f0cf5d4 100644
--- a/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/AwsLegalTagUtils.java
+++ b/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/AwsLegalTagUtils.java
@@ -21,10 +21,10 @@ import java.util.List;
 
 import com.amazonaws.services.s3.AmazonS3;
 import org.opengroup.osdu.core.aws.cognito.AWSCognitoClient;
-import org.opengroup.osdu.core.aws.dynamodb.DynamoDBQueryHelper;
+import org.opengroup.osdu.core.aws.dynamodb.DynamoDBQueryHelperFactory;
+import org.opengroup.osdu.core.aws.dynamodb.DynamoDBQueryHelperV2;
 import org.opengroup.osdu.core.aws.s3.S3Config;
-import org.opengroup.osdu.core.common.model.legal.Properties;
-import org.springframework.beans.factory.annotation.Value;
+
 
 public class AwsLegalTagUtils extends LegalTagUtils {
     private static final String FILE_NAME = "Legal_COO.json";
@@ -77,6 +77,7 @@ public class AwsLegalTagUtils extends LegalTagUtils {
         doc.setDescription("Expired integration test tag");
         doc.setName(integrationTagTestName);
         doc.setId(Integer.toString(integrationTagTestName.hashCode()));
+        doc.setDataPartitionId(getMyDataPartition());
 
         org.opengroup.osdu.core.common.model.legal.Properties properties = new org.opengroup.osdu.core.common.model.legal.Properties();
         List countryOfOrigin = new ArrayList();
@@ -92,11 +93,8 @@ public class AwsLegalTagUtils extends LegalTagUtils {
         properties.setExportClassification("EAR99");
         doc.setProperties(properties);
 
-        String tablePrefix = String.format("%s%s", System.getenv(TABLE_PREFIX), "-");
-        String dynamoDbRegion = System.getenv(DYNAMO_DB_REGION);
-        String dynamoDbEndpoint = System.getenv(DYNAMO_DB_ENDPOINT);
-
-        DynamoDBQueryHelper queryHelper = new DynamoDBQueryHelper(dynamoDbEndpoint, dynamoDbRegion, tablePrefix);
+        DynamoDBQueryHelperFactory dynamoDBQueryHelperFactory = new DynamoDBQueryHelperFactory();
+        DynamoDBQueryHelperV2 queryHelper = dynamoDBQueryHelperFactory.getQueryHelperForPartition(getMyDataPartition(), "legal/legalTable");
 
         // delete legal tag if it exists
         if(queryHelper.keyExistsInTable(LegalDoc.class, doc)){
@@ -105,4 +103,6 @@ public class AwsLegalTagUtils extends LegalTagUtils {
 
         queryHelper.save(doc);
     }
+
+
 }
diff --git a/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/LegalDoc.java b/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/LegalDoc.java
index f8e0c0b1d0a2b6dc8aa9cfca3472ab38fc571b7c..6b20d3dc93dfcefbb1e160e9963090a97be7108e 100644
--- a/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/LegalDoc.java
+++ b/testing/legal-test-aws/src/test/java/org/opengroup/osdu/legal/util/LegalDoc.java
@@ -14,10 +14,7 @@
 
 package org.opengroup.osdu.legal.util;
 
-import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBAttribute;
-import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBHashKey;
-import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBTable;
-import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBTypeConverted;
+import com.amazonaws.services.dynamodbv2.datamodeling.*;
 import org.opengroup.osdu.core.common.model.legal.Properties;
 
 
@@ -27,6 +24,9 @@ public class LegalDoc {
     @DynamoDBHashKey(attributeName = "Id")
     private String id;
 
+    @DynamoDBRangeKey(attributeName = "dataPartitionId")
+    private String dataPartitionId;
+
     @DynamoDBAttribute(attributeName = "Name")
     private String name;
 
@@ -49,6 +49,14 @@ public class LegalDoc {
         this.id = id;
     }
 
+    public String getDataPartitionId(){
+        return dataPartitionId;
+    }
+
+    public void setDataPartitionId(String tenant){
+        this.dataPartitionId = tenant;
+    }
+
     public void setName(String name) {
         this.name = name;
     }