Merge pull request #468 from bossenti/add-get-view-definition

implement `get_view_definition()` for AthenaDialect

Author: laughingman7743

pyathena/sqlalchemy/base.py

@@ -19,7 +19,7 @@
 )
 
 import botocore
-from sqlalchemy import exc, schema, types, util
+from sqlalchemy import exc, schema, text, types, util
 from sqlalchemy.engine import Engine, reflection
 from sqlalchemy.engine.default import DefaultDialect
 from sqlalchemy.sql.compiler import (
@@ -1102,6 +1102,20 @@ def has_table(
         except exc.NoSuchTableError:
             return False
 
+    @reflection.cache
+    def get_view_definition(
+        self, connection: Connection, view_name: str, schema: Optional[str] = None, **kw
+    ) -> str:
+        raw_connection = self._raw_connection(connection)
+        schema = schema if schema else raw_connection.schema_name  # type: ignore
+        query = f"""SHOW CREATE VIEW "{schema}"."{view_name}";"""
+        try:
+            res = connection.scalars(text(query))
+        except exc.OperationalError as e:
+            raise exc.NoSuchTableError(f"{schema}.{view_name}") from e
+        else:
+            return "\n".join([r for r in res])
+
     @reflection.cache
     def get_columns(
         self, connection: "Connection", table_name: str, schema: Optional[str] = None, **kw

tests/pyathena/sqlalchemy/test_base.py

@@ -1796,3 +1796,27 @@ def test_insert_from_select_cte_follows_insert_one(self, engine):
             ).strip()
         )
         assert actual == [(2, "bar")]
+
+    def test_get_view_definition(self, engine):
+        engine, conn = engine
+        insp = sqlalchemy.inspect(engine)
+        actual = insp.get_view_definition(schema=ENV.schema, view_name="v_one_row")
+        assert (
+            actual
+            == textwrap.dedent(
+                f"""
+            CREATE VIEW {ENV.schema}.v_one_row AS
+            SELECT number_of_rows
+            FROM
+              {ENV.schema}.one_row
+            """
+            ).strip()
+        )
+
+    def test_get_view_definition_missing_view(self, engine):
+        engine, conn = engine
+        insp = sqlalchemy.inspect(engine)
+        pytest.raises(
+            NoSuchTableError,
+            lambda: insp.get_view_definition(schema=ENV.schema, view_name="test_view"),
+        )

tests/resources/queries/create_table.sql.jinja2

@@ -122,3 +122,4 @@ CREATE EXTERNAL TABLE IF NOT EXISTS {{ schema }}.parquet_with_compression (
 STORED AS PARQUET
 LOCATION '{{ s3_staging_dir }}{{ schema }}/parquet_with_compression/'
 TBLPROPERTIES ('parquet.compress'='SNAPPY');
+CREATE OR REPLACE VIEW {{ schema }}.v_one_row AS SELECT number_of_rows FROM {{ schema }}.one_row