Merge pull request #239 from digma-ai/feature/flask_client_side_discovery

Feature/flask client side discovery

Author: doppleware

src/analyticsCodeLens.ts

@@ -239,8 +239,11 @@ class CodelensProvider implements vscode.CodeLensProvider<vscode.CodeLens>
                 const lenses = await this.getLensForCodeLocationObject(methodInfo,
                     uniqueEndpoints,documentInfo.usageData.getAll(),documentInfo.insights.all.filter(x=>x.scope=="EntrySpan"|| x.scope=="Span"),
                                         );
+                
+                const uniqueLenses =[...new Map(lenses.map(item =>
+                    [item.command!.title, item])).values()];
 
-                for (const lens of lenses){
+                for (const lens of uniqueLenses){
                     codelens.push(lens);
                 }         
 

src/services/codeInspector.ts

@@ -67,7 +67,7 @@ export class CodeInspector {
         return;
     }
 
-    private async getType(
+    public async getType(
         usageDocument: vscode.TextDocument,
         usagePosition: vscode.Position,
     ): Promise<Definition | undefined> {

src/services/documentInfoProvider.ts

@@ -183,7 +183,7 @@ export class DocumentInfoProvider implements vscode.Disposable
                 const symbolTrees = await this.symbolProvider.getSymbolTree(doc);
                 const tokens = await this.symbolProvider.getTokens(doc);
                 const symbolInfos = await this.symbolProvider.getMethods(doc, tokens, symbolTrees);
-                const endpoints = await this.symbolProvider.getEndpoints(doc, symbolInfos, tokens, symbolTrees, this);
+                const endpoints = await this.symbolProvider.getEndpoints(doc, symbolInfos, tokens, symbolTrees );
                 const spans = await this.symbolProvider.getSpans(doc, symbolInfos, tokens);
                 const paramsExtractor = await this.symbolProvider.getParametersExtractor(doc);
                 const symbolAliasExtractor = await this.symbolProvider.getSymbolAliasExtractor(doc);

src/services/languages/csharp/AspNetCoreMvcEndpointExtractor.ts

@@ -1,7 +1,7 @@
 import * as vscode from 'vscode';
 import { DocumentInfoProvider } from './../../documentInfoProvider';
 import { CodeInspector } from '../../codeInspector';
-import { SymbolTree } from './../symbolProvider';
+import { SymbolProvider, SymbolTree } from './../symbolProvider';
 import { Token, TokenType } from '../tokens';
 import { EndpointInfo, IEndpointExtractor, SymbolInfo } from '../extractors';
 import { convertRange } from '../../utils';
@@ -16,7 +16,7 @@ export class AspNetCoreMvcEndpointExtractor implements IEndpointExtractor {
         symbolInfo: SymbolInfo[],
         tokens: Token[],
         symbolTrees: SymbolTree[] | undefined,
-        documentInfoProvider: DocumentInfoProvider,
+        symbolProvider: SymbolProvider
     ): Promise<EndpointInfo[]> {
         const results: EndpointInfo[] = [];
 
@@ -28,7 +28,7 @@ export class AspNetCoreMvcEndpointExtractor implements IEndpointExtractor {
             const derivesFromControllerBase = await this._codeInspector.derivesFrom(
                 classDefinition,
                 'ControllerBase',
-                documentInfoProvider.symbolProvider,
+                symbolProvider,
                 this.findParentToken,
             );
             if(!derivesFromControllerBase) {

src/services/languages/extractors.ts

@@ -100,7 +100,7 @@ export interface IEndpointExtractor {
         symbolInfos: SymbolInfo[],
         tokens: Token[],
         symbolTrees: SymbolTree[] | undefined,
-        documentInfoProvider: DocumentInfoProvider,
+        symbolProvider: SymbolProvider
     ): Promise<EndpointInfo[]>;
 }
 

src/services/languages/languageExtractor.ts

@@ -5,6 +5,7 @@ import { IMethodExtractor, IParametersExtractor, IEndpointExtractor, ISpanExtrac
 import { BasicParametersExtractor } from './defaultImpls';
 import { IModulePathToUriConverter } from './modulePathToUriConverters';
 import { ICodeObjectIdParser, CommonCodeObjectIdParser } from '../codeObject';
+import { FlaskEndpointExtractor } from './python/flaskEndpointExtractor';
 
 export interface ILanguageExtractor {
     requiredExtensionLoaded: boolean;
@@ -44,7 +45,7 @@ export abstract class LanguageExtractor implements ILanguageExtractor {
     }
 
     public getEndpointExtractors(codeInspector: CodeInspector): IEndpointExtractor[] {
-        return [];
+        return [new FlaskEndpointExtractor(codeInspector)];
     }
 
     public abstract getSpanExtractors(codeInspector: CodeInspector): ISpanExtractor[];

src/services/languages/python/fastapiEndpointExtractor.ts

@@ -1,6 +1,6 @@
 import * as vscode from 'vscode';
 import { DocumentInfoProvider } from './../../documentInfoProvider';
-import { SymbolTree } from './../symbolProvider';
+import { SymbolProvider, SymbolTree } from './../symbolProvider';
 import { Token, TokenType } from '../tokens';
 import { EndpointInfo, IEndpointExtractor, SymbolInfo } from '../extractors';
 import { EndpointSchema } from '../../analyticsProvider';
@@ -12,7 +12,7 @@ export class FastapiEndpointExtractor implements IEndpointExtractor
         symbolInfo: SymbolInfo[],
         tokens: Token[],
         symbolTrees: SymbolTree[] | undefined,
-        documentInfoProvider: DocumentInfoProvider,
+        symbolProvider: SymbolProvider,
     ): Promise<EndpointInfo[]> {
         // Ensure fastapi module was imported
         if(!tokens.any(t => t.text == 'fastapi' && t.type == TokenType.module))

src/services/languages/python/flaskEndpointExtractor.ts

@@ -0,0 +1,112 @@
+import * as vscode from 'vscode';
+import { DocumentInfoProvider } from './../../documentInfoProvider';
+import { SymbolProvider, SymbolTree } from './../symbolProvider';
+import { Token, TokenType } from '../tokens';
+import { EndpointInfo, IEndpointExtractor, SymbolInfo } from '../extractors';
+import { EndpointSchema } from '../../analyticsProvider';
+import { CodeInspector } from '../../codeInspector';
+
+export class FlaskEndpointExtractor implements IEndpointExtractor
+{
+    readonly routeMethods = ['route', 'get','post', 'patch'];
+    readonly genericRouteMethod = 'route';
+    readonly flaskAppType = 'Flask';
+
+
+    constructor(private _codeInspector: CodeInspector) {}
+    
+    async extractEndpoints(
+        document: vscode.TextDocument,
+        symbolInfo: SymbolInfo[],
+        tokens: Token[],
+        symbolTrees: SymbolTree[] | undefined,
+        symbolProvider: SymbolProvider
+    ): Promise<EndpointInfo[]> {
+
+        // Ensure flask module was imported
+        if(!tokens.any(t => t.text == 'flask' && t.type == TokenType.module))
+            return [];
+        
+        // Search for "@app.get" decorators
+        const results: EndpointInfo[] = [];
+        for(let i=0; i<tokens.length-1; i++)
+        {
+            const appToken = tokens[i];
+            const methodToken = tokens[i+1];
+            let route = '';
+            let method = '';
+            
+            //Check if this is a router variable
+            if ( appToken.type === TokenType.variable 
+                    && methodToken.type === TokenType.method && this.routeMethods.includes( methodToken.text) ){
+               
+                const routerLine = document.getText(new vscode.Range(
+                    appToken.range.start, 
+                    new vscode.Position(methodToken.range.end.line, 1000)));
+                    
+                //extract the prefix
+                let match = new RegExp(`^(?:${appToken.text}.(route|get|post|update|patch))\\s*\\((?:'|")(.*?)(?:'|")(?:.*[\\s,]*\\bmethods\\b\\s*=\\s*\\[(.*?)\\])?`).exec(routerLine);
+
+                if (!match || match.length<3){
+                    return [];
+                }
+                const tracerDefTypeDef = await this._codeInspector.getTypeFromSymbolProvider(document, appToken.range.start,symbolProvider,x=>true);
+
+                if (!tracerDefTypeDef || tracerDefTypeDef!=this.flaskAppType){
+                    return [];
+                }
+
+               
+                route = match[2];
+                method = match[1];
+
+                if (method==this.genericRouteMethod){
+
+                    //Handle form of @app.route('/', methods: ["POST", "GET"])
+                    //In this scenario the regex would capture the methods 
+                    //As the fourth group
+                    if (match.length==4 && match[3]){
+                        const methods = match[3].replaceAll('"','');
+                        method=methods;
+                    }
+                    //Otherwise the default for the "route" form is GET
+                    else{
+                        method='GET';
+
+                    }
+                }
+                
+
+            }
+            else{
+                continue;
+            }
+                            
+            const relevantFunc = symbolInfo.firstOrDefault(s => s.range.contains(methodToken.range));
+            if (!relevantFunc){
+                continue;
+            }
+            
+            let folder = vscode.workspace.getWorkspaceFolder(document.uri);
+            let folderPrefix = folder?.uri.path.replaceAll('\\',"/").split('/').slice(0,-1).join('/');
+            let relevantPath = document.uri.path.substring(folderPrefix!.length);
+            let pathParts = relevantPath.split('/').filter(x=>x);
+            const methods = method.split(",").map(x=>x.toUpperCase().trim());
+
+            for (var apiMethod of methods){
+                for (let j=0;j<pathParts.length-1;j++){
+                    let possibleRoot = pathParts[j];
+                    results.push(new EndpointInfo(
+                        possibleRoot + '$_$' + EndpointSchema.HTTP + "HTTP " + apiMethod + ' ' + route,
+                        apiMethod, 
+                        route,
+                        relevantFunc.range,
+                        document.uri));
+                }
+            }
+  
+
+        }
+        return results;
+    }
+}

src/services/languages/symbolProvider.ts

@@ -103,8 +103,7 @@ export class SymbolProvider
         document: vscode.TextDocument,
         symbolInfos: SymbolInfo[],
         tokens: Token[],
-        symbolTrees: SymbolTree[] | undefined,
-        documentInfoProvider: DocumentInfoProvider,
+        symbolTrees: SymbolTree[] | undefined
     ): Promise<EndpointInfo[]> {
         const supportedLanguage = await this.getSupportedLanguageExtractor(document);
         if(!supportedLanguage) {
@@ -113,7 +112,7 @@ export class SymbolProvider
 
         const endpointExtractors = supportedLanguage.getEndpointExtractors(this._codeInspector);
         const extractedEndpoints = await Promise.all(
-            endpointExtractors.map(async (x) => await x.extractEndpoints(document, symbolInfos, tokens, symbolTrees, documentInfoProvider))
+            endpointExtractors.map(async (x) => await x.extractEndpoints(document, symbolInfos, tokens, symbolTrees,  this))
         );
         const endpoints = extractedEndpoints.flat();
         return endpoints;