From 00d9852cab6f335735ae768b98be0813e506c701 Mon Sep 17 00:00:00 2001
From: Lewis Davies <lewis2106@gmail.com>
Date: Fri, 15 Dec 2023 11:20:59 +0000
Subject: [PATCH] Improve compatibiltiy with custom model aliases (#26)

---
 dbt_project.yml                      |  2 +-
 macros/raise_ref_not_found_error.sql | 12 ++++++------
 macros/ref.sql                       | 16 ++++++++++++----
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/dbt_project.yml b/dbt_project.yml
index e9394c3..f13eccf 100644
--- a/dbt_project.yml
+++ b/dbt_project.yml
@@ -1,5 +1,5 @@
 name: 'upstream_prod'
-version: '0.7.2'
+version: '0.7.3'
 config-version: 2
 
 require-dbt-version: [">=1.5.0", "<2.0.0"]
diff --git a/macros/raise_ref_not_found_error.sql b/macros/raise_ref_not_found_error.sql
index 415406b..66ad442 100644
--- a/macros/raise_ref_not_found_error.sql
+++ b/macros/raise_ref_not_found_error.sql
@@ -1,15 +1,15 @@
-{% macro raise_ref_not_found_error(current_model, relation) %}
-    {{ return(adapter.dispatch("raise_ref_not_found_error", "upstream_prod")(current_model, relation)) }}
+{% macro raise_ref_not_found_error(current_model, database, schema, relation) %}
+    {{ return(adapter.dispatch("raise_ref_not_found_error", "upstream_prod")(current_model, database, schema, relation)) }}
 {% endmacro %}
 
-{% macro default__raise_ref_not_found_error(current_model, relation) %}
+{% macro default__raise_ref_not_found_error(current_model, database, schema, relation) %}
 
     {% set error_msg -%}
 [{{ current_model }}] upstream_prod couldnt find the specified model:
 
-DATABASE: {{ relation.database }}
-SCHEMA:   {{ relation.schema }}
-RELATION: {{ relation.identifier }}
+DATABASE: {{ database }}
+SCHEMA:   {{ schema }}
+RELATION: {{ relation }}
 
 Check your variable settings in the README or create a GitHub issue for more help.
     {%- endset %}
diff --git a/macros/ref.sql b/macros/ref.sql
index e9e9445..89ef514 100644
--- a/macros/ref.sql
+++ b/macros/ref.sql
@@ -30,8 +30,8 @@
         {{ return(parent_ref) }}
     -- Try deferring to prod for non-selected upstream models
     {% else %}
+        {% set parent_node = upstream_prod.find_model_node(parent_model, version) %}
         {% if env_schemas == true %}
-            {% set parent_node = upstream_prod.find_model_node(parent_model, version) %}
             {% set custom_schema_name = parent_node.config.schema %}
             {% set parent_schema = generate_schema_name(custom_schema_name, parent_node, True) | trim %}
         -- No prod_schema = one-DB-per-env setup with same schema structure in all
@@ -43,7 +43,15 @@
         {% endif %}
         {% set parent_database = prod_database or parent_ref.database %}
 
-        {% set prod_rel = adapter.get_relation(parent_database, parent_schema, parent_ref.table) %}
+        -- Check whether the relations have been materialised in both envs
+        /***************
+        prod_rel_name helps the package find the correct prod relation for projects using a custom 
+        generate_alias_name macro. It assumes that custom aliases are only used in dev envs and prod
+        relations always have the same name as the model (+ version suffix when needed).
+        It's hacky but it seems to work. 
+        ***************/
+        {% set prod_rel_name = parent_node.path.split("/")[-1].replace(".sql", "") %}
+        {% set prod_rel = adapter.get_relation(parent_database, parent_schema, prod_rel_name) %}
         {% set dev_rel = load_relation(parent_ref) %}
         {% set prod_exists = prod_rel is not none %}
         {% set dev_exists = dev_rel is not none %}
@@ -70,10 +78,10 @@
                 {{ log("[" ~ current_model ~ "] " ~ parent_ref.table ~ " not found in prod, falling back to default target", info=True) }}
                 {{ return(dev_rel) }}
             {% else %}
-                {{ upstream_prod.raise_ref_not_found_error(current_model, dev_rel) }}
+                {{ upstream_prod.raise_ref_not_found_error(current_model, parent_ref.database, parent_ref.schema, parent_ref.identifier) }}
             {% endif %}
         {% else %}
-            {{ upstream_prod.raise_ref_not_found_error(current_model, prod_rel) }}
+            {{ upstream_prod.raise_ref_not_found_error(current_model, parent_database, parent_schema, prod_rel_name) }}
         {% endif %}
 
     {% endif %}