MongoDBUpdateVisitor.java

/*
 * JBoss, Home of Professional Open Source.
 * See the COPYRIGHT.txt file distributed with this work for information
 * regarding copyright ownership.  Some portions may be licensed
 * to Red Hat, Inc. under one or more contributor license agreements.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301 USA.
 */
package org.teiid.translator.mongodb;

import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Stack;

import org.teiid.language.*;
import org.teiid.language.AndOr.Operator;
import org.teiid.language.Update;
import org.teiid.language.visitor.CollectorVisitor;
import org.teiid.metadata.RuntimeMetadata;
import org.teiid.translator.TranslatorException;
import org.teiid.translator.mongodb.MergeDetails.Association;
import org.teiid.translator.mongodb.MongoDBUpdateExecution.RowInfo;

import com.mongodb.*;

public class MongoDBUpdateVisitor extends MongoDBSelectVisitor {

	protected LinkedHashMap<String, Object> columnValues = new LinkedHashMap<String, Object>();
	private DB mongoDB;
	private BasicDBObject pull;
	private Condition condition;
	protected Stack<DBObject> onGoingPullCriteria = new Stack<DBObject>();
	protected TranslatorException pullException;

	public MongoDBUpdateVisitor(MongoDBExecutionFactory executionFactory, RuntimeMetadata metadata, DB mongoDB) {
		super(executionFactory, metadata);
		this.mongoDB = mongoDB;
	}

	@Override
	public void visit(Insert obj) {
        append(obj.getTable());

        List<ColumnReference> columns = obj.getColumns();
        List<Expression> values = ((ExpressionValueSource)obj.getValueSource()).getValues();

		try {
		    IDRef pk = null;
			for (int i = 0; i < columns.size(); i++) {
				String colName = getColumnName(columns.get(i));
				Expression expr = values.get(i);
				Object value = resolveExpressionValue(colName, expr);
				
                if (this.mongoDoc.isPartOfPrimaryKey(colName)) {
                    if (pk == null) {
                        pk = new IDRef();
                    }
                    pk.addColumn(colName, value);
                }
                else {
                    this.columnValues.put(colName, value);    
                }

		        // Update he mongo document to keep track the reference values.
		        this.mongoDoc.updateReferenceColumnValue(obj.getTable().getName(), colName, value);

		        // if this FK column, replace with reference rather than simple key value
		        if (this.mongoDoc.isPartOfForeignKey(colName)) {
		            MergeDetails ref = this.mongoDoc.getFKReference(colName);
		            this.columnValues.put(colName, ref.clone());
		        }				
			}
            if (pk != null) {
                this.columnValues.put("_id", pk.getValue()); //$NON-NLS-1$
            }			
		} catch (TranslatorException e) {
			this.exceptions.add(e);
		}
	}

	private Object resolveExpressionValue(String colName, Expression expr) throws TranslatorException {
		Object value = null;
		if (expr instanceof Literal) {
			value = this.executionFactory.convertToMongoType(((Literal) expr).getValue(), this.mongoDB, colName);
		}
		else if (expr instanceof org.teiid.language.Array) {
		    org.teiid.language.Array contents = (org.teiid.language.Array)expr;
		    List<Expression> arrayExprs = contents.getExpressions();
		    value = new BasicDBList();
		    for (Expression exp:arrayExprs) {
		        if (exp instanceof Literal) {
		            ((BasicDBList)value).add(this.executionFactory.convertToMongoType(((Literal) exp).getValue(), this.mongoDB, colName));
		        }
		        else {
		            this.exceptions.add(new TranslatorException(MongoDBPlugin.Util.gs(MongoDBPlugin.Event.TEIID18001)));
		        }
		    }
		}
		else {
			this.exceptions.add(new TranslatorException(MongoDBPlugin.Util.gs(MongoDBPlugin.Event.TEIID18001)));
		}
		return value;
	}

