From 87178287c7869985dd0c046a5bf6b0bbac763edd Mon Sep 17 00:00:00 2001
From: Chenyu Li <chenyu.li@dbtlabs.com>
Date: Wed, 31 Jan 2024 14:38:40 -0800
Subject: [PATCH] Fix python model parsing from meder (#9162)

* fix: parsing f-strings in python models

* add unittest

* changelong

* pre-commit

---------

Co-authored-by: Meder Kamalov <meder@kamalov.net>
Co-authored-by: Kshitij Aranke <kshitij.aranke@dbtlabs.com>
---
 .../unreleased/Fixes-20231128-102111.yaml     |  6 ++++++
 core/dbt/parser/models.py                     |  6 ++++++
 tests/unit/test_parser.py                     | 21 +++++++++++++++++++
 3 files changed, 33 insertions(+)
 create mode 100644 .changes/unreleased/Fixes-20231128-102111.yaml

diff --git a/.changes/unreleased/Fixes-20231128-102111.yaml b/.changes/unreleased/Fixes-20231128-102111.yaml
new file mode 100644
index 00000000000..c3371708a11
--- /dev/null
+++ b/.changes/unreleased/Fixes-20231128-102111.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Fix parsing f-strings in python models
+time: 2023-11-28T10:21:11.596121-08:00
+custom:
+  Author: mederka
+  Issue: "6976"
diff --git a/core/dbt/parser/models.py b/core/dbt/parser/models.py
index b5cbba5583a..6092140839b 100644
--- a/core/dbt/parser/models.py
+++ b/core/dbt/parser/models.py
@@ -134,6 +134,12 @@ def visit_Call(self, node: ast.Call) -> None:
                 for value in obj.values:
                     if isinstance(value, ast.Call):
                         self.visit_Call(value)
+            # support dbt function calls in f-strings
+            elif isinstance(obj, ast.JoinedStr):
+                for value in obj.values:
+                    if isinstance(value, ast.FormattedValue) and isinstance(value.value, ast.Call):
+                        self.visit_Call(value.value)
+
         # visit node.func.value if we are at an call attr
         if isinstance(node.func, ast.Attribute):
             self.attribute_helper(node.func)
diff --git a/tests/unit/test_parser.py b/tests/unit/test_parser.py
index ef750cb5915..c949756eba8 100644
--- a/tests/unit/test_parser.py
+++ b/tests/unit/test_parser.py
@@ -798,6 +798,7 @@ def model(dbt, session):
     a_dict = {'test2': dbt.ref('test2')}
     df5 = {'test2': dbt.ref('test3')}
     df6 = [dbt.ref("test4")]
+    f"{dbt.ref('test5')}"
 
     df = df0.limit(2)
     return df
@@ -861,6 +862,17 @@ def model(dbt, session):
     return dbt.ref("some_model"), dbt.ref("some_other_model")
 """
 
+python_model_f_string = """
+# my_python_model.py
+import pandas as pd
+
+def model(dbt, fal):
+    dbt.config(materialized="table")
+    print(f"my var: {dbt.config.get('my_var')}") # Prints "my var: None"
+    df: pd.DataFrame = dbt.ref("some_model")
+    return df
+"""
+
 python_model_no_return = """
 def model(dbt, session):
     dbt.config(materialized='table')
@@ -985,6 +997,7 @@ def test_python_model_parse(self):
                 RefArgs("test2"),
                 RefArgs("test3"),
                 RefArgs("test4"),
+                RefArgs("test5"),
             ],
             sources=[["test", "table1"]],
         )
@@ -1001,6 +1014,14 @@ def test_python_model_config(self):
         node = list(self.parser.manifest.nodes.values())[0]
         self.assertEqual(node.config.to_dict()["config_keys_used"], ["param_1", "param_2"])
 
+    def test_python_model_f_string_config(self):
+        block = self.file_block_for(python_model_f_string, "nested/py_model.py")
+        self.parser.manifest.files[block.file.file_id] = block.file
+
+        self.parser.parse_file(block)
+        node = list(self.parser.manifest.nodes.values())[0]
+        self.assertEqual(node.config.to_dict()["config_keys_used"], ["my_var"])
+
     def test_python_model_config_with_defaults(self):
         block = self.file_block_for(python_model_config_with_defaults, "nested/py_model.py")
         self.parser.manifest.files[block.file.file_id] = block.file