Infer CREATE/DROP SCHEMA migrations (#359)

Update the `raw_migation` event trigger function to capture `CREATE
SCHEMA` and `DROP SCHEMA` statements.

Author: andrew-farries

pkg/state/state.go

@@ -276,7 +276,11 @@ BEGIN
 		RETURN;
 	END IF;
 
-	IF tg_event = 'sql_drop' and tg_tag != 'ALTER TABLE' THEN
+  IF tg_event = 'sql_drop' AND tg_tag = 'DROP SCHEMA' THEN
+    -- Take the schema name from the drop schema command
+    SELECT object_identity INTO schemaname FROM pg_event_trigger_dropped_objects();
+
+  ELSIF tg_event = 'sql_drop' and tg_tag != 'ALTER TABLE' THEN
 		-- Guess the schema from drop commands
 		SELECT schema_name INTO schemaname FROM pg_catalog.pg_event_trigger_dropped_objects() WHERE schema_name IS NOT NULL;
 
@@ -286,7 +290,11 @@ BEGIN
 			RETURN;
 		END IF;
 
-		SELECT schema_name INTO schemaname FROM pg_catalog.pg_event_trigger_ddl_commands() WHERE schema_name IS NOT NULL;
+		IF tg_tag = 'CREATE SCHEMA' THEN
+			SELECT object_identity INTO schemaname FROM pg_event_trigger_ddl_commands();
+		ELSE
+			SELECT schema_name INTO schemaname FROM pg_catalog.pg_event_trigger_ddl_commands() WHERE schema_name IS NOT NULL;
+		END IF;
 	END IF;
 
 	IF schemaname IS NULL THEN

pkg/state/state_test.go

@@ -197,6 +197,25 @@ func TestInferredMigration(t *testing.T) {
 					},
 				},
 			},
+			{
+				name: "create/drop schema",
+				sqlStmts: []string{
+					"CREATE SCHEMA foo",
+					"DROP SCHEMA foo",
+				},
+				wantMigrations: []migrations.Migration{
+					{
+						Operations: migrations.Operations{
+							&migrations.OpRawSQL{Up: "CREATE SCHEMA foo"},
+						},
+					},
+					{
+						Operations: migrations.Operations{
+							&migrations.OpRawSQL{Up: "DROP SCHEMA foo"},
+						},
+					},
+				},
+			},
 		}
 
 		for _, tt := range tests {