	@Override
	public void visit(Update obj) {
		this.condition = obj.getWhere();
        append(obj.getTable());

        List<SetClause> changes = obj.getChanges();
        try {
            IDRef pk = null;
			for (SetClause clause:changes) {
				String colName = getColumnName(clause.getSymbol());
				// make sure user not updating the linked keys
				if (this.mongoDoc.isMerged()) {
                    if (this.mongoDoc.getMergeKey().getAssociation() == Association.ONE
                            && this.mongoDoc.isPartOfPrimaryKey(colName)) {
                        throw new TranslatorException(MongoDBPlugin.Event.TEIID18035, MongoDBPlugin.Util.gs(
                                MongoDBPlugin.Event.TEIID18035, colName, obj.getTable().getName()));
                    } else if (this.mongoDoc.getMergeKey().getAssociation() == Association.MANY
                            && this.mongoDoc.isPartOfForeignKey(colName)) {
                        throw new TranslatorException(MongoDBPlugin.Event.TEIID18036, MongoDBPlugin.Util.gs(
                                MongoDBPlugin.Event.TEIID18036, colName, obj.getTable().getName()));
                    }
				}
				Expression expr = clause.getValue();
				Object value = resolveExpressionValue(colName, expr);

	            if (this.mongoDoc.isPartOfPrimaryKey(colName)) {
	                if (pk == null) {
	                    pk = new IDRef();
	                }
	                pk.addColumn(colName, value);
	            }
	            else {
	                this.columnValues.put(colName, value);    
	            }

                // Update the mongo document to keep track the reference values.
                this.mongoDoc.updateReferenceColumnValue(obj.getTable().getName(), colName, value);

                // if this FK column, replace with reference rather than simple key value
                if (this.mongoDoc.isPartOfForeignKey(colName)) {
                    MergeDetails ref = this.mongoDoc.getFKReference(colName);
                    this.columnValues.put(colName, ref.clone());
                }				
			}
            if (pk != null) {
                this.columnValues.put("_id", pk.getValue()); //$NON-NLS-1$
            }
		} catch (TranslatorException e) {
			this.exceptions.add(e);
		}

        append(obj.getWhere());

        if (!this.onGoingExpression.isEmpty()) {
        	this.match = (DBObject)this.onGoingExpression.pop();
        }
	}

	@Override
	public void visit(Delete obj) {
		this.condition = obj.getWhere();
		append(obj.getTable());
        append(obj.getWhere());

        if (!this.onGoingExpression.isEmpty()) {
        	this.match = (DBObject)this.onGoingExpression.pop();
        }
	}

	public BasicDBObject getInsert(LinkedHashMap<String, DBObject> embeddedDocuments) {
		BasicDBObject insert = new BasicDBObject();
		for (String key:this.columnValues.keySet()) {
			Object obj = this.columnValues.get(key);

			if (obj instanceof MergeDetails) {
				obj =  ((MergeDetails)obj).getValue();
			}

			if (key.equals("_id")) { //$NON-NLS-1$
			    insert.append("_id", obj); //$NON-NLS-1$
			}
			if (!this.mongoDoc.isPartOfPrimaryKey(key)) {
    			if (this.mongoDoc.isPartOfForeignKey(key)) {
    				if (obj instanceof BasicDBObject) {
    					insert.append(key, ((BasicDBObject) obj).get(key));
    				}
    				else {
    					insert.append(key, obj);
    				}
    			}
    			else {
    				insert.append(key, obj);
    			}
			}
		}

		if (this.mongoDoc.hasEmbeddedDocuments()) {
			for (String docName:this.mongoDoc.getEmbeddedDocumentNames()) {
				DBObject embedDoc = embeddedDocuments.get(docName);
				if (embedDoc != null) {
					insert.append(docName, embedDoc);
				}
			}
		}
		return insert;
	}
	
	public BasicDBObject getUpdate(LinkedHashMap<String, DBObject> embeddedDocuments) throws TranslatorException {
		BasicDBObject update = new BasicDBObject();

		for (String key:this.columnValues.keySet()) {
			Object obj = this.columnValues.get(key);

			if (obj instanceof MergeDetails) {
				MergeDetails ref = ((MergeDetails)obj);

				if (this.mongoDoc.isMerged()) {
					// do not allow updating the main document reference where this embedded document is embedded.
					if (ref.getParentTable().equals(this.mongoDoc.getMergeTable().getName())) {
						throw new TranslatorException(MongoDBPlugin.Util.gs(MongoDBPlugin.Event.TEIID18007, 
						        ref.getParentTable(), this.mongoDoc.getDocumentName()));
					}
				}

				//update.append(key, ref.getDBRef(db, true));
				if (this.mongoDoc.isPartOfForeignKey(key)) {
					if (ref.getValue() instanceof BasicDBObject) {
						update.append(key, ((BasicDBObject) ref.getValue()).get(key));
					}
					else {
						update.append(key, ref.getValue());
					}
				}
				else {
					update.append(key, ref.getValue());
				}

				// also update the embedded document
				if (this.mongoDoc.hasEmbeddedDocuments()) {
					for (MergeDetails docKey: this.mongoDoc.getEmbeddedReferences()) {
						if (ref.getParentTable().equals(docKey.getEmbeddedTable())) {
							DBObject embedDoc = embeddedDocuments.get(docKey.getName());
							if (embedDoc == null || ref.getValue() == null) {
								update.append(docKey.getName(), null);
							}
							else {
								update.append(docKey.getName(), embedDoc);
							}
						}
					}
				}
			}
			else {
				if (this.mongoDoc.isMerged()) {
					if (this.mongoDoc.getMergeAssociation() == Association.MANY) {
						update.append(this.mongoDoc.getDocumentName()+".$."+key, obj); //$NON-NLS-1$
					}
					else {
						update.append(this.mongoDoc.getDocumentName()+"."+key, obj); //$NON-NLS-1$
					}
				}
				else {
					if (this.mongoDoc.isPartOfPrimaryKey(key)) {
						if (hasCompositePrimaryKey(this.mongoDoc.getTargetTable())) {
							update.append("_id."+key, obj);//$NON-NLS-1$
						}
						else {
							update.append("_id", obj); //$NON-NLS-1$
						}
					}
					else {
						update.append(key, obj);
					}
				}
			}
		}
		return update;
	}

