internal/codegen/golang: omit unused structs from output

Author: andrewmbenton

internal/cmd/shim.go

@@ -108,6 +108,7 @@ func pluginGoCode(s config.SQLGo) *plugin.GoCode {
 		OutputFilesSuffix:           s.OutputFilesSuffix,
 		InflectionExcludeTableNames: s.InflectionExcludeTableNames,
 		QueryParameterLimit:         s.QueryParameterLimit,
+		OmitUnusedStructs:           s.OmitUnusedStructs,
 	}
 }
 

internal/codegen/golang/gen.go

@@ -108,6 +108,11 @@ func Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenR
 	if err != nil {
 		return nil, err
 	}
+
+	if req.Settings.Go.OmitUnusedStructs {
+		enums, structs = filterUnusedStructs(enums, structs, queries)
+	}
+
 	return generate(req, enums, structs, queries)
 }
 
@@ -288,3 +293,42 @@ func usesBatch(queries []Query) bool {
 	}
 	return false
 }
+
+func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query) ([]Enum, []Struct) {
+	keepTypes := make(map[string]struct{})
+
+	for _, query := range queries {
+		if !query.Arg.isEmpty() {
+			keepTypes[strings.TrimPrefix(query.Arg.Type(), "Null")] = struct{}{}
+			if query.Arg.IsStruct() {
+				for _, field := range query.Arg.Struct.Fields {
+					keepTypes[strings.TrimPrefix(field.Type, "Null")] = struct{}{}
+				}
+			}
+		}
+		if query.hasRetType() {
+			keepTypes[strings.TrimPrefix(query.Ret.Type(), "Null")] = struct{}{}
+			if query.Ret.IsStruct() {
+				for _, field := range query.Ret.Struct.Fields {
+					keepTypes[strings.TrimPrefix(field.Type, "Null")] = struct{}{}
+				}
+			}
+		}
+	}
+
+	keepEnums := make([]Enum, 0, len(enums))
+	for _, enum := range enums {
+		if _, ok := keepTypes[enum.Name]; ok {
+			keepEnums = append(keepEnums, enum)
+		}
+	}
+
+	keepStructs := make([]Struct, 0, len(structs))
+	for _, st := range structs {
+		if _, ok := keepTypes[st.Name]; ok {
+			keepStructs = append(keepStructs, st)
+		}
+	}
+
+	return keepEnums, keepStructs
+}

internal/config/config.go

@@ -145,6 +145,7 @@ type SQLGo struct {
 	OutputFilesSuffix           string            `json:"output_files_suffix,omitempty" yaml:"output_files_suffix"`
 	InflectionExcludeTableNames []string          `json:"inflection_exclude_table_names,omitempty" yaml:"inflection_exclude_table_names"`
 	QueryParameterLimit         *int32            `json:"query_parameter_limit,omitempty" yaml:"query_parameter_limit"`
+	OmitUnusedStructs           bool              `json:"omit_unused_structs,omitempty" yaml:"omit_unused_structs"`
 }
 
 type SQLJSON struct {

internal/config/v_one.go

@@ -49,6 +49,7 @@ type v1PackageSettings struct {
 	StrictFunctionChecks      bool       `json:"strict_function_checks" yaml:"strict_function_checks"`
 	StrictOrderBy             *bool      `json:"strict_order_by" yaml:"strict_order_by"`
 	QueryParameterLimit       *int32     `json:"query_parameter_limit,omitempty" yaml:"query_parameter_limit"`
+	OmitUnusedStructs         bool       `json:"omit_unused_structs,omitempty" yaml:"omit_unused_structs"`
 }
 
 func v1ParseConfig(rd io.Reader) (Config, error) {
@@ -168,6 +169,7 @@ func (c *V1GenerateSettings) Translate() Config {
 					OutputQuerierFileName:     pkg.OutputQuerierFileName,
 					OutputFilesSuffix:         pkg.OutputFilesSuffix,
 					QueryParameterLimit:       pkg.QueryParameterLimit,
+					OmitUnusedStructs:         pkg.OmitUnusedStructs,
 				},
 			},
 			StrictFunctionChecks: pkg.StrictFunctionChecks,

internal/endtoend/testdata/codegen_json/gen/codegen.json

@@ -41,7 +41,8 @@
       "emit_pointers_for_null_types": false,
       "query_parameter_limit": 1,
       "output_batch_file_name": "",
-      "json_tags_id_uppercase": false
+      "json_tags_id_uppercase": false,
+      "omit_unused_structs": false
     },
     "json": {
       "out": "gen",