Repository: dbt-athena/dbt-athena (python)

Issue:
Error running dbt docs generate with Federated Query database
I'm trying to run `dbt docs generate` with Federated Query data using the Redshift Connector and it returns the following error:

```
[0m15:45:16.272788 [debug] [ThreadPool]: Acquiring new athena connection 'awsdatacatalog.information_schema'
[0m15:45:16.273878 [debug] [ThreadPool]: Acquiring new athena connection 'redshiftCatalog_dev.information_schema'
[0m15:45:16.274842 [debug] [ThreadPool]: Opening a new connection, currently in state init
[0m15:45:16.275793 [debug] [ThreadPool]: Opening a new connection, currently in state init
[0m15:45:17.416799 [debug] [ThreadPool]: On redshiftCatalog_dev.information_schema: Close
[0m15:45:17.419386 [warn ] [MainThread]: Encountered an error while generating catalog: 'catalog-id'
[0m15:45:20.555092 [debug] [ThreadPool]: On awsdatacatalog.information_schema: Close
[0m15:45:20.707223 [error] [MainThread]: dbt encountered 1 failure while writing the catalog
[0m15:45:20.709353 [info ] [MainThread]: Catalog written to /app/target/catalog.json
```

I'm aware of #247 but believe this is an issue with using non-glue data catalogs.

The catalog file and html are generated and aparently look fine, but the error is returned nonetheless.




Pull request:
fix: docs generate support federated query catalogs
### Description

<!--- Add a little description on what you plan to tackle with this PR -->
<!--- Add resolve #issue-number in case you resolve an open issue -->
Resolve #306 

Updated `_get_one_catalog` to use Glue API only if catalog is of type `GLUE`, otherwise it runs the `get_catalog` macro which was changed back to query information_schema.

## Models used to test - Optional
<!--- Add here the models that you use to test the changes -->


## Checklist
- [ ] You followed [contributing section](https://github.com/dbt-athena/dbt-athena#contributing)
- [ ] You kept your Pull Request small and focused on a single feature or bug fix.
- [ ] You added unit testing when necessary
- [ ] You added functional testing when necessary


Patch:
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index 4d6dccc..676ec2d 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -44,7 +44,13 @@ from dbt.adapters.athena.relation import (
     get_table_type,
 )
 from dbt.adapters.athena.s3 import S3DataNaming
-from dbt.adapters.athena.utils import clean_sql_comment, get_catalog_id, get_chunks
+from dbt.adapters.athena.utils import (
+    AthenaCatalogType,
+    clean_sql_comment,
+    get_catalog_id,
+    get_catalog_type,
+    get_chunks,
+)
 from dbt.adapters.base import ConstraintSupport, available
 from dbt.adapters.base.relation import BaseRelation, InformationSchema
 from dbt.adapters.sql import SQLAdapter
@@ -411,6 +417,29 @@ class AthenaAdapter(SQLAdapter):
             for idx, col in enumerate(table["StorageDescriptor"]["Columns"] + table.get("PartitionKeys", []))
         ]
 
+    def _get_one_table_for_non_glue_catalog(
+        self, table: TableTypeDef, schema: str, database: str
+    ) -> List[Dict[str, Any]]:
+        table_catalog = {
+            "table_database": database,
+            "table_schema": schema,
+            "table_name": table["Name"],
+            "table_type": RELATION_TYPE_MAP[table.get("TableType", "EXTERNAL_TABLE")].value,
+            "table_comment": table.get("Parameters", {}).get("comment", ""),
+        }
+        return [
+            {
+                **table_catalog,
+                **{
+                    "column_name": col["Name"],
+                    "column_index": idx,
+                    "column_type": col["Type"],
+                    "column_comment": col.get("Comment", ""),
+                },
+            }
+            for idx, col in enumerate(table["Columns"] + table.get("PartitionKeys", []))
+        ]
+
     def _get_one_catalog(
         self,
         information_schema: InformationSchema,
@@ -418,29 +447,55 @@ class AthenaAdapter(SQLAdapter):
         manifest: Manifest,
     ) -> agate.Table:
         data_catalog = self._get_data_catalog(information_schema.path.database)
-        catalog_id = get_catalog_id(data_catalog)
+        data_catalog_type = get_catalog_type(data_catalog)
+
         conn = self.connections.get_thread_connection()
         client = conn.handle
