SK-2679: Add typed context accessors for Conditional Data Access (Map support)

README.md

@@ -2817,7 +2817,7 @@ ctx.put("department", "finance");
 credentials.setContext(ctx);
 ```
 
-> **Note:** `getContext()` returns `Object` — callers should use `instanceof` if they need to inspect the type.
+> **Note:** `getContext()` returns the context as a `String` (or `null` if a Map was set). Use `getContextAsMap()` to retrieve a Map context, or `getContextAsObject()` to retrieve either as the underlying `Object`.
 
 Context map keys must contain only alphanumeric characters and underscores (`[a-zA-Z0-9_]`). Invalid keys will throw a `SkyflowException`.
 

src/main/java/com/skyflow/config/Credentials.java

@@ -33,10 +33,19 @@ public void setRoles(ArrayList<String> roles) {
         this.roles = roles;
     }
 
-    public Object getContext() {
+    public String getContext() {
+        return context instanceof String ? (String) context : null;
+    }
+
+    public Object getContextAsObject() {
         return context;
     }
 
+    @SuppressWarnings("unchecked")
+    public Map<String, Object> getContextAsMap() {
+        return context instanceof Map ? (Map<String, Object>) context : null;
+    }
+
     public void setContext(String context) {
         this.context = context;
     }

src/main/java/com/skyflow/serviceaccount/util/BearerToken.java

@@ -17,6 +17,8 @@
 import com.skyflow.utils.logger.LogUtil;
 import io.jsonwebtoken.Jwts;
 
+import com.skyflow.utils.validations.Validations;
+
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.FileReader;
@@ -220,7 +222,10 @@ public BearerTokenBuilder setCtx(String ctx) {
             return this;
         }
 
-        public BearerTokenBuilder setCtx(Map<String, Object> ctx) {
+        public BearerTokenBuilder setCtx(Map<String, Object> ctx) throws SkyflowException {
+            if (ctx != null && !ctx.isEmpty()) {
+                Validations.validateCtxMapKeys(ctx);
+            }
             this.ctx = ctx;
             return this;
         }

src/main/java/com/skyflow/serviceaccount/util/SignedDataTokens.java

@@ -13,6 +13,8 @@
 import com.skyflow.utils.logger.LogUtil;
 import io.jsonwebtoken.Jwts;
 
+import com.skyflow.utils.validations.Validations;
+
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.FileReader;
@@ -199,7 +201,10 @@ public SignedDataTokensBuilder setCtx(String ctx) {
             return this;
         }
 
-        public SignedDataTokensBuilder setCtx(Map<String, Object> ctx) {
+        public SignedDataTokensBuilder setCtx(Map<String, Object> ctx) throws SkyflowException {
+            if (ctx != null && !ctx.isEmpty()) {
+                Validations.validateCtxMapKeys(ctx);
+            }
             this.ctx = ctx;
             return this;
         }

src/main/java/com/skyflow/utils/Utils.java

@@ -47,28 +47,27 @@ public static String getVaultURL(String clusterId, Env env) {
         return sb.toString();
     }
 
-    @SuppressWarnings("unchecked")
     public static String generateBearerToken(Credentials credentials) throws SkyflowException {
         if (credentials.getPath() != null) {
             BearerToken.BearerTokenBuilder builder = BearerToken.builder()
                     .setCredentials(new File(credentials.getPath()))
                     .setRoles(credentials.getRoles());
-            Object ctx = credentials.getContext();
-            if (ctx instanceof String) {
-                builder.setCtx((String) ctx);
-            } else if (ctx instanceof Map) {
-                builder.setCtx((Map<String, Object>) ctx);
+            Map<String, Object> ctxMap = credentials.getContextAsMap();
+            if (ctxMap != null) {
+                builder.setCtx(ctxMap);
+            } else {
+                builder.setCtx(credentials.getContext());
             }
             return builder.build().getBearerToken();
         } else if (credentials.getCredentialsString() != null) {
             BearerToken.BearerTokenBuilder builder = BearerToken.builder()
                     .setCredentials(credentials.getCredentialsString())
                     .setRoles(credentials.getRoles());
-            Object ctx = credentials.getContext();
-            if (ctx instanceof String) {
-                builder.setCtx((String) ctx);
-            } else if (ctx instanceof Map) {
-                builder.setCtx((Map<String, Object>) ctx);
+            Map<String, Object> ctxMap = credentials.getContextAsMap();
+            if (ctxMap != null) {
+                builder.setCtx(ctxMap);
+            } else {
+                builder.setCtx(credentials.getContext());
             }
             return builder.build().getBearerToken();
         } else {

src/main/java/com/skyflow/utils/validations/Validations.java

@@ -45,6 +45,19 @@ public class Validations {
     private Validations() {
     }
 
+    public static void validateCtxMapKeys(Map<?, ?> ctxMap) throws SkyflowException {
+        Pattern ctxKeyPattern = Pattern.compile(Constants.CONTEXT_KEY_REGEX);
+        for (Object key : ctxMap.keySet()) {
+            if (key == null || !ctxKeyPattern.matcher(key.toString()).matches()) {
+                String keyStr = key == null ? "null" : key.toString();
+                LogUtil.printErrorLog(Utils.parameterizedString(
+                        ErrorLogs.INVALID_CONTEXT_MAP_KEY.getLog(), keyStr));
+                throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(),
+                        Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), keyStr));
+            }
+        }
+    }
+
     public static void validateVaultConfig(VaultConfig vaultConfig) throws SkyflowException {
         String vaultId = vaultConfig.getVaultId();
         String clusterId = vaultConfig.getClusterId();
@@ -162,7 +175,7 @@ public static void validateCredentials(Credentials credentials) throws SkyflowEx
         String credentialsString = credentials.getCredentialsString();
         String token = credentials.getToken();
         String apiKey = credentials.getApiKey();
-        Object context = credentials.getContext();
+        Object context = credentials.getContextAsObject();
         ArrayList<String> roles = credentials.getRoles();
 
         if (path != null) nonNullMembers++;

src/test/java/com/skyflow/config/CredentialsTests.java

@@ -281,6 +281,9 @@ public void testValidMapContextInCredentials() {
             ctxMap.put("user_id", "user_12345");
             credentials.setContext(ctxMap);
             Validations.validateCredentials(credentials);
+            Assert.assertNull(credentials.getContext());
+            Assert.assertEquals(ctxMap, credentials.getContextAsObject());
+            Assert.assertEquals(ctxMap, credentials.getContextAsMap());
         } catch (SkyflowException e) {
             Assert.fail(INVALID_EXCEPTION_THROWN);
         }

src/test/java/com/skyflow/serviceaccount/util/BearerTokenTests.java

@@ -249,4 +249,79 @@ public void testInvalidTokenURIInCredentialsForCredentials() throws SkyflowExcep
             Assert.assertEquals(ErrorMessage.InvalidTokenUri.getMessage(), e.getMessage());
         }
     }
+
+    @Test
+    public void testBearerTokenBuilderWithNestedMapContext() {
+        try {
+            File file = new File(credentialsFilePath);
+            Map<String, Object> nestedContext = new HashMap<>();
+            Map<String, Object> user = new HashMap<>();
+            user.put("role", "admin");
+            user.put("level", 5);
+            nestedContext.put("user", user);
+            nestedContext.put("project_id", "proj_123");
+            BearerToken.builder().setCredentials(file).setCtx(nestedContext).build();
+        } catch (Exception e) {
+            Assert.fail(INVALID_EXCEPTION_THROWN);
+        }
+    }
+
+    @Test
+    public void testBearerTokenBuilderWithMapContextAndCredentialsString() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("role", "admin");
+            BearerToken.builder().setCredentials(credentialsString).setCtx(mapContext).build();
+        } catch (Exception e) {
+            Assert.fail(INVALID_EXCEPTION_THROWN);
+        }
+    }
+
+    @Test
+    public void testBearerTokenWithInvalidCtxMapKeyContainingHyphen() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid-key", "value");
+            BearerToken.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid-key"),
+                    e.getMessage()
+            );
+        }
+    }
+
+    @Test
+    public void testBearerTokenWithInvalidCtxMapKeyContainingDot() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid.key", "value");
+            BearerToken.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid.key"),
+                    e.getMessage()
+            );
+        }
+    }
+
+    @Test
+    public void testBearerTokenWithInvalidCtxMapKeyContainingSpace() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid key", "value");
+            BearerToken.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid key"),
+                    e.getMessage()
+            );
+        }
+    }
 }

src/test/java/com/skyflow/serviceaccount/util/SignedDataTokensTests.java

@@ -238,4 +238,83 @@ public void testSignedDataTokenResponse() {
             Assert.fail(INVALID_EXCEPTION_THROWN);
         }
     }
+
+    @Test
+    public void testSignedDataTokensBuilderWithNestedMapContext() {
+        try {
+            File file = new File(credentialsFilePath);
+            Map<String, Object> nestedContext = new HashMap<>();
+            Map<String, Object> user = new HashMap<>();
+            user.put("role", "admin");
+            user.put("level", 5);
+            nestedContext.put("user", user);
+            nestedContext.put("project_id", "proj_123");
+            SignedDataTokens.builder()
+                    .setCredentials(file).setCtx(nestedContext).setDataTokens(dataTokens).setTimeToLive(ttl)
+                    .build();
+        } catch (Exception e) {
+            Assert.fail(INVALID_EXCEPTION_THROWN);
+        }
+    }
+
+    @Test
+    public void testSignedDataTokensBuilderWithMapContextAndCredentialsString() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("role", "admin");
+            SignedDataTokens.builder()
+                    .setCredentials(credentialsString).setCtx(mapContext).setDataTokens(dataTokens).setTimeToLive(ttl)
+                    .build();
+        } catch (Exception e) {
+            Assert.fail(INVALID_EXCEPTION_THROWN);
+        }
+    }
+
+    @Test
+    public void testSignedDataTokensWithInvalidCtxMapKeyContainingHyphen() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid-key", "value");
+            SignedDataTokens.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid-key"),
+                    e.getMessage()
+            );
+        }
+    }
+
+    @Test
+    public void testSignedDataTokensWithInvalidCtxMapKeyContainingDot() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid.key", "value");
+            SignedDataTokens.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid.key"),
+                    e.getMessage()
+            );
+        }
+    }
+
+    @Test
+    public void testSignedDataTokensWithInvalidCtxMapKeyContainingSpace() {
+        try {
+            Map<String, Object> mapContext = new HashMap<>();
+            mapContext.put("invalid key", "value");
+            SignedDataTokens.builder().setCtx(mapContext);
+            Assert.fail(EXCEPTION_NOT_THROWN);
+        } catch (SkyflowException e) {
+            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
+            Assert.assertEquals(
+                    Utils.parameterizedString(ErrorMessage.InvalidContextMapKey.getMessage(), "invalid key"),
+                    e.getMessage()
+            );
+        }
+    }
 }