stager log support

elementary-data · Jan 26, 2023 · 36b8210 · 36b8210
1 parent 2776ad7
commit 36b8210
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 19 deletions.
diff --git a/edr_stager_dbt_project/macros/get_elementary_dbt_pkg_version.sql b/edr_stager_dbt_project/macros/get_elementary_dbt_pkg_version.sql
@@ -2,15 +2,15 @@
   {% set database, schema = target_database(), target.schema %}
   {% set invocations_relation = adapter.get_relation(database, schema, 'dbt_invocations') %}
   {% if not invocations_relation %}
-    {% do print('') %}
+    {% do elementary_internal.edr_stager_log('') %}
     {% do return(none) %}
   {% endif %}
 
   {% set get_pkg_version_query %}
     select elementary_version from {{ invocations_relation }} order by generated_at desc limit 1
   {% endset %}
   {% set result = dbt.run_query(get_pkg_version_query)[0][0] %}
-  {% do print(result or '') %}
+  {% do elementary_internal.edr_stager_log(result or '') %}
 {% endmacro %}
 
 {% macro target_database() -%}

diff --git a/edr_stager_dbt_project/macros/logs.sql b/edr_stager_dbt_project/macros/logs.sql
@@ -0,0 +1,5 @@
+{% macro edr_stager_log(msg, info=True) %}
+    {%- if execute %}
+        {% do log('edr_stager: ' ~ msg, info=info) %}
+    {%- endif %}
+{% endmacro %}
diff --git a/entrypoint.py b/entrypoint.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import subprocess
@@ -11,6 +12,9 @@
 logging.basicConfig(format="%(asctime)s %(message)s", level=logging.DEBUG)
 
 
+EDR_STAGER_PREFIX = "edr_stager: "
+
+
 def install_dbt(adapter: str):
     dbt_pkg_name = f"dbt-{adapter}"
     logging.info(f"Installing {dbt_pkg_name}")
@@ -36,23 +40,34 @@ def setup_env(
 def install_edr(adapter: str, project_dir: Optional[str]):
     logging.info("Getting Elementary dbt package version.")
     try:
-        dbt_pkg_ver = (
-            subprocess.run(
-                [
-                    "dbt",
-                    "-q",
-                    "run-operation",
-                    "get_elementary_dbt_pkg_version",
-                    "--project-dir",
-                    "/edr_stager_dbt_project",
-                ],
-                check=True,
-                capture_output=True,
-                cwd=project_dir,
-            )
-            .stdout.decode()
-            .strip()
-        )
+        dbt_pkg_ver = None
+        command_results = subprocess.run(
+            [
+                "dbt",
+                "--log-format",
+                "json",
+                "run-operation",
+                "get_elementary_dbt_pkg_version",
+                "--project-dir",
+                "/edr_stager_dbt_project",
+            ],
+            check=True,
+            capture_output=True,
+            cwd=project_dir,
+        ).stdout.decode("utf-8")
+
+        for log_line in command_results.splitlines():
+            try:
+                log = json.loads(log_line)
+                message = log.get("info", {}).get("msg") or log.get("data", {}).get(
+                    "msg"
+                )
+                if message.startswith(EDR_STAGER_PREFIX):
+                    dbt_pkg_ver = message[len(EDR_STAGER_PREFIX) :]
+                    break
+            except Exception:
+                pass
+
     except subprocess.CalledProcessError as err:
         logging.error(f"Failed to get Elementary dbt package version: {vars(err)}")
         raise