-        with boto3_client_lock:
-            glue_client = client.session.client("glue", region_name=client.region_name, config=get_boto3_config())
-
-        catalog = []
-        paginator = glue_client.get_paginator("get_tables")
-        for schema, relations in schemas.items():
-            kwargs = {
-                "DatabaseName": schema,
-                "MaxResults": 100,
-            }
-            # If the catalog is `awsdatacatalog` we don't need to pass CatalogId as boto3 infers it from the account Id.
-            if catalog_id:
-                kwargs["CatalogId"] = catalog_id
+        if data_catalog_type == AthenaCatalogType.GLUE:
+            with boto3_client_lock:
+                glue_client = client.session.client("glue", region_name=client.region_name, config=get_boto3_config())
+
+            catalog = []
+            paginator = glue_client.get_paginator("get_tables")
+            for schema, relations in schemas.items():
+                kwargs = {
+                    "DatabaseName": schema,
+                    "MaxResults": 100,
+                }
+                # If the catalog is `awsdatacatalog` we don't need to pass CatalogId as boto3
+                # infers it from the account Id.
+                catalog_id = get_catalog_id(data_catalog)
+                if catalog_id:
+                    kwargs["CatalogId"] = catalog_id
+
+                for page in paginator.paginate(**kwargs):
+                    for table in page["TableList"]:
+                        if relations and table["Name"] in relations:
+                            catalog.extend(self._get_one_table_for_catalog(table, information_schema.path.database))
+            table = agate.Table.from_object(catalog)
+        else:
+            with boto3_client_lock:
+                athena_client = client.session.client(
+                    "athena", region_name=client.region_name, config=get_boto3_config()
+                )
 
-            for page in paginator.paginate(**kwargs):
-                for table in page["TableList"]:
-                    if relations and table["Name"] in relations:
-                        catalog.extend(self._get_one_table_for_catalog(table, information_schema.path.database))
+            catalog = []
+            paginator = athena_client.get_paginator("list_table_metadata")
+            for schema, relations in schemas.items():
+                for page in paginator.paginate(
+                    CatalogName=information_schema.path.database,
+                    DatabaseName=schema,
+                    MaxResults=50,  # Limit supported by this operation
+                ):
+                    for table in page["TableMetadataList"]:
+                        if relations and table["Name"] in relations:
+                            catalog.extend(
+                                self._get_one_table_for_non_glue_catalog(
+                                    table, schema, information_schema.path.database
+                                )
+                            )
+            table = agate.Table.from_object(catalog)
 
-        table = agate.Table.from_object(catalog)
         filtered_table = self._catalog_filter_table(table, manifest)
         return self._join_catalog_table_owners(filtered_table, manifest)
 
