Restrict upload of particular files types in FILES related list- Salesforce lightning

Sometimes we need to restrict the 'Type' of a file which end-user can upload to  'Files' related list of any salesforce record in lightning experience for example

1. Allow only word, pdf, images and text files  
2. Restrict .zip, .exe files etc (any particular file format) 

Here I'll be implementing and explaining:

Requirement: Account object has 'Banking' record type. User should not be able to upload '.pdf' files in 'Files' related list for banking account records.

Quick overview video

Before we proceed to implementation, we should know a few important things: 

1. In lightning, 'Attachment' has been replaced by 'Files/Content'. 'Files' related list also shows attachments (uploaded previously via classic experience) if there is any. 
2. If a document is uploaded in Lightning experience, it becomes part of Files/Content object. If the user switches to classic experience and uploads any document to 'Notes and Attachment' related list then it goes to 'Attachment' object. 
3.  ContentDocument: This object record we don’t create. It gets created when we create Content Version who is a child of ContentDocument. But Id of this record will be required to do other stuff.
4. ContentVersion: This is where our attachment will be inserted but still it will not be visible under Attachment/Files related list. For that, we need to create ContentDocumentLink.
5. ContentDocumentLink: It will share the files with users, records, groups etc and create files under your object records You can create multiple records to attach the same files under multiple records.

What we need to do to meet the requirement?

1. Create a trigger on 'ContentDocumentLink' object and utilizing after insert event.
2. Loop through ContentDocumentLink records.
3. Check 'LinkedEntityId' i.e. if the document is associated with an Account record or something else.
4. Check for RecordType
5. Finally restricting by using .addError(); 

 Trigger Code: 

// Trigger on ContentDocumentLink object : Created by Sudhanshu Gupta - SalesForce Technical Warrior

trigger trig_ContentDocumentLink_AfterInsert on ContentDocumentLink (after insert) {

   if(Trigger.isAfter && Trigger.isInsert){

      System.debug('SFDC Tech Warrior... Trigger started after insert.');

      TW_ContentDocumentLink_Handler.onAfterInsert(Trigger.New);  

    }

}

********************************************************************************

Helper Class Code:

/* 

Class : TW_ContentDocumentLink_Handler

Description: Called from trigger - 'trig_ContentDocumentLink_AfterInsert', after insert case.   

in Banking type account (Record Type) Restrict pdf files upload through 'Upload Files' button present in Notes and Attachments  

Owner : Sudhanshu Gupta Technical Warrior- 2020

*/

public class TW_ContentDocumentLink_Handler {

     //This method restrict .pdf file upload to 'Banking' type account records

     public static void onAfterInsert(list<ContentDocumentLink> lstCntLinks) {

        String strObjPrefix; // to check parent record prefix value like for Account object: 001, for Opportunity: 006 

        Set<Id> setCntDocIds = new set<Id>();

        set<Id> setAgmtIds = new set<Id>();//Store Parent record i.e. Agreement Ids

        map<Id, Account> mapAgmt;

        try{

            for(ContentDocumentLink clIterator : lstCntLinks) {

                strObjPrefix = String.valueOf(clIterator.LinkedEntityId).substring(0, 3); // Return first letters of record id like 001 for Account 

                if(strObjPrefix == Account.sObjectType.getDescribe().getKeyPrefix()) {

                    setCntDocIds.add(clIterator.ContentDocumentId);// Content Document Id

                    setAgmtIds.add(clIterator.LinkedEntityId);// Agreement Id - Parent record Id

                }

            }

            Id recordTypeIdBanking = Schema.SObjectType.Account.getRecordTypeInfosByName().get('Banking').getRecordTypeId();// Account object- 'Banking' Reord Type id

            if(setCntDocIds.size() > 0 && setAgmtIds.size() > 0 ) {           

                mapAgmt = new map<Id, Account>([SELECT Id, Name,RecordTypeId FROM Account WHERE Id IN :setAgmtIds and RecordTypeId = :recordTypeIdBanking]);

                //Fetching parent record details and restricting only if Record type is Account 

                if(mapAgmt.size() > 0){

                    map<Id, ContentDocument> mapContentDocuments = new map<Id, ContentDocument>([SELECT Id, Title, FileExtension FROM ContentDocument WHERE Id IN :setCntDocIds]);

                    list<ContentDocument> lstCntDocsToUpdate = new list<ContentDocument>();        

                    for(ContentDocumentLink cdlIterator : lstCntLinks) {

                        ContentDocument objCntDoc = mapContentDocuments.get(cdlIterator.ContentDocumentId);

                        // Allow all files except : pdf 

                        if(objCntDoc.FileExtension == 'pdf'){

                            cdlIterator.addError('You can not upload pdf files for Banking type accounts.');//Showing error   

                            // This line will abort file upload option. 

                            // Error message will be displayed only when 'ContentDocumentLink' object has some RecordType created (any other than Master)

                            // else Generic message will be dispayed - 'Can not add 1 file to Account'.   

                        }                        

                    }  

                }

            }

        }

        Catch (Exception ex){

            system.debug('Exception in TW_ContentDocumentLink_Handler class :' + ex.getMessage());

        }

    }

}

*********************************************************************************

Follow me @LinkedIn - sud-gupta

Keep Learning ... Keep Growing !!!