TEIID-2570 ensuring blocked lookups consider the keyvalue as well.

teiid · Jul 5, 2013 · 79b8fa4 · 79b8fa4
1 parent 7de4543
commit 79b8fa4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 9 deletions.
diff --git a/engine/src/main/java/org/teiid/query/tempdata/TempTableDataManager.java b/engine/src/main/java/org/teiid/query/tempdata/TempTableDataManager.java
@@ -685,7 +685,7 @@ public Object lookupCodeValue(CommandContext context, String codeTableName,
 		returnElementName = returnElementName.toUpperCase();
     	String matTableName = CODE_PREFIX + codeTableName + ElementSymbol.SEPARATOR + keyElementName + ElementSymbol.SEPARATOR + returnElementName;
 
-    	TupleSource ts = context.getCodeLookup(matTableName);
+    	TupleSource ts = context.getCodeLookup(matTableName, keyValue);
     	if (ts == null) {
 	    	QueryMetadataInterface metadata = context.getMetadata();
 
@@ -711,7 +711,7 @@ public Object lookupCodeValue(CommandContext context, String codeTableName,
 	    	ts.closeSource();
 	    	return result;
     	} catch (BlockedException e) {
-    		context.putCodeLookup(matTableName, ts);
+    		context.putCodeLookup(matTableName, keyValue, ts);
     		throw e;
     	}
     }

diff --git a/engine/src/main/java/org/teiid/query/util/CommandContext.java b/engine/src/main/java/org/teiid/query/util/CommandContext.java
@@ -96,6 +96,25 @@ private static class VDBState {
 		private DQPWorkContext dqpWorkContext;
 	}
 
+	private static class LookupKey implements Comparable<LookupKey> {
+		String matTableName;
+		Comparable keyValue;
+
+		public LookupKey(String matTableName, Object keyValue) {
+			this.matTableName = matTableName;
+			this.keyValue = (Comparable) keyValue;
+		}
+
+		@Override
+		public int compareTo(LookupKey arg0) {
+			int comp = matTableName.compareTo(arg0.matTableName);
+			if (comp != 0) {
+				return comp;
+			}
+			return keyValue.compareTo(arg0.keyValue);
+		}
+	}
+
 	private static class GlobalState implements Cloneable {
 	    private WeakReference<RequestWorkItem> processorID;
 
@@ -157,7 +176,7 @@ private static class GlobalState implements Cloneable {
 
 		private AuthorizationValidator authorizationValidator;
 
-		private Map<String, TupleSource> lookups;
+		private Map<LookupKey, TupleSource> lookups;
 	}
 
 	private GlobalState globalState = new GlobalState();
@@ -933,19 +952,19 @@ public Object getSessionVariable(String key) {
 	public AuthorizationValidator getAuthorizationValidator() {
 		return this.globalState.authorizationValidator;
 	}
-
-	public TupleSource getCodeLookup(String matTableName) {
+	
+	public TupleSource getCodeLookup(String matTableName, Object keyValue) {
 		if (this.globalState.lookups != null) {
-			return this.globalState.lookups.remove(matTableName);
+			return this.globalState.lookups.remove(new LookupKey(matTableName, keyValue));
 		}
 		return null;
 	}
 
-	public void putCodeLookup(String matTableName, TupleSource ts) {
+	public void putCodeLookup(String matTableName, Object keyValue, TupleSource ts) {
 		if (this.globalState.lookups == null) {
-			this.globalState.lookups = new HashMap<String, TupleSource>();
+			this.globalState.lookups = new TreeMap<LookupKey, TupleSource>();
 		}
-		this.globalState.lookups.put(matTableName, ts);
+		this.globalState.lookups.put(new LookupKey(matTableName, keyValue), ts);
 	}