Repository: dbt-labs/dbt-core

Issue:
[CT-1034] [Bug] Unexpected behavior when chaining methods on dbt-ref'ed/sourced dataframes
### Is this a new bug in dbt-core?

- [X] I believe this is a new bug in dbt-core
- [X] I have searched the existing issues, and I could not find an existing issue for this bug

### Current Behavior

Community Slack report here:

https://getdbt.slack.com/archives/C03QUA7DWCW/p1660223963454249

I reproduced easily -- the first run had `df = df.limit(100)` line uncommented and no chained call on the prior line, the second run is what's shown in the screenshot:

![image](https://user-images.githubusercontent.com/54814569/184193825-e0afb6e2-03bb-452d-b5ab-b6f15f4fd3ad.png)


### Expected Behavior

Would not expect method chaining to cause errors like this on `dbt.ref(...)` or `dbt.source(...)`. We should clearly document this behavior and give a better error method if there's some underlying reason why we can't easily fix this, otherwise we should fix it.

### Steps To Reproduce

See screenshot above. Basically add a `.limit(N)` to your reference call, e.g. `df = dbt.ref(model_name).limit(100)`.

### Relevant log output

```shell
@lostmygithubaccount ➜ /workspaces/codyspace/nyc_taxi_with_python (main ✗) $ dbt run -s int_yellow 
17:10:04  Running with dbt=1.3.0-b1
17:10:04  Found 5 models, 0 tests, 0 snapshots, 1 analysis, 571 macros, 0 operations, 0 seed files, 1 source, 0 exposures, 0 metrics
17:10:04  
17:10:05  Concurrency: 8 threads (target='dev')
17:10:05  
17:10:05  1 of 1 START python table model dbt_cody.int_yellow ............................ [RUN]
17:10:10  1 of 1 ERROR creating python table model dbt_cody.int_yellow ................... [ERROR in 4.58s]
17:10:10  
17:10:10  Finished running 1 table model in 0 hours 0 minutes and 5.57 seconds (5.57s).
17:10:10  
17:10:10  Completed with 1 error and 0 warnings:
17:10:10  
17:10:10  Database Error in model int_yellow (models/intermediate/int_yellow.py)
17:10:10    100357 (P0000): Python Interpreter Error:
17:10:10    Traceback (most recent call last):
17:10:10      File "_udf_code.py", line 83, in main
17:10:10      File "_udf_code.py", line 7, in model
17:10:10      File "_udf_code.py", line 53, in <lambda>
17:10:10      File "_udf_code.py", line 21, in ref
17:10:10    KeyError: 'stg_yellow'
17:10:10     in function INT_YELLOW__DBT_SP with handler main
17:10:10    compiled Code at target/run/nyc_taxi_with_python/models/intermediate/int_yellow.py
17:10:10  
17:10:10  Done. PASS=0 WARN=0 ERROR=1 SKIP=0 TOTAL=1
```


### Environment

```markdown
- OS: linux
- Python: 3.10
- dbt: 1.3 beta, installed via main branches
```


### Which database adapter are you using with dbt?

snowflake

### Additional Context

cc: @ChenyuLInx 

_No response_


Discussion:
@lostmygithubaccount  I looked into this issue a bit and noticed that this is related to how we traverse the parsed AST. 
We can go with either options you purposed, through proper message is slightly easier. But properly fix this is totally doable, we can switch from just traverse the AST tree of python models to also use regex to parse refs and sources. Will estimate this in the upcoming BLG and we can decide which way we want to go.

Note to myself: some example regex for `dbt.ref`
```
(dbt\.ref\()[^\)\}]*(\))
```
Without having tested at all locally — I'd guess that the regex approach risks being significantly slower, and potentially liable to abuse, versus using Python's AST. I'd lean slightly in the direction of just documenting that `dbt.ref()` + `dbt.source()` cannot be chained with other method calls. It's an extra line of code for the user, not a big usability issue.
As the one who experienced and reported the issue, I personally think that it's preferable to give a technical solution here.
It's probably going to confuse a lot of users and that they _will_ stumble upon it given that `dbt.ref()` and `dbt.source()` are among the most used functions that will be used in Python models.
@jtcohen6 thinking again it should be possible to do all of this using the AST NodeVisitor. When I get to it, if it is really tricky to support syntax like `dbt.ref(model_name).limit(100)` or `fancy_transform(dbt.ref(model_name))`, I will make sure we raise a clear error.
@elongl definitely agree that'd be ideal, but to @jtcohen6's point we should be conscious of the implications. will defer to @ChenyuLInx to further investigate the options and performance/other technical impacts before we make a call. 

@elongl appreciate you reporting this! it is something that's going to trip people up and good to catch early in the beta
As part of estimation we're including doing a very basic benchmark of doing a full ast parse (as opposed to the top-level we do currently)

Resolving patch:
diff --git a/core/dbt/parser/models.py b/core/dbt/parser/models.py
--- a/core/dbt/parser/models.py
+++ b/core/dbt/parser/models.py
@@ -32,6 +32,10 @@
 from dbt.exceptions import ParsingException, validator_error_message, UndefinedMacroException
 
 
+dbt_function_key_words = set(["ref", "source", "config", "get"])
+dbt_function_full_names = set(["dbt.ref", "dbt.source", "dbt.config", "dbt.config.get"])
+
+
 class PythonValidationVisitor(ast.NodeVisitor):
     def __init__(self):
         super().__init__()
@@ -108,13 +112,40 @@ def _get_call_literals(self, node):
         return arg_literals, kwarg_literals
 
     def visit_Call(self, node: ast.Call) -> None:
-
-        func_name = self._flatten_attr(node.func)
-        if func_name in ["dbt.ref", "dbt.source", "dbt.config", "dbt.config.get"]:
-            # drop the dot-dbt prefix
-            func_name = func_name.split(".")[-1]
-            args, kwargs = self._get_call_literals(node)
-            self.dbt_function_calls.append((func_name, args, kwargs))
+        # check weather the current call could be a dbt function call
+        if isinstance(node.func, ast.Attribute) and node.func.attr in dbt_function_key_words:
+            func_name = self._flatten_attr(node.func)
+            # check weather the current call really is a dbt function call
+            if func_name in dbt_function_full_names:
+                # drop the dot-dbt prefix
+                func_name = func_name.split(".")[-1]
+                args, kwargs = self._get_call_literals(node)
+                self.dbt_function_calls.append((func_name, args, kwargs))
+
+        # no matter what happened above, we should keep visiting the rest of the tree
+        # visit args and kwargs to see if there's call in it
+        for obj in node.args + [kwarg.value for kwarg in node.keywords]:
+            if isinstance(obj, ast.Call):
+                self.visit_Call(obj)
+            # support dbt.ref in list args, kwargs
+            elif isinstance(obj, ast.List) or isinstance(obj, ast.Tuple):
+                for el in obj.elts:
+                    if isinstance(el, ast.Call):
+                        self.visit_Call(el)
+            # support dbt.ref in dict args, kwargs
+            elif isinstance(obj, ast.Dict):
+                for value in obj.values:
+                    if isinstance(value, ast.Call):
+                        self.visit_Call(value)
+        # visit node.func.value if we are at an call attr
+        if isinstance(node.func, ast.Attribute):
+            self.attribute_helper(node.func)
+
+    def attribute_helper(self, node: ast.Attribute) -> None:
+        while isinstance(node, ast.Attribute):
+            node = node.value  # type: ignore
+        if isinstance(node, ast.Call):
+            self.visit_Call(node)
 
     def visit_Import(self, node: ast.Import) -> None:
         for n in node.names:


Test patch:
diff --git a/test/unit/test_parser.py b/test/unit/test_parser.py
--- a/test/unit/test_parser.py
+++ b/test/unit/test_parser.py
@@ -549,9 +549,16 @@ def model(dbt, session):
     from torch import b
     import textblob.text
     import sklearn
-    df = dbt.ref("my_sql_model")
-    df = dbt.ref("my_sql_model")
-    df = dbt.ref("my_sql_model_2")
+    df0 = pandas(dbt.ref("a_model"))
+    df1 = dbt.ref("my_sql_model").task.limit(2)
+    df2 = dbt.ref("my_sql_model_1")
+    df3 = dbt.ref("my_sql_model_2")
+    df4 = dbt.source("test", 'table1').limit(max = [max(dbt.ref('something'))])
+    df5 = [dbt.ref('test1')]
+    
+    a_dict = {'test2' : dbt.ref('test2')}
+    df5 = anotherfunction({'test2' : dbt.ref('test3')})
+    df6 = [somethingelse.ref(dbt.ref("test4"))]
 
     df = df.limit(2)
     return df   
@@ -582,7 +589,8 @@ def model(dbt, session):
             checksum=block.file.checksum,
             unrendered_config={'materialized': 'table', 'packages':python_packages},
             config_call_dict={'materialized': 'table', 'packages':python_packages},
-            refs=[['my_sql_model'], ['my_sql_model'], ['my_sql_model_2']]
+            refs=[['a_model'], ['my_sql_model'], ['my_sql_model_1'], ['my_sql_model_2'], ['something'], ['test1'], ['test2'], ['test3'], ['test4']],
+            sources = [['test', 'table1']],
         )
         assertEqualNodes(node, expected)
         file_id = 'snowplow://' + normalize('models/nested/py_model.py')