go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/db/dbutil/query_formatter.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package dbutil
     9  
    10  import (
    11  	"bytes"
    12  	"fmt"
    13  	"reflect"
    14  	"text/template"
    15  
    16  	"go.charczuk.com/sdk/db"
    17  )
    18  
    19  // QueryFormatter is a type that can format queries.
    20  type QueryFormatter struct {
    21  	vars           map[string]any
    22  	templates      []string
    23  	funcs          template.FuncMap
    24  	models         map[string]any
    25  	modelTypeMeta  map[string]*db.TypeMeta
    26  	modelTableName map[string]string
    27  
    28  	t *template.Template
    29  }
    30  
    31  // WithTemplates adds snippets or template to the formatter.
    32  func (qf *QueryFormatter) WithTemplates(templates ...string) *QueryFormatter {
    33  	qf.templates = append(qf.templates, templates...)
    34  	return qf
    35  }
    36  
    37  // WithFunc adds a view func to the formatter.
    38  func (qf *QueryFormatter) WithFunc(name string, fn any) *QueryFormatter {
    39  	if qf.funcs == nil {
    40  		qf.funcs = make(template.FuncMap)
    41  	}
    42  	qf.funcs[name] = fn
    43  	return qf
    44  }
    45  
    46  // WithModels registers model types for use in view func helpers.
    47  func (qf *QueryFormatter) WithModels(models ...any) *QueryFormatter {
    48  	if qf.models == nil {
    49  		qf.models = make(map[string]any)
    50  	}
    51  	if qf.modelTypeMeta == nil {
    52  		qf.modelTypeMeta = make(map[string]*db.TypeMeta)
    53  	}
    54  	if qf.modelTableName == nil {
    55  		qf.modelTableName = make(map[string]string)
    56  	}
    57  	for _, m := range models {
    58  		modelName := reflect.TypeOf(m).Name()
    59  		qf.models[modelName] = m
    60  		qf.modelTypeMeta[modelName] = db.TypeMetaFor(m)
    61  		qf.modelTableName[modelName] = db.TableName(m)
    62  	}
    63  	return qf
    64  }
    65  
    66  // MustFormatQuery formats a query with a given variadic set of options and panics on error.
    67  func (qf *QueryFormatter) MustFormat(body string, vars any) (output string) {
    68  	var err error
    69  	output, err = qf.FormatQuery(body, vars)
    70  	if err != nil {
    71  		panic(err)
    72  	}
    73  	return output
    74  }
    75  
    76  func (qf *QueryFormatter) Initialize() *QueryFormatter {
    77  	if qf.t == nil {
    78  		qf.t = template.New("")
    79  		if qf.funcs == nil {
    80  			qf.funcs = make(template.FuncMap)
    81  		}
    82  		qf.funcs["columns"] = func(modelName string) (string, error) {
    83  			columns, ok := qf.modelTypeMeta[modelName]
    84  			if !ok {
    85  				return "", fmt.Errorf("invalid model; %s not found", modelName)
    86  			}
    87  			return db.ColumnNamesCSV(columns.Columns()), nil
    88  		}
    89  		qf.funcs["columns_alias"] = func(modelName, alias string) (string, error) {
    90  			columns, ok := qf.modelTypeMeta[modelName]
    91  			if !ok {
    92  				return "", fmt.Errorf("invalid model; %s not found", modelName)
    93  			}
    94  			return db.ColumnNamesFromAliasCSV(columns.Columns(), alias), nil
    95  		}
    96  		qf.funcs["columns_prefix_alias"] = func(modelName, prefix, alias string) (string, error) {
    97  			columns, ok := qf.modelTypeMeta[modelName]
    98  			if !ok {
    99  				return "", fmt.Errorf("invalid model; %s not found", modelName)
   100  			}
   101  			return db.ColumnNamesWithPrefixFromAliasCSV(columns.Columns(), prefix, alias), nil
   102  		}
   103  		qf.funcs["table"] = func(modelName string) (string, error) {
   104  			tableName, ok := qf.modelTableName[modelName]
   105  			if !ok {
   106  				return "", fmt.Errorf("invalid model; %s not found", modelName)
   107  			}
   108  			return tableName, nil
   109  		}
   110  		qf.t = qf.t.Funcs(qf.funcs)
   111  		var err error
   112  		for _, template := range qf.templates {
   113  			qf.t, err = qf.t.Parse(template)
   114  			if err != nil {
   115  				err = fmt.Errorf("cannot parse template; %s: %w", template, err)
   116  				panic(err)
   117  			}
   118  		}
   119  	}
   120  	return qf
   121  }
   122  
   123  // FormatQuery formats a query with a given variadic set of options.
   124  func (qf *QueryFormatter) FormatQuery(body string, vars any) (output string, err error) {
   125  	qf.Initialize()
   126  	var tmpl *template.Template
   127  	tmpl, err = qf.t.Parse(body)
   128  	if err != nil {
   129  		err = fmt.Errorf("cannot parse query; %s: %w", body, err)
   130  		return
   131  	}
   132  	buf := new(bytes.Buffer)
   133  	if err = tmpl.Execute(buf, vars); err != nil {
   134  		err = fmt.Errorf("cannot execute query; %s: %w", body, err)
   135  		return
   136  	}
   137  	output = buf.String()
   138  	return
   139  }