	public BasicDBObject getPullQuery() throws TranslatorException {
	    if (this.pullException != null) {
	        throw this.pullException;
	    }
		if (this.pull == null) {
		    if (this.onGoingPullCriteria.isEmpty()) {
		        this.pull = new BasicDBObject();
		    }
		    else {
		        this.pull =  new BasicDBObject(this.mongoDoc.getTable().getName(), this.onGoingPullCriteria.pop());
		    }
		}
		return this.pull;
	}

	public boolean updateMerge(BasicDBList previousRows, RowInfo parentKey, BasicDBList updated) throws TranslatorException {
	    boolean update = false;
		for (int i = 0; i < previousRows.size(); i++) {
			BasicDBObject row = (BasicDBObject)previousRows.get(i);
			if (this.match == null && getPullQuery() == null || ExpressionEvaluator.matches(this.executionFactory, this.mongoDB, this.condition, row, parentKey)) {
			    update = true;
				for (String key:this.columnValues.keySet()) {
					Object obj = this.columnValues.get(key);

					if (obj instanceof MergeDetails) {
						MergeDetails ref = ((MergeDetails)obj);
						row.put(key, ref.getValue());
					}
					else {
						row.put(key, obj);
					}
				}
			}
			updated.add(row);
		}
		return update;
	}
	
    public boolean updateDelete(BasicDBList previousRows, RowInfo parentKey, BasicDBList updated) throws TranslatorException {
        for (int i = 0; i < previousRows.size(); i++) {
            BasicDBObject row = (BasicDBObject)previousRows.get(i);
            if (this.match == null && getPullQuery() == null
                    || ExpressionEvaluator.matches(this.executionFactory, this.mongoDB, this.condition, row, parentKey)) {
                //do not add
            }
            else {
                updated.add(row);    
            }
        }
        return updated.size() != previousRows.size();
    }	
	
	public boolean updateMerge(BasicDBObject previousRow, RowInfo parentKey) throws TranslatorException {
	    boolean update = false;
		if (this.match == null || ExpressionEvaluator.matches(this.executionFactory, this.mongoDB, this.condition, previousRow, parentKey)) {
			for (String key:this.columnValues.keySet()) {
				Object obj = this.columnValues.get(key);
				
				update = true;
				
				if (obj instanceof MergeDetails) {
					MergeDetails ref = ((MergeDetails)obj);
					previousRow.put(key, ref.getValue());
				}
				else {
					previousRow.put(key, obj);
				}
			}
		}
		return update;
	}
	
    @Override
    public void visit(Comparison obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }

