From 96d7ffa5282f8d4fb199fcbe473fd6761f737430 Mon Sep 17 00:00:00 2001
From: Daniel Cortez <daniel@jyde.io>
Date: Wed, 23 Feb 2022 17:59:51 +0100
Subject: [PATCH] Prepend "-- " to default `query_comment`

Athena DDL does not always respect /* ... */ block quotations. This is
the case with (at least) `create external table ...` statements.

Currently, if you try to run `create external table ...` statements, the
following error is returned: `pyathena.error.OperationalError: FAILED:
ParseException line 1:0 cannot recognize input near '/' '*' '{'`

Here I monkey-patched _QueryComment.add to prepend "-- " to the
`query_comment` and replace any newlines with " ".

This allows the default `query_comment` to be added to `create external
table` statements.

We might want to do this so that dbt-labs/dbt-external-tables can add
Athena support as is proposed in dbt-labs/dbt-external-tables/pull/133

Without this modification, the only way to run `dbt run-operation
stage_external_sources --vars "ext_full_refresh: true"` without an error
is to set `query_comment:` to nothing, a specific value, or a macro.
---
 dbt/adapters/athena/__init__.py      |  1 +
 dbt/adapters/athena/query_headers.py | 27 +++++++++++++++++++++++++++
 2 files changed, 28 insertions(+)
 create mode 100644 dbt/adapters/athena/query_headers.py

diff --git a/dbt/adapters/athena/__init__.py b/dbt/adapters/athena/__init__.py
index d4d84206..aea8cc4f 100644
--- a/dbt/adapters/athena/__init__.py
+++ b/dbt/adapters/athena/__init__.py
@@ -1,6 +1,7 @@
 from dbt.adapters.athena.connections import AthenaConnectionManager
 from dbt.adapters.athena.connections import AthenaCredentials
 from dbt.adapters.athena.impl import AthenaAdapter
+import dbt.adapters.athena.query_headers
 
 from dbt.adapters.base import AdapterPlugin
 from dbt.include import athena
diff --git a/dbt/adapters/athena/query_headers.py b/dbt/adapters/athena/query_headers.py
new file mode 100644
index 00000000..d5ed74db
--- /dev/null
+++ b/dbt/adapters/athena/query_headers.py
@@ -0,0 +1,27 @@
+import dbt.adapters.base.query_headers
+
+class _QueryComment(dbt.adapters.base.query_headers._QueryComment):
+    """
+    Athena DDL does not always respect /* ... */ block quotations. 
+    This function is the same as _QueryComment.add except that 
+    a leading "-- " is prepended to the query_comment and any newlines
+    in the query_comment are replaced with " ". This allows the default
+    query_comment to be added to `create external table` statements.
+    """
+    def add(self, sql: str) -> str:
+        if not self.query_comment:
+            return sql
+
+        if self.append:
+            # replace last ';' with '<comment>;'
+            sql = sql.rstrip()
+            if sql[-1] == ";":
+                sql = sql[:-1]
+                return "{}\n-- /* {} */;".format(sql, self.query_comment.strip().replace("\n", " "))
+
+            return "{}\n-- /* {} */".format(sql, self.query_comment.strip().replace("\n", " "))
+
+        return "-- /* {} */\n{}".format(self.query_comment.strip().replace("\n", " "), sql)
+
+
+dbt.adapters.base.query_headers._QueryComment = _QueryComment