ability to include or exclude custom_sql files

examples/simple/ent.yml

@@ -28,4 +28,7 @@ codegen:
     alias: ExampleViewerAlias
   subscriptionType:
     path: src/graphql/resolvers/subscription_type
-    name: SubscriptionType
+    name: SubscriptionType
+databaseMigration:
+  customSQLInclude:
+    - 63ec20382c27

internal/codegen/config.go

@@ -197,6 +197,13 @@ func (cfg *Config) getCodegenConfig() *CodegenConfig {
 	return nil
 }
 
+func (cfg *Config) getDatabaseMigrationConfig() *DatabaseMigrationConfig {
+	if cfg.config != nil && cfg.config.DatabaseMigration != nil {
+		return cfg.config.DatabaseMigration
+	}
+	return nil
+}
+
 func (cfg *Config) ShouldUseRelativePaths() bool {
 	if codegen := cfg.getCodegenConfig(); codegen != nil {
 		return codegen.RelativeImports
@@ -476,6 +483,20 @@ func (cfg *Config) TransformLoadMethodX() string {
 	return fmt.Sprintf("%sX", cfg.TransformLoadMethod())
 }
 
+func (cfg *Config) CustomSQLInclude() []string {
+	if dbMigration := cfg.getDatabaseMigrationConfig(); dbMigration != nil {
+		return dbMigration.CustomSQLInclude
+	}
+	return nil
+}
+
+func (cfg *Config) CustomSQLExclude() []string {
+	if dbMigration := cfg.getDatabaseMigrationConfig(); dbMigration != nil {
+		return dbMigration.CustomSQLExclude
+	}
+	return nil
+}
+
 const DEFAULT_PRETTIER_GLOB = "src/**/*.ts"
 const PRETTIER_FILE_CHUNKS = 20
 
@@ -603,15 +624,17 @@ func parseConfig(absPathToRoot string) (*ConfigurableConfig, error) {
 }
 
 type ConfigurableConfig struct {
-	Codegen                            *CodegenConfig `yaml:"codegen"`
-	CustomGraphQLJSONPath              string         `yaml:"customGraphQLJSONPath"`
-	DynamicScriptCustomGraphQLJSONPath string         `yaml:"dynamicScriptCustomGraphQLJSONPath"`
-	GlobalSchemaPath                   string         `yaml:"globalSchemaPath"`
+	Codegen                            *CodegenConfig           `yaml:"codegen"`
+	DatabaseMigration                  *DatabaseMigrationConfig `yaml:"databaseMigration"`
+	CustomGraphQLJSONPath              string                   `yaml:"customGraphQLJSONPath"`
+	DynamicScriptCustomGraphQLJSONPath string                   `yaml:"dynamicScriptCustomGraphQLJSONPath"`
+	GlobalSchemaPath                   string                   `yaml:"globalSchemaPath"`
 }
 
 func (cfg *ConfigurableConfig) Clone() *ConfigurableConfig {
 	return &ConfigurableConfig{
 		Codegen:                            cloneCodegen(cfg.Codegen),
+		DatabaseMigration:                  cloneDatabaseMigration(cfg.DatabaseMigration),
 		CustomGraphQLJSONPath:              cfg.CustomGraphQLJSONPath,
 		DynamicScriptCustomGraphQLJSONPath: cfg.DynamicScriptCustomGraphQLJSONPath,
 		GlobalSchemaPath:                   cfg.GlobalSchemaPath,
@@ -656,6 +679,13 @@ func cloneCodegen(cfg *CodegenConfig) *CodegenConfig {
 	return cfg.Clone()
 }
 
+func cloneDatabaseMigration(cfg *DatabaseMigrationConfig) *DatabaseMigrationConfig {
+	if cfg == nil {
+		return nil
+	}
+	return cfg.Clone()
+}
+
 func (cfg *CodegenConfig) Clone() *CodegenConfig {
 	return &CodegenConfig{
 		DefaultEntPolicy:           clonePrivacyConfig(cfg.DefaultEntPolicy),
@@ -739,3 +769,15 @@ func (cfg *PrettierConfig) Clone() *PrettierConfig {
 		Glob:   cfg.Glob,
 	}
 }
+
+type DatabaseMigrationConfig struct {
+	CustomSQLInclude []string `yaml:"customSQLInclude"`
+	CustomSQLExclude []string `yaml:"customSQLExclude"`
+}
+
+func (cfg *DatabaseMigrationConfig) Clone() *DatabaseMigrationConfig {
+	return &DatabaseMigrationConfig{
+		CustomSQLInclude: cfg.CustomSQLInclude,
+		CustomSQLExclude: cfg.CustomSQLExclude,
+	}
+}

internal/db/db_schema.go

@@ -759,7 +759,7 @@ func RunAlembicCommand(cfg *codegen.Config, command string, args ...string) erro
 }
 
 func (s *dbSchema) writeSchemaFile(cfg *codegen.Config) error {
-	data, err := s.getSchemaForTemplate()
+	data, err := s.getSchemaForTemplate(cfg)
 	if err != nil {
 		return err
 	}
@@ -774,8 +774,10 @@ func (s *dbSchema) writeSchemaFile(cfg *codegen.Config) error {
 	)
 }
 
-func (s *dbSchema) getSchemaForTemplate() (*dbSchemaTemplate, error) {
-	ret := &dbSchemaTemplate{}
+func (s *dbSchema) getSchemaForTemplate(cfg *codegen.Config) (*dbSchemaTemplate, error) {
+	ret := &dbSchemaTemplate{
+		Config: cfg,
+	}
 
 	for _, table := range s.Tables {
 
@@ -1456,6 +1458,7 @@ type dbDataInfo struct {
 
 // wrapper object to represent the list of tables that will be passed to a schema template file
 type dbSchemaTemplate struct {
+	Config *codegen.Config
 	Tables []dbSchemaTableInfo
 	Edges  []dbEdgeInfo
 	Data   []dbDataInfo

internal/db/db_schema.tmpl

@@ -45,5 +45,24 @@ metadata.info["data"] = {
 }
 {{end}}
 
+{{ if .Config.CustomSQLInclude -}}
+  metadata.info["custom_sql_include"] = {
+    'public': [
+      {{ range .Config.CustomSQLInclude -}}
+        '{{.}}',
+      {{end -}}
+    ],
+  }
+{{ end -}}
+{{ if .Config.CustomSQLExclude -}}
+  metadata.info["custom_sql_exclude"] = {
+    'public': [
+      {{ range .Config.CustomSQLExclude -}}
+        '{{.}}',
+      {{end -}}
+    ],
+  }
+{{ end}}
+
 def get_metadata():
   return metadata

python/auto_schema/auto_schema/runner.py

@@ -305,6 +305,27 @@ def squash(self, squash):
 
 
     def _get_custom_sql(self, connection, dialect) -> io.StringIO:
+        custom_sql_include_all = config.metadata.info.setdefault('custom_sql_include', {})
+        custom_sql_exclude_all= config.metadata.info.setdefault('custom_sql_exclude', {})
+        custom_sql_include_list = custom_sql_include_all.setdefault('public', [])
+        custom_sql_exclude_list = custom_sql_exclude_all.setdefault('public', [])
+
+        custom_sql_include = None
+        custom_sql_exclude = None
+
+        if len(custom_sql_include_list) > 0:
+            custom_sql_include = set(custom_sql_include_list)
+        if len(custom_sql_exclude_list) > 0:
+            custom_sql_exclude = set(custom_sql_exclude_list)
+
+        def include_rev(name):
+            if custom_sql_include is not None and name in custom_sql_include:
+                return True
+            if custom_sql_exclude is not None and name in custom_sql_exclude:
+                return False
+            return True
+
+
         script_directory = self.cmd.get_script_directory()
         revs = script_directory.walk_revisions()
 
@@ -347,7 +368,7 @@ def wrapper(self, *args, **kwargs):
                 # run upgrade(), we capture what's being changed via the dispatcher and see if it's custom sql 
                 rev.module.upgrade()
 
-                if isinstance(last_obj, ops.ExecuteSQL) or isinstance(last_obj, alembicops.ExecuteSQLOp):
+                if (isinstance(last_obj, ops.ExecuteSQL) or isinstance(last_obj, alembicops.ExecuteSQLOp) and include_rev(rev.revision)):
                     custom_sql_buffer.write("-- custom sql for rev %s\n" % rev.revision)
                     custom_sql_buffer.write(temp_buffer.getvalue())
 

python/auto_schema/setup.py

@@ -8,7 +8,7 @@
 setuptools.setup(
     name="auto_schema_test",  # auto_schema_test to test
     # 0.0.29 production
-    version="0.0.26",  # 0.0.26 was last test version
+    version="0.0.27",  # 0.0.26 was last test version
     author="Ola Okelola",
     author_email="[email protected]",
     description="auto schema for a db",

ts/src/core/config.ts

@@ -48,6 +48,8 @@ export interface ConfigWithCodegen extends Config {
   // config for codegen
   codegen?: CodegenConfig;
 
+  databaseMigration?: DatabaseMigrationConfig;
+
   // because of swc issues, we might not be able to
   // parse custom graphql via decorators so we put this
   // in a json file for now
@@ -64,6 +66,11 @@ export interface ConfigWithCodegen extends Config {
   globalSchemaPath?: string;
 }
 
+interface DatabaseMigrationConfig {
+  custom_sql_include?: string[];
+  custom_sql_exclude?: string[];
+}
+
 interface CodegenConfig {
   defaultEntPolicy?: PrivacyConfig;
   defaultActionPolicy?: PrivacyConfig;