        try {
            // this for the normal where clause
            ColumnDetail leftExpr = getExpressionAlias(obj.getLeftExpression());

            append(obj.getRightExpression());

            Object rightExpr = this.onGoingExpression.pop();
            if (this.expressionMap.get(rightExpr) != null) {
                rightExpr = this.expressionMap.get(rightExpr).getProjectedName();
            }
            // build pull criteria for delete; the pull criteria only applies in merge scenario
            // and only columns in the embedded document.
            boolean buildPullQuery = (includeInPullCriteria(obj.getLeftExpression()) && includeInPullCriteria(obj.getRightExpression()));

            if (!buildPullQuery) {
                QueryBuilder query = leftExpr.getQueryBuilder();
                buildComparisionQuery(obj, rightExpr, query);
                this.onGoingExpression.push(query.get());
            }
            else {
                QueryBuilder pullQuery = leftExpr.getPullQueryBuilder();
                buildComparisionQuery(obj, rightExpr, pullQuery);
                this.onGoingPullCriteria.push(pullQuery.get());
            }

            if (obj.getLeftExpression() instanceof ColumnReference) {
                ColumnReference column = (ColumnReference)obj.getLeftExpression();
                this.mongoDoc.updateReferenceColumnValue(column.getTable().getName(), column.getName(), rightExpr);
            }
        } catch (TranslatorException e) {
            this.exceptions.add(e);
        }
    }

    private boolean includeInPullCriteria(Expression expr) throws TranslatorException {
        if (!this.mongoDoc.isMerged()) {
            return false;
        }
        Collection<ColumnReference> columns = CollectorVisitor.collectElements(expr);
        for (ColumnReference column:columns) {
            if (this.mongoDoc.isPartOfForeignKey(column.getName())) {
                return false;
            }
        }
        return true;
    }

    @Override
    public void visit(AndOr obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }

        append(obj.getLeftCondition());
        append(obj.getRightCondition());

        boolean valid = false;
        if (this.onGoingExpression.size() >= 2) {
            DBObject right = (DBObject)this.onGoingExpression.pop();
            DBObject left = (DBObject) this.onGoingExpression.pop();

            switch(obj.getOperator()) {
            case AND:
                this.onGoingExpression.push(QueryBuilder.start().and(left, right).get());
                break;
            case OR:
                this.onGoingExpression.push(QueryBuilder.start().or(left, right).get());
                break;
            }
            valid = true;
        }

        if (this.onGoingPullCriteria.size() >= 2) {
            DBObject pullRight = this.onGoingPullCriteria.pop();
            DBObject pullLeft = this.onGoingPullCriteria.pop();
            switch(obj.getOperator()) {
            case AND:
                this.onGoingPullCriteria.push(QueryBuilder.start().and(pullLeft, pullRight).get());
                break;
            case OR:
                this.onGoingPullCriteria.push(QueryBuilder.start().or(pullLeft, pullRight).get());
                break;
            }
            valid = true;
        }
        if (!valid && obj.getOperator() == Operator.OR) {
            this.pullException = new TranslatorException(MongoDBPlugin.Event.TEIID18029, MongoDBPlugin.Util.gs(MongoDBPlugin.Event.TEIID18029));
        }
    }

    @Override
    public void visit(Function obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }
        this.pullException = new TranslatorException(MongoDBPlugin.Event.TEIID18028, MongoDBPlugin.Util.gs(MongoDBPlugin.Event.TEIID18028));
    }

    @Override
    public void visit(In obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }
        try {
            boolean buildPullQuery = includeInPullCriteria(obj.getLeftExpression());
            if (buildPullQuery) {
                ColumnDetail exprAlias = getExpressionAlias(obj.getLeftExpression());
                this.onGoingPullCriteria.push(buildInQuery(obj, exprAlias.getPullQueryBuilder()).get());
            }
            else {
                ColumnDetail exprAlias = getExpressionAlias(obj.getLeftExpression());
                this.onGoingExpression.push(buildInQuery(obj, exprAlias.getQueryBuilder()).get());
            }
        } catch (TranslatorException e) {
            this.exceptions.add(e);
        }
    }

    @Override
    public void visit(IsNull obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }

        try {
            boolean buildPullQuery = includeInPullCriteria(obj.getExpression());
            if (buildPullQuery) {
                ColumnDetail exprAlias = getExpressionAlias(obj.getExpression());
                this.onGoingPullCriteria.push(buildIsNullQuery(obj, exprAlias.getPullQueryBuilder()).get());
            }
            else {
                ColumnDetail exprAlias = getExpressionAlias(obj.getExpression());
                this.onGoingExpression.push(buildIsNullQuery(obj, exprAlias.getQueryBuilder()).get());
            }
        } catch (TranslatorException e) {
            this.exceptions.add(e);
        }
    }

    @Override
    public void visit(Like obj) {
        if (!this.mongoDoc.isMerged() || this.mongoDoc.isMerged() && this.mongoDoc.getMergeAssociation() != Association.MANY) {
            super.visit(obj);
            return;
        }

        try {
            boolean buildPullQuery = includeInPullCriteria(obj.getLeftExpression());
            if (buildPullQuery) {
                ColumnDetail exprAlias = getExpressionAlias(obj.getLeftExpression());
                this.onGoingPullCriteria.push(buildLikeQuery(obj, exprAlias.getPullQueryBuilder()).get());
            }
            else {
                ColumnDetail exprAlias = getExpressionAlias(obj.getLeftExpression());
                this.onGoingExpression.push(buildLikeQuery(obj, exprAlias.getQueryBuilder()).get());
            }
        } catch (TranslatorException e) {
            this.exceptions.add(e);
        }
    }
}