Add configuration parameters for JSON maximum line size as well as fi…

…le reader buffer size.
cube2222 · Oct 15, 2022 · 6644557 · 6644557
1 parent cc43f25
commit 6644557
Show file tree

Hide file tree

Showing 9 changed files with 62 additions and 30 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -78,6 +78,12 @@ octosql "SELECT * FROM plugins.plugins"`,
 			defer profile.Start(profile.TraceProfile, profile.ProfilePath(".")).Stop()
 		}
 		ctx := cmd.Context()
+		cfg, err := config.Read()
+		if err != nil {
+			return fmt.Errorf("couldn't read config: %w", err)
+		}
+		ctx = config.ContextWithConfig(ctx, cfg)
+
 		debug.SetGCPercent(1000)
 
 		logs.InitializeFileLogger()
@@ -94,11 +100,6 @@ octosql "SELECT * FROM plugins.plugins"`,
 			}
 		}()
 
-		cfg, err := config.Read()
-		if err != nil {
-			return fmt.Errorf("couldn't read config: %w", err)
-		}
-
 		installedPlugins, err := pluginManager.ListInstalledPlugins()
 		if err != nil {
 			return fmt.Errorf("couldn't list installed plugins: %w", err)
@@ -188,15 +189,15 @@ octosql "SELECT * FROM plugins.plugins"`,
 		if err != nil {
 			return fmt.Errorf("couldn't get file extension handlers: %w", err)
 		}
-		fileHandlers := map[string]func(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error){
+		fileHandlers := map[string]func(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error){
 			"csv":     csv.Creator(','),
 			"json":    json.Creator,
 			"lines":   lines.Creator,
 			"parquet": parquet.Creator,
 			"tsv":     csv.Creator('\t'),
 		}
 		for ext, pluginName := range fileExtensionHandlers {
-			fileHandlers[ext] = func(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+			fileHandlers[ext] = func(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
 				db, err := databases[pluginName]()
 				if err != nil {
 					return nil, physical.Schema{}, fmt.Errorf("couldn't get plugin %s database for plugin extensions %s: %w", pluginName, ext, err)
@@ -513,13 +514,13 @@ func typecheckExpr(ctx context.Context, expr logical.Expression, env physical.En
 }
 
 type fileTypeDatabaseCreator struct {
-	creator func(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error)
+	creator func(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error)
 }
 
 func (f *fileTypeDatabaseCreator) ListTables(ctx context.Context) ([]string, error) {
 	return nil, nil
 }
 
 func (f *fileTypeDatabaseCreator) GetTable(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
-	return f.creator(name, options)
+	return f.creator(ctx, name, options)
 }
diff --git a/config/config.go b/config/config.go
@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -23,6 +24,7 @@ var octosqlHomeDir = func() string {
 
 type Config struct {
 	Databases []DatabaseConfig `yaml:"databases"`
+	Files     FilesConfig      `yaml:"files"`
 }
 
 type DatabaseConfig struct {
@@ -32,20 +34,34 @@ type DatabaseConfig struct {
 	Config  yaml.Node                            `yaml:"config"`
 }
 
+type FilesConfig struct {
+	JSON            JSONConfig `yaml:"json"`
+	BufferSizeBytes int        `yaml:"buffer_size_bytes"`
+}
+
+type JSONConfig struct {
+	MaxLineSizeBytes int `yaml:"max_line_size_bytes"`
+}
+
 func Read() (*Config, error) {
+	var config Config
 	data, err := ioutil.ReadFile(filepath.Join(octosqlHomeDir, "octosql.yml"))
-	if err != nil {
-		if os.IsNotExist(err) {
-			return &Config{
-				Databases: nil,
-			}, nil
-		}
+	if err != nil && os.IsNotExist(err) {
+		config = Config{}
+	} else if err != nil {
 		return nil, fmt.Errorf("couldn't read config file: %w", err)
+	} else {
+		if err := yaml.Unmarshal(data, &config); err != nil {
+			return nil, fmt.Errorf("couldn't unmarshal yaml configuration: %w", err)
+		}
 	}
 
-	var config Config
-	if err := yaml.Unmarshal(data, &config); err != nil {
-		return nil, fmt.Errorf("couldn't unmarshal yaml configuration: %w", err)
+	// TODO: Refactor to a sensibly structured default value system.
+	if config.Files.BufferSizeBytes == 0 {
+		config.Files.BufferSizeBytes = 4096 * 1024
+	}
+	if config.Files.JSON.MaxLineSizeBytes == 0 {
+		config.Files.JSON.MaxLineSizeBytes = 1024 * 1024
 	}
 
 	return &config, nil
@@ -89,3 +105,15 @@ func (r *PluginReference) UnmarshalText(text []byte) error {
 func (r *PluginReference) String() string {
 	return fmt.Sprintf("%s/%s", r.Repository, r.Name)
 }
+
+// TODO: Using a custom context everywhere would be better.
+
+type contextKey struct{}
+
+func ContextWithConfig(ctx context.Context, config *Config) context.Context {
+	return context.WithValue(ctx, contextKey{}, config)
+}
+
+func FromContext(ctx context.Context) *Config {
+	return ctx.Value(contextKey{}).(*Config)
+}
diff --git a/datasources/csv/impl.go b/datasources/csv/impl.go
@@ -14,9 +14,9 @@ import (
 	"github.com/cube2222/octosql/physical"
 )
 
-func Creator(separator rune) func(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
-	return func(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
-		f, err := files.OpenLocalFile(context.Background(), name, files.WithPreview())
+func Creator(separator rune) func(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+	return func(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+		f, err := files.OpenLocalFile(ctx, name, files.WithPreview())
 		if err != nil {
 			return nil, physical.Schema{}, fmt.Errorf("couldn't open local file: %w", err)
 		}

diff --git a/datasources/json/execution.go b/datasources/json/execution.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/valyala/fastjson"
 
+	"github.com/cube2222/octosql/config"
 	. "github.com/cube2222/octosql/execution"
 	"github.com/cube2222/octosql/execution/files"
 	"github.com/cube2222/octosql/octosql"
@@ -28,7 +29,7 @@ func (d *DatasourceExecuting) Run(ctx ExecutionContext, produce ProduceFn, metaS
 	defer f.Close()
 
 	sc := bufio.NewScanner(f)
-	sc.Buffer(nil, 1024*1024)
+	sc.Buffer(nil, config.FromContext(ctx).Files.JSON.MaxLineSizeBytes)
 
 	// All async processing in this function and all jobs created by it use this context.
 	// This means that returning from this function will properly clean up all async processors.

diff --git a/datasources/json/impl.go b/datasources/json/impl.go
@@ -15,8 +15,8 @@ import (
 	"github.com/cube2222/octosql/physical"
 )
 
-func Creator(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
-	f, err := files.OpenLocalFile(context.Background(), name, files.WithPreview())
+func Creator(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+	f, err := files.OpenLocalFile(ctx, name, files.WithPreview())
 	if err != nil {
 		return nil, physical.Schema{}, fmt.Errorf("couldn't open local file: %w", err)
 	}

diff --git a/datasources/lines/impl.go b/datasources/lines/impl.go
@@ -10,8 +10,8 @@ import (
 	"github.com/cube2222/octosql/physical"
 )
 
-func Creator(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
-	f, err := files.OpenLocalFile(context.Background(), name, files.WithPreview())
+func Creator(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+	f, err := files.OpenLocalFile(ctx, name, files.WithPreview())
 	if err != nil {
 		return nil, physical.Schema{}, fmt.Errorf("couldn't open local file: %w", err)
 	}

diff --git a/datasources/parquet/impl.go b/datasources/parquet/impl.go
@@ -12,7 +12,7 @@ import (
 	"github.com/cube2222/octosql/physical"
 )
 
-func Creator(name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
+func Creator(ctx context.Context, name string, options map[string]string) (physical.DatasourceImplementation, physical.Schema, error) {
 	f, err := os.Open(name)
 	if err != nil {
 		return nil, physical.Schema{}, fmt.Errorf("couldn't open file: %w", err)

diff --git a/execution/files/files.go b/execution/files/files.go
@@ -10,6 +10,8 @@ import (
 	"sync"
 
 	"github.com/nxadm/tail"
+
+	"github.com/cube2222/octosql/config"
 )
 
 type customCloser struct {
@@ -113,7 +115,7 @@ func OpenLocalFile(ctx context.Context, path string, opts ...OpenFileOption) (io
 			return nil, fmt.Errorf("couldn't open file: %w", err)
 		}
 		return &customCloser{
-			Reader: bufio.NewReaderSize(f, 4096*1024),
+			Reader: bufio.NewReaderSize(f, config.FromContext(ctx).Files.BufferSizeBytes),
 			close:  f.Close,
 		}, nil
 	} else {

diff --git a/physical/physical.go b/physical/physical.go
@@ -54,7 +54,7 @@ type DatasourceRepository struct {
 	// Bo inaczej będzie na start strasznie dużo rzeczy ładować niepotrzebnych dla wszystkich
 	// skonfigurowanych baz danych.
 	Databases    map[string]func() (Database, error)
-	FileHandlers map[string]func(name string, options map[string]string) (DatasourceImplementation, Schema, error)
+	FileHandlers map[string]func(ctx context.Context, name string, options map[string]string) (DatasourceImplementation, Schema, error)
 }
 
 type Database interface {
@@ -78,7 +78,7 @@ func (dr *DatasourceRepository) GetDatasource(ctx context.Context, name string,
 	if index := strings.LastIndex(name, "."); index != -1 {
 		extension := name[index+1:]
 		if handler, ok := dr.FileHandlers[extension]; ok {
-			return handler(name, options)
+			return handler(ctx, name, options)
 		}
 	}