Merge pull request #611 from erizocosmico/feature/config-bblfsh-size

function: make max bblfsh blob size configurable

Author: ajnavarro

docs/using-gitbase/configuration.md

@@ -16,6 +16,7 @@
 | `GITBASE_CACHESIZE_MB`       | size of the cache for git objects specified as MB                                  |
 | `GITBASE_CONNECTION_TIMEOUT` | timeout in seconds used for client connections on write and reads. No timeout by default.     |
 | `GITBASE_USER_FILE`          | JSON file with user credentials                                                    |
+| `GITBASE_MAX_UAST_BLOB_SIZE`          | Max size of blobs to send to be parsed by bblfsh. Default: 5242880 (5MB)                                                    |
 
 ### Jaeger tracing variables
 

internal/function/uast.go

@@ -21,14 +21,18 @@ import (
 )
 
 const (
-	uastCacheSize        = "GITBASE_UAST_CACHE_SIZE"
+	uastCacheSizeKey     = "GITBASE_UAST_CACHE_SIZE"
 	defaultUASTCacheSize = 10000
+
+	uastMaxBlobSizeKey     = "GITBASE_MAX_UAST_BLOB_SIZE"
+	defaultUASTMaxBlobSize = 5 * 1024 * 1024 // 5MB
 )
 
 var uastCache *lru.Cache
+var uastMaxBlobSize int
 
 func init() {
-	s := os.Getenv(uastCacheSize)
+	s := os.Getenv(uastCacheSizeKey)
 	size, err := strconv.Atoi(s)
 	if err != nil || size <= 0 {
 		size = defaultUASTCacheSize
@@ -38,6 +42,11 @@ func init() {
 	if err != nil {
 		panic(fmt.Errorf("cannot initialize UAST cache: %s", err))
 	}
+
+	uastMaxBlobSize, err = strconv.Atoi(os.Getenv(uastMaxBlobSizeKey))
+	if err != nil {
+		uastMaxBlobSize = defaultUASTMaxBlobSize
+	}
 }
 
 // uastFunc shouldn't be used as an sql.Expression itself.
@@ -167,6 +176,18 @@ func (u *uastFunc) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err err
 		return nil, nil
 	}
 
+	if uastMaxBlobSize >= 0 && len(bytes) > uastMaxBlobSize {
+		logrus.WithFields(logrus.Fields{
+			"max":  uastMaxBlobSize,
+			"size": len(bytes),
+		}).Warnf(
+			"uast will be skipped, file is too big to send to bblfsh."+
+				"This can be configured using %s environment variable",
+			uastMaxBlobSizeKey,
+		)
+		return nil, nil
+	}
+
 	lang, err := exprToString(ctx, u.Lang, row)
 	if err != nil {
 		return nil, err

internal/function/uast_test.go

@@ -62,6 +62,34 @@ func TestUASTMode(t *testing.T) {
 	}
 }
 
+func TestUASTMaxBlobSize(t *testing.T) {
+	ctx, cleanup := setup(t)
+	defer cleanup()
+
+	fn := NewUASTMode(
+		expression.NewGetField(0, sql.Text, "", false),
+		expression.NewGetField(1, sql.Blob, "", false),
+		expression.NewGetField(2, sql.Text, "", false),
+	)
+
+	u, _ := bblfshFixtures(t, ctx)
+
+	require := require.New(t)
+	row := sql.NewRow("annotated", []byte(testCode), "Python")
+	result, err := fn.Eval(ctx, row)
+	require.NoError(err)
+
+	assertUASTBlobs(t, ctx, u["annotated"], result)
+
+	uastMaxBlobSize = 2
+
+	result, err = fn.Eval(ctx, row)
+	require.NoError(err)
+	require.Nil(result)
+
+	uastMaxBlobSize = defaultUASTMaxBlobSize
+}
+
 func TestUAST(t *testing.T) {
 	ctx, cleanup := setup(t)
 	defer cleanup()