Repository: dbt-labs/dbt-core

Issue:
[PP] KeyError in delete_schema_mssa_links 
From [Slack thread](https://getdbt.slack.com/archives/C2JRRQDTL/p1626380641235500):

> Yes I was editing a test that applied to a source. I only ran into this issue for custom generic tests that apply to sources

```
2021-07-15 20:41:44.027209 (MainThread): Encountered an error:
2021-07-15 20:41:44.027763 (MainThread): 'sources'
2021-07-15 20:41:44.032374 (MainThread): Traceback (most recent call last):
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/main.py", line 125, in main
    results, succeeded = handle_and_check(args)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/main.py", line 203, in handle_and_check
    task, res = run_from_args(parsed)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/main.py", line 256, in run_from_args
    results = task.run()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/task/runnable.py", line 406, in run
    self._runtime_initialize()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/task/runnable.py", line 122, in _runtime_initialize
    super()._runtime_initialize()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/task/runnable.py", line 79, in _runtime_initialize
    self.load_manifest()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/task/runnable.py", line 66, in load_manifest
    self.manifest = ManifestLoader.get_full_manifest(self.config)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/manifest.py", line 170, in get_full_manifest
    manifest = loader.load()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/manifest.py", line 212, in load
    project_parser_files = partial_parsing.get_parsing_files()
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 131, in get_parsing_files
    self.update_in_saved(file_id)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 201, in update_in_saved
    self.update_macro_in_saved(new_source_file, old_source_file)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 254, in update_macro_in_saved
    self.handle_macro_file_links(old_source_file, follow_references=True)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 360, in handle_macro_file_links
    self.schedule_macro_nodes_for_parsing(referencing_nodes)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 390, in schedule_macro_nodes_for_parsing
    self.delete_schema_mssa_links(schema_file, key, node_patch)
  File "/Users/<user>/Library/Python/3.7/lib/python/site-packages/dbt/parser/partial.py", line 578, in delete_schema_mssa_links
    prefix = key_to_prefix[dict_key]
KeyError: 'sources' 
```


Discussion:
This code should not be reached for sources, only for model/seed/snapshot/analysis files. Will need to figure out how it got there.
The code to remove related nodes (etc) when test macros are re-done was not handling source tests.

Resolving patch:
diff --git a/core/dbt/parser/partial.py b/core/dbt/parser/partial.py
--- a/core/dbt/parser/partial.py
+++ b/core/dbt/parser/partial.py
@@ -404,12 +404,21 @@ def schedule_macro_nodes_for_parsing(self, unique_ids):
                         patch_list = []
                         if key in schema_file.dict_from_yaml:
                             patch_list = schema_file.dict_from_yaml[key]
-                        node_patch = self.get_schema_element(patch_list, name)
-                        if node_patch:
-                            self.delete_schema_mssa_links(schema_file, key, node_patch)
-                            self.merge_patch(schema_file, key, node_patch)
-                            if unique_id in schema_file.node_patches:
-                                schema_file.node_patches.remove(unique_id)
+                        patch = self.get_schema_element(patch_list, name)
+                        if patch:
+                            if key in ['models', 'seeds', 'snapshots']:
+                                self.delete_schema_mssa_links(schema_file, key, patch)
+                                self.merge_patch(schema_file, key, patch)
+                                if unique_id in schema_file.node_patches:
+                                    schema_file.node_patches.remove(unique_id)
+                            elif key == 'sources':
+                                # re-schedule source
+                                if 'overrides' in patch:
+                                    # This is a source patch; need to re-parse orig source
+                                    self.remove_source_override_target(patch)
+                                self.delete_schema_source(schema_file, patch)
+                                self.remove_tests(schema_file, 'sources', patch['name'])
+                                self.merge_patch(schema_file, 'sources', patch)
                 else:
                     file_id = node.file_id
                     if file_id in self.saved_files and file_id not in self.file_diff['deleted']:


Test patch:
diff --git a/test/integration/068_partial_parsing_tests/extra-files/schema-sources4.yml b/test/integration/068_partial_parsing_tests/extra-files/schema-sources4.yml
--- a/test/integration/068_partial_parsing_tests/extra-files/schema-sources4.yml
+++ b/test/integration/068_partial_parsing_tests/extra-files/schema-sources4.yml
@@ -11,6 +11,7 @@ sources:
               - not_null:
                   severity: "{{ 'error' if target.name == 'prod' else 'warn' }}"
               - unique
+              - every_value_is_blue
           - name: first_name
           - name: last_name
           - name: email
diff --git a/test/integration/068_partial_parsing_tests/extra-files/sources-tests1.sql b/test/integration/068_partial_parsing_tests/extra-files/sources-tests1.sql
new file mode 100644
--- /dev/null
+++ b/test/integration/068_partial_parsing_tests/extra-files/sources-tests1.sql
@@ -0,0 +1,9 @@
+
+{% test every_value_is_blue(model, column_name) %}
+
+    select *
+    from {{ model }}
+    where {{ column_name }} = 9999
+
+{% endtest %}
+
diff --git a/test/integration/068_partial_parsing_tests/extra-files/sources-tests2.sql b/test/integration/068_partial_parsing_tests/extra-files/sources-tests2.sql
new file mode 100644
--- /dev/null
+++ b/test/integration/068_partial_parsing_tests/extra-files/sources-tests2.sql
@@ -0,0 +1,9 @@
+
+{% test every_value_is_blue(model, column_name) %}
+
+    select *
+    from {{ model }}
+    where {{ column_name }} != 99
+
+{% endtest %}
+
diff --git a/test/integration/068_partial_parsing_tests/macros-b/.gitignore b/test/integration/068_partial_parsing_tests/macros-b/.gitignore
new file mode 100644
--- /dev/null
+++ b/test/integration/068_partial_parsing_tests/macros-b/.gitignore
@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore
diff --git a/test/integration/068_partial_parsing_tests/test_partial_parsing.py b/test/integration/068_partial_parsing_tests/test_partial_parsing.py
--- a/test/integration/068_partial_parsing_tests/test_partial_parsing.py
+++ b/test/integration/068_partial_parsing_tests/test_partial_parsing.py
@@ -211,6 +211,7 @@ def project_config(self):
             'config-version': 2,
             'data-paths': ['seed'],
             'test-paths': ['tests'],
+            'macro-paths': ['macros-b'],
             'analysis-paths': ['analysis'],
             'seeds': {
                 'quote_columns': False,
@@ -237,6 +238,8 @@ def tearDown(self):
             os.remove(normalize('tests/my_test.sql'))
         if os.path.exists(normalize('analysis/my_analysis.sql')):
             os.remove(normalize('analysis/my_analysis.sql'))
+        if os.path.exists(normalize('macros-b/tests.sql')):
+            os.remove(normalize('macros-b/tests.sql'))
 
 
     @use_profile('postgres')
@@ -244,6 +247,7 @@ def test_postgres_pp_sources(self):
         # initial run
         self.run_dbt(['clean'])
         shutil.copyfile('extra-files/raw_customers.csv', 'seed/raw_customers.csv')
+        shutil.copyfile('extra-files/sources-tests1.sql', 'macros-b/tests.sql')
         results = self.run_dbt(["run"])
         self.assertEqual(len(results), 1)
 
@@ -324,7 +328,7 @@ def test_postgres_pp_sources(self):
         shutil.copyfile('extra-files/my_test.sql', 'tests/my_test.sql')
         results = self.run_dbt(["--partial-parse", "test"])
         manifest = get_manifest()
-        self.assertEqual(len(manifest.nodes), 8)
+        self.assertEqual(len(manifest.nodes), 9)
         test_id = 'test.test.my_test'
         self.assertIn(test_id, manifest.nodes)
 
@@ -337,13 +341,17 @@ def test_postgres_pp_sources(self):
         os.remove(normalize('tests/my_test.sql'))
         results = self.run_dbt(["--partial-parse", "test"])
         manifest = get_manifest()
-        self.assertEqual(len(manifest.nodes), 8)
+        self.assertEqual(len(manifest.nodes), 9)
 
         # Remove analysis
         os.remove(normalize('analysis/my_analysis.sql'))
         results = self.run_dbt(["--partial-parse", "run"])
         manifest = get_manifest()
-        self.assertEqual(len(manifest.nodes), 7)
+        self.assertEqual(len(manifest.nodes), 8)
+
+        # Change source test
+        shutil.copyfile('extra-files/sources-tests2.sql', 'macros-b/tests.sql')
+        results = self.run_dbt(["--partial-parse", "run"])
 
 
 class TestPartialParsingDependency(DBTIntegrationTest):