DataContext template using IQueryable

DataContext.tt

<#@ template language="C#" debug="false" hostspecific="true"#>
<#@ include file="EF.Utility.CS.ttinclude"#>
<#@ output extension=".cs"#>
<#
var loader = new MetadataLoader(this);
var region = new CodeRegion(this);
var inputFile = @"TradingPartner.edmx";
var itemCollection = loader.CreateEdmItemCollection(inputFile);
var container = itemCollection.GetItems<EntityContainer>().FirstOrDefault();
var fileManager = EntityFrameworkTemplateFileManager.Create(this);

Code = new CodeGenerationTools(this);
Tools = new MetadataTools(this);
ObjectNamespace = Code.VsNamespaceSuggestion();
ModelNamespace = loader.GetModelNamespace(inputFile);

WriteHeader(fileManager);

// Create interface
fileManager.StartNewFile("I" + Code.Escape(container) + ".cs");
BeginNamespace();
#>
using System;
using System.Collections.Generic;
using System.Linq;
<#
    if (ObjectNamespace != ModelNamespace)
    {
#>
using <#=ModelNamespace#>;
<#
    }
#>

<#=Accessibility.ForType(container)#> interface I<#=Code.Escape(container)#> : IDisposable
{
    void Add<TEntity>(TEntity entity) where TEntity : class;

    void Add<TEntity>(IEnumerable<TEntity> entities) where TEntity : class;

    void Remove<TEntity>(TEntity entity) where TEntity : class;

    void Remove<TEntity>(IEnumerable<TEntity> entities) where TEntity : class;

    void Update<TEntity>(TEntity entity) where TEntity : class;
<#
foreach (var entitySet in container.BaseEntitySets.OfType<EntitySet>())
{
#>

    IQueryable<<#=Code.Escape(entitySet.ElementType)#>> <#=Code.Escape(entitySet)#> { get; }
<#
}

foreach (var edmFunction in container.FunctionImports)
{
    WriteFunctionDefinition(edmFunction, false);
}
#>

    int SaveChanges();
}
<#
EndNamespace();

// Create implementation
fileManager.StartNewFile(Code.Escape(container) + ".cs");
BeginNamespace();
#>
using System;
using System.Collections.Generic;
using System.Data.Entity;
<#
    if (container.FunctionImports.Any())
    {
#>
using System.Data.Entity.Core.Objects;
<#
    }
#>
using System.Data.Entity.Infrastructure;
using System.Linq;
<#
    if (ObjectNamespace != ModelNamespace)
    {
#>
using <#=ModelNamespace#>;
<#
    }
#>