diff --git a/dbt/adapters/athena/relation.py b/dbt/adapters/athena/relation.py
index bb07393..680f9e0 100644
--- a/dbt/adapters/athena/relation.py
+++ b/dbt/adapters/athena/relation.py
@@ -79,6 +79,7 @@ class AthenaSchemaSearchMap(Dict[InformationSchema, Dict[str, Set[Optional[str]]
 
 RELATION_TYPE_MAP = {
     "EXTERNAL_TABLE": TableType.TABLE,
+    "EXTERNAL": TableType.TABLE,  # type returned by federated query tables
     "MANAGED_TABLE": TableType.TABLE,
     "VIRTUAL_VIEW": TableType.VIEW,
     "table": TableType.TABLE,
diff --git a/dbt/adapters/athena/utils.py b/dbt/adapters/athena/utils.py
index 778fb4c..dcd7491 100644
--- a/dbt/adapters/athena/utils.py
+++ b/dbt/adapters/athena/utils.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Generator, List, Optional, TypeVar
 
 from mypy_boto3_athena.type_defs import DataCatalogTypeDef
@@ -9,7 +10,17 @@ def clean_sql_comment(comment: str) -> str:
 
 
 def get_catalog_id(catalog: Optional[DataCatalogTypeDef]) -> Optional[str]:
-    return catalog["Parameters"]["catalog-id"] if catalog else None
+    return catalog["Parameters"]["catalog-id"] if catalog and catalog["Type"] == AthenaCatalogType.GLUE.value else None
+
+
+class AthenaCatalogType(Enum):
+    GLUE = "GLUE"
+    LAMBDA = "LAMBDA"
+    HIVE = "HIVE"
+
+
+def get_catalog_type(catalog: Optional[DataCatalogTypeDef]) -> Optional[AthenaCatalogType]:
+    return AthenaCatalogType(catalog["Type"]) if catalog else None
 
 
 T = TypeVar("T")


Test patch:
diff --git a/tests/unit/constants.py b/tests/unit/constants.py
index 0513ab6..8fe9aa8 100644
--- a/tests/unit/constants.py
+++ b/tests/unit/constants.py
@@ -1,6 +1,7 @@
 CATALOG_ID = "12345678910"
 DATA_CATALOG_NAME = "awsdatacatalog"
 SHARED_DATA_CATALOG_NAME = "9876543210"
+FEDERATED_QUERY_CATALOG_NAME = "federated_query_data_source"
 DATABASE_NAME = "test_dbt_athena"
 BUCKET = "test-dbt-athena"
 AWS_REGION = "eu-west-1"
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index 30d39e4..f14ae52 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -4,7 +4,10 @@ from unittest.mock import patch
 
 import agate
 import boto3
+import botocore
 import pytest
+
+# from botocore.client.BaseClient import _make_api_call
 from moto import mock_athena, mock_glue, mock_s3, mock_sts
 from moto.core import DEFAULT_ACCOUNT_ID
 
@@ -14,6 +17,7 @@ from dbt.adapters.athena.column import AthenaColumn
 from dbt.adapters.athena.connections import AthenaCursor, AthenaParameterFormatter
 from dbt.adapters.athena.exceptions import S3LocationException
 from dbt.adapters.athena.relation import AthenaRelation, TableType
+from dbt.adapters.athena.utils import AthenaCatalogType
 from dbt.clients import agate_helper
 from dbt.contracts.connection import ConnectionState
 from dbt.contracts.files import FileHash
@@ -28,6 +32,7 @@ from .constants import (
     BUCKET,
     DATA_CATALOG_NAME,
     DATABASE_NAME,
+    FEDERATED_QUERY_CATALOG_NAME,
     S3_STAGING_DIR,
     SHARED_DATA_CATALOG_NAME,
 )
@@ -66,6 +71,7 @@ class TestAthenaAdapter:
             ("awsdatacatalog", "quux"),
             ("awsdatacatalog", "baz"),
             (SHARED_DATA_CATALOG_NAME, "foo"),
+            (FEDERATED_QUERY_CATALOG_NAME, "foo"),
         }
         self.mock_manifest.nodes = {
             "model.root.model1": CompiledNode(
@@ -212,6 +218,42 @@ class TestAthenaAdapter:
                 raw_code="select * from source_table",
                 language="",
             ),
+            "model.root.model5": CompiledNode(
+                name="model5",
+                database=FEDERATED_QUERY_CATALOG_NAME,
+                schema="foo",
+                resource_type=NodeType.Model,
+                unique_id="model.root.model5",
+                alias="bar",
+                fqn=["root", "model5"],
+                package_name="root",
+                refs=[],
+                sources=[],
+                depends_on=DependsOn(),
+                config=NodeConfig.from_dict(
+                    {
+                        "enabled": True,
+                        "materialized": "table",
+                        "persist_docs": {},
+                        "post-hook": [],
+                        "pre-hook": [],
+                        "vars": {},
+                        "meta": {"owner": "data-engineers"},
+                        "quoting": {},
+                        "column_types": {},
+                        "tags": [],
+                    }
+                ),
+                tags=[],
+                path="model5.sql",
+                original_file_path="model5.sql",
+                compiled=True,
+                extra_ctes_injected=False,
+                extra_ctes=[],
+                checksum=FileHash.from_contents(""),
+                raw_code="select * from source_table",
+                language="",
+            ),
         }
 
     @property
@@ -612,9 +654,84 @@ class TestAthenaAdapter:
         for row in actual.rows.values():
             assert row.values() in expected_rows
 
+    @mock_athena
+    def test__get_one_catalog_federated_query_catalog(self, mock_aws_service):
+        mock_aws_service.create_data_catalog(
+            catalog_name=FEDERATED_QUERY_CATALOG_NAME, catalog_type=AthenaCatalogType.LAMBDA
+        )
+        mock_information_schema = mock.MagicMock()
+        mock_information_schema.path.database = FEDERATED_QUERY_CATALOG_NAME
+
+        # Original botocore _make_api_call function
+        orig = botocore.client.BaseClient._make_api_call
+
+        # Mocking this as list_table_metadata and creating non glue tables is not supported by moto.
+        # Followed this guide: http://docs.getmoto.org/en/latest/docs/services/patching_other_services.html
+        def mock_athena_list_table_metadata(self, operation_name, kwarg):
+            if operation_name == "ListTableMetadata":
+                return {
+                    "TableMetadataList": [
+                        {
+                            "Name": "bar",
+                            "TableType": "EXTERNAL_TABLE",
+                            "Columns": [
+                                {
+                                    "Name": "id",
+                                    "Type": "string",
+                                },
+                                {
+                                    "Name": "country",
+                                    "Type": "string",
+                                },
+                            ],
+                            "PartitionKeys": [
+                                {
+                                    "Name": "dt",
+                                    "Type": "date",
+                                },
+                            ],
+                        }
+                    ],
+                }
+            # If we don't want to patch the API call
+            return orig(self, operation_name, kwarg)
+
+        self.adapter.acquire_connection("dummy")
+        with patch("botocore.client.BaseClient._make_api_call", new=mock_athena_list_table_metadata):
+            actual = self.adapter._get_one_catalog(
+                mock_information_schema,
+                {
+                    "foo": {"bar"},
+                },
+                self.mock_manifest,
+            )
+
+        expected_column_names = (
+            "table_database",
+            "table_schema",
+            "table_name",
+            "table_type",
+            "table_comment",
+            "column_name",
+            "column_index",
+            "column_type",
+            "column_comment",
+            "table_owner",
+        )
+        expected_rows = [
+            (FEDERATED_QUERY_CATALOG_NAME, "foo", "bar", "table", None, "id", 0, "string", None, "data-engineers"),
+            (FEDERATED_QUERY_CATALOG_NAME, "foo", "bar", "table", None, "country", 1, "string", None, "data-engineers"),
+            (FEDERATED_QUERY_CATALOG_NAME, "foo", "bar", "table", None, "dt", 2, "date", None, "data-engineers"),
+        ]
+
+        assert actual.column_names == expected_column_names
+        assert len(actual.rows) == len(expected_rows)
+        for row in actual.rows.values():
+            assert row.values() in expected_rows
+
     def test__get_catalog_schemas(self):
         res = self.adapter._get_catalog_schemas(self.mock_manifest)
-        assert len(res.keys()) == 2
+        assert len(res.keys()) == 3
 
         information_schema_0 = list(res.keys())[0]
         assert information_schema_0.name == "INFORMATION_SCHEMA"
@@ -632,6 +749,14 @@ class TestAthenaAdapter:
         assert set(relations.keys()) == {"foo"}
         assert list(relations.values()) == [{"bar"}]
 
+        information_schema_1 = list(res.keys())[2]
+        assert information_schema_1.name == "INFORMATION_SCHEMA"
+        assert information_schema_1.schema is None
+        assert information_schema_1.database == FEDERATED_QUERY_CATALOG_NAME
+        relations = list(res.values())[1]
+        assert set(relations.keys()) == {"foo"}
+        assert list(relations.values()) == [{"bar"}]
+
     @mock_athena
     @mock_sts
     def test__get_data_catalog(self, mock_aws_service):
@@ -696,7 +821,7 @@ class TestAthenaAdapter:
         self, parent_list_relations_without_caching, mock_aws_service
     ):
         data_catalog_name = "other_data_catalog"
-        mock_aws_service.create_data_catalog(data_catalog_name, "HIVE")
+        mock_aws_service.create_data_catalog(data_catalog_name, AthenaCatalogType.HIVE)
         schema_relation = self.adapter.Relation.create(
             database=data_catalog_name,
             schema=DATABASE_NAME,
diff --git a/tests/unit/utils.py b/tests/unit/utils.py
index 60ec02e..097f6ce 100644
--- a/tests/unit/utils.py
+++ b/tests/unit/utils.py
@@ -5,6 +5,7 @@ from typing import Optional
 import agate
 import boto3
 
+from dbt.adapters.athena.utils import AthenaCatalogType
 from dbt.config.project import PartialProject
 
 from .constants import AWS_REGION, BUCKET, CATALOG_ID, DATA_CATALOG_NAME, DATABASE_NAME
@@ -146,10 +147,18 @@ class TestAdapterConversions:
 
 class MockAWSService:
     def create_data_catalog(
-        self, catalog_name: str = DATA_CATALOG_NAME, catalog_type: str = "GLUE", catalog_id: str = CATALOG_ID
+        self,
+        catalog_name: str = DATA_CATALOG_NAME,
+        catalog_type: AthenaCatalogType = AthenaCatalogType.GLUE,
+        catalog_id: str = CATALOG_ID,
     ):
         athena = boto3.client("athena", region_name=AWS_REGION)
-        athena.create_data_catalog(Name=catalog_name, Type=catalog_type, Parameters={"catalog-id": catalog_id})
+        parameters = {}
+        if catalog_type == AthenaCatalogType.GLUE:
+            parameters = {"catalog-id": catalog_id}
+        else:
+            parameters = {"catalog": catalog_name}
+        athena.create_data_catalog(Name=catalog_name, Type=catalog_type.value, Parameters=parameters)
 
     def create_database(self, name: str = DATABASE_NAME, catalog_id: str = CATALOG_ID):
         glue = boto3.client("glue", region_name=AWS_REGION)