<#=Accessibility.ForType(container)#> partial class <#=Code.Escape(container)#> : DbContext, I<#=Code.Escape(container)#>
{
    public <#=Code.Escape(container)#>(string nameOrConnectionString)
        : base(nameOrConnectionString)
    {
<#
        WriteLazyLoadingEnabled(container);
#>
    }

    protected override void OnModelCreating(DbModelBuilder modelBuilder)
    {
        throw new UnintentionalCodeFirstException();
    }

    public void Add<TEntity>(TEntity entity) where TEntity : class
    {
        this.Set<TEntity>().Add(entity);
    }

    public void Add<TEntity>(IEnumerable<TEntity> entities) where TEntity : class
    {
        this.Set<TEntity>().AddRange(entities);
    }

    public void Remove<TEntity>(TEntity entity) where TEntity : class
    {
        this.Set<TEntity>().Remove(entity);
    }

    public void Remove<TEntity>(IEnumerable<TEntity> entities) where TEntity : class
    {
        this.Set<TEntity>().RemoveRange(entities); ;
    }

    public void Update<TEntity>(TEntity entity) where TEntity : class
    {
        this.Entry(entity).State = EntityState.Modified;
    }
<#
foreach (var entitySet in container.BaseEntitySets.OfType<EntitySet>())
{
#>

    public virtual IQueryable<<#=Code.Escape(entitySet.ElementType)#>> <#=Code.Escape(entitySet)#> 
    {
        get { return this.Set<<#=Code.Escape(entitySet.ElementType)#>>(); } 
    }
<#
}

foreach (var edmFunction in container.FunctionImports)
{
    WriteFunctionRepresentation(edmFunction, false);
}
#>
}
<#
EndNamespace();

// Write output
fileManager.Process();
#>
<#+
string ModelNamespace { get; set; }
string ObjectNamespace { get; set; }
CodeGenerationTools Code { get; set; }
MetadataTools Tools { get; set; }

string GetResourceString(string resourceName)
{
	if(_resourceManager == null)
	{
		_resourceManager = new System.Resources.ResourceManager("System.Data.Entity.Design", typeof(System.Data.Entity.Design.MetadataItemCollectionFactory).Assembly);
	}
	
    return _resourceManager.GetString(resourceName, null);
}

System.Resources.ResourceManager _resourceManager;

void WriteLazyLoadingEnabled(EntityContainer container)
{
   string lazyLoadingAttributeValue = null;
   var lazyLoadingAttributeName = MetadataConstants.EDM_ANNOTATION_09_02 + ":LazyLoadingEnabled";
   if(MetadataTools.TryGetStringMetadataPropertySetting(container, lazyLoadingAttributeName, out lazyLoadingAttributeValue))
   {
       bool isLazyLoading;
       if(bool.TryParse(lazyLoadingAttributeValue, out isLazyLoading) && !isLazyLoading)
       {
#>
        this.Configuration.LazyLoadingEnabled = false;
<#+
       }
   }
}

void WriteHeader(EntityFrameworkTemplateFileManager fileManager)
{
    fileManager.StartHeader();
#>
//------------------------------------------------------------------------------
// <auto-generated>
// <#=GetResourceString("Template_GeneratedCodeCommentLine1")#>
//
// <#=GetResourceString("Template_GeneratedCodeCommentLine2")#>
// <#=GetResourceString("Template_GeneratedCodeCommentLine3")#>
// </auto-generated>
//------------------------------------------------------------------------------

<#+
    fileManager.EndBlock();
}

void BeginNamespace()
{
    var region = new CodeRegion(this);

    if (!String.IsNullOrEmpty(ObjectNamespace))
    {
#>
namespace <#=Code.EscapeNamespace(ObjectNamespace)#>
{
<#+
        PushIndent(CodeRegion.GetIndent(1));
    }
}

void EndNamespace()
{
    if (!String.IsNullOrEmpty(ObjectNamespace))
    {
        PopIndent();
#>
}
<#+
    }
}

void WriteFunctionDefinition(EdmFunction edmFunction, bool includeMergeOption)
{
    var parameters = FunctionImportParameter.Create(edmFunction.Parameters, Code, Tools);
    var paramList = String.Join(", ", parameters.Select(p => p.FunctionParameterType + " " + p.FunctionParameterName).ToArray());
    var returnType = edmFunction.ReturnParameter == null ? null : Tools.GetElementType(edmFunction.ReturnParameter.TypeUsage);
    var processedReturn = returnType == null ? "int" : "IEnumerable<" + MultiSchemaEscape(returnType) + ">";

    if (includeMergeOption)
    {
        paramList = Code.StringAfter(paramList, ", ") + "MergeOption mergeOption";
    }
#>

    <#=processedReturn#> <#=Code.Escape(edmFunction)#>(<#=paramList#>);
<#+
    if(!includeMergeOption && returnType != null && returnType.EdmType.BuiltInTypeKind == BuiltInTypeKind.EntityType)
    {
        WriteFunctionDefinition(edmFunction, true);
    }
}

void WriteFunctionRepresentation(EdmFunction edmFunction, bool includeMergeOption)
{
    var parameters = FunctionImportParameter.Create(edmFunction.Parameters, Code, Tools);
    var paramList = String.Join(", ", parameters.Select(p => p.FunctionParameterType + " " + p.FunctionParameterName).ToArray());
    var returnType = edmFunction.ReturnParameter == null ? null : Tools.GetElementType(edmFunction.ReturnParameter.TypeUsage);
    var processedReturn = returnType == null ? "int" : "IEnumerable<" + MultiSchemaEscape(returnType) + ">";

    if (includeMergeOption)
    {
        paramList = Code.StringAfter(paramList, ", ") + "MergeOption mergeOption";
    }
#>

    <#=AccessibilityAndVirtual(Accessibility.ForMethod(edmFunction))#> <#=processedReturn#> <#=Code.Escape(edmFunction)#>(<#=paramList#>)
    {
<#+
        if(returnType != null && (returnType.EdmType.BuiltInTypeKind == BuiltInTypeKind.EntityType ||
                                  returnType.EdmType.BuiltInTypeKind == BuiltInTypeKind.ComplexType))
        {
#>
        ((IObjectContextAdapter)this).ObjectContext.MetadataWorkspace.LoadFromAssembly(typeof(<#=MultiSchemaEscape(returnType)#>).Assembly);

<#+
        }

        foreach (var parameter in parameters.Where(p => p.NeedsLocalVariable))
        {
            var isNotNull = parameter.IsNullableOfT ? parameter.FunctionParameterName + ".HasValue" : parameter.FunctionParameterName + " != null";
            var notNullInit = "new ObjectParameter(\"" + parameter.EsqlParameterName + "\", " + parameter.FunctionParameterName + ")";
            var nullInit = "new ObjectParameter(\"" + parameter.EsqlParameterName + "\", typeof(" + parameter.RawClrTypeName + "))";
#>
        var <#=parameter.LocalVariableName#> = <#=isNotNull#> ?
            <#=notNullInit#> :
            <#=nullInit#>;

<#+
        }

        var genericArg = returnType == null ? "" : "<" + MultiSchemaEscape(returnType) + ">";
        var callParams = Code.StringBefore(", ", String.Join(", ", parameters.Select(p => p.ExecuteParameterName).ToArray()));

        if (includeMergeOption)
        {
            callParams = ", mergeOption" + callParams;
        }
#>
        return ((IObjectContextAdapter)this).ObjectContext.ExecuteFunction<#=genericArg#>("<#=edmFunction.Name#>"<#=callParams#>);
    }
<#+
    if(!includeMergeOption && returnType != null && returnType.EdmType.BuiltInTypeKind == BuiltInTypeKind.EntityType)
    {
        WriteFunctionRepresentation(edmFunction, true);
    }
}

string AccessibilityAndVirtual(string accessibility)
{
    return accessibility + (accessibility != "private" ? " virtual" : "");
}

string MultiSchemaEscape(TypeUsage usage)
{
    var type = usage.EdmType as StructuralType;

    return type != null && type.NamespaceName != ModelNamespace ?
        Code.CreateFullName(Code.EscapeNamespace(type.NamespaceName), Code.Escape(type)) :
        Code.Escape(usage);
}
#>