openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/ingestion/source/dashboard/tableau/models.py

@@ -15,7 +15,7 @@ Tableau Source Model module
 
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Extra, Field, validator
+from pydantic import BaseModel, ConfigDict, Field, validator
 
 from metadata.generated.schema.entity.data.chart import ChartType
 
@@ -25,11 +25,10 @@ class TableauBaseModel(BaseModel):
     Tableau basic configurations
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     id: str
-    name: Optional[str]
+    name: Optional[str] = None
 
     def __hash__(self):
         return hash(self.id)
@@ -54,8 +53,7 @@ class TableauTag(BaseModel):
     Aux class for Tag object of the tableau_api_lib response
     """
 
-    class Config:
-        frozen = True
+    model_config = ConfigDict(frozen=True)
 
     label: str
 
@@ -65,7 +63,7 @@ class TableauOwner(TableauBaseModel):
     Aux class for Owner object of the tableau_api_lib response
     """
 
-    email: Optional[str]
+    email: Optional[str] = None
 
 
 def transform_tags(raw: Union[Dict[str, Any], List[TableauTag]]) -> List[TableauTag]:
@@ -83,57 +81,58 @@ class CustomSQLTable(TableauBaseModel):
     https://help.tableau.com/current/api/metadata_api/en-us/reference/customsqltable.doc.html
     """
 
-    query: Optional[str]
+    query: Optional[str] = None
 
 
 class UpstreamColumn(BaseModel):
     id: str
-    name: Optional[str]
-    remoteType: Optional[str]
+    name: Optional[str] = None
+    remoteType: Optional[str] = None
 
 
 class DatasourceField(BaseModel):
     id: str
-    name: Optional[str]
-    upstreamColumns: Optional[List[Union[UpstreamColumn, None]]]
-    description: Optional[str]
+    name: Optional[str] = None
+    upstreamColumns: Optional[List[Union[UpstreamColumn, None]]] = None
+    description: Optional[str] = None
 
 
 class UpstreamTableColumn(BaseModel):
     id: str
-    name: Optional[str]
+    name: Optional[str] = None
 
 
 class TableauDatabase(BaseModel):
     id: str
-    name: Optional[str]
+    name: Optional[str] = None
 
 
 class UpstreamTable(BaseModel):
     id: str
     luid: str
-    name: Optional[str]
-    fullName: Optional[str]
-    schema_: Optional[str] = Field(..., alias="schema")
-    columns: Optional[List[UpstreamTableColumn]]
-    database: Optional[TableauDatabase]
-    referencedByQueries: Optional[List[CustomSQLTable]]
+    name: Optional[str] = None
+    fullName: Optional[str] = None
+    schema_: Optional[str] = Field(None, alias="schema")
+    columns: Optional[List[UpstreamTableColumn]] = None
+    database: Optional[TableauDatabase] = None
+    referencedByQueries: Optional[List[CustomSQLTable]] = None
 
 
 class DataSource(BaseModel):
     id: str
-    name: Optional[str]
-    fields: Optional[List[DatasourceField]]
-    upstreamTables: Optional[List[UpstreamTable]]
+    name: Optional[str] = None
+    fields: Optional[List[DatasourceField]] = None
+    upstreamTables: Optional[List[UpstreamTable]] = None
+    upstreamDatasources: Optional[List["DataSource"]] = None
 
 
 class TableauDatasources(BaseModel):
-    nodes: Optional[List[DataSource]]
-    totalCount: Optional[int]
+    nodes: Optional[List[DataSource]] = None
+    totalCount: Optional[int] = None
 
 
 class TableauDatasourcesConnection(BaseModel):
-    embeddedDatasourcesConnection: Optional[TableauDatasources]
+    embeddedDatasourcesConnection: Optional[TableauDatasources] = None
 
 
 class TableauChart(TableauBaseModel):
@@ -141,7 +140,7 @@ class TableauChart(TableauBaseModel):
     Aux class for Chart object of the tableau_api_lib response
     """
 
-    owner: Optional[TableauOwner]
+    owner: Optional[TableauOwner] = None
     tags: Optional[List[TableauTag]] = []
     _extract_tags = validator("tags", pre=True, allow_reuse=True)(transform_tags)
     contentUrl: Optional[str] = ""
@@ -153,14 +152,13 @@ class TableauDashboard(TableauBaseModel):
     Aux class for Dashboard object of the tableau_api_lib response
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
-    project: Optional[TableauBaseModel]
-    description: Optional[str]
-    owner: Optional[TableauOwner]
+    project: Optional[TableauBaseModel] = None
+    description: Optional[str] = None
+    owner: Optional[TableauOwner] = None
     tags: Optional[List[TableauTag]] = []
     _extract_tags = validator("tags", pre=True, allow_reuse=True)(transform_tags)
-    webpageUrl: Optional[str]
-    charts: Optional[List[TableauChart]]
+    webpageUrl: Optional[str] = None
+    charts: Optional[List[TableauChart]] = None
     dataModels: List[DataSource] = []

metadata/ingestion/source/dashboard/tableau/queries.py

@@ -23,6 +23,20 @@ workbooks(filter:{{luid: "{workbook_id}"}}){{
     nodes {{
       id
       name
+      upstreamDatasources{{
+        id
+        name
+        fields {{
+          id
+          name
+          upstreamColumns{{
+            id
+            name
+            remoteType
+          }}
+          description
+        }}
+      }}
       fields {{
         id
         name

metadata/ingestion/source/database/athena/connection.py

@@ -12,10 +12,13 @@
 """
 Source connection handler
 """
+from functools import partial
 from typing import Optional
 from urllib.parse import quote_plus
 
 from sqlalchemy.engine import Engine
+from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy.inspection import inspect
 
 from metadata.clients.aws_client import AWSClient
 from metadata.generated.schema.entity.automations.workflow import (
@@ -28,7 +31,11 @@ from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
 )
-from metadata.ingestion.connections.test_connections import test_connection_db_common
+from metadata.ingestion.connections.test_connections import (
+    execute_inspector_func,
+    test_connection_engine_step,
+    test_connection_steps,
+)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 
 
@@ -55,7 +62,7 @@ def get_connection_url(connection: AthenaConnection) -> str:
         url += ":"
     url += f"@athena.{connection.awsConfig.awsRegion}.amazonaws.com:443"
 
-    url += f"?s3_staging_dir={quote_plus(connection.s3StagingDir)}"
+    url += f"?s3_staging_dir={quote_plus(str(connection.s3StagingDir))}"
     if connection.workgroup:
         url += f"&work_group={connection.workgroup}"
     if aws_session_token:
@@ -92,9 +99,31 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    test_connection_db_common(
+
+    def get_test_schema(inspector: Inspector):
+        all_schemas = inspector.get_schema_names()
+        return all_schemas[0] if all_schemas else None
+
+    def custom_executor_for_table():
+        inspector = inspect(engine)
+        test_schema = get_test_schema(inspector)
+        return inspector.get_table_names(test_schema) if test_schema else []
+
+    def custom_executor_for_view():
+        inspector = inspect(engine)
+        test_schema = get_test_schema(inspector)
+        return inspector.get_view_names(test_schema) if test_schema else []
+
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, engine),
+        "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
+        "GetTables": custom_executor_for_table,
+        "GetViews": custom_executor_for_view,
+    }
+
+    test_connection_steps(
         metadata=metadata,
-        engine=engine,
-        service_connection=service_connection,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
     )

metadata/ingestion/source/database/athena/metadata.py

@@ -12,12 +12,9 @@
 """Athena source module"""
 
 import traceback
-from copy import deepcopy
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Iterable, Optional, Tuple
 
 from pyathena.sqlalchemy.base import AthenaDialect
-from sqlalchemy import types
-from sqlalchemy.engine import reflection
 from sqlalchemy.engine.reflection import Inspector
 
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
@@ -42,18 +39,33 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source import sqa_types
 from metadata.ingestion.source.database.athena.client import AthenaLakeFormationClient
-from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
+from metadata.ingestion.source.database.athena.utils import (
+    _get_column_type,
+    get_columns,
+    get_table_options,
+    get_view_definition,
+)
 from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.external_table_lineage_mixin import (
+    ExternalTableLineageMixin,
+)
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.sqlalchemy_utils import is_complex_type
+from metadata.utils.sqlalchemy_utils import get_all_table_ddls, get_table_ddl
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
+AthenaDialect._get_column_type = _get_column_type  # pylint: disable=protected-access
+AthenaDialect.get_columns = get_columns
+AthenaDialect.get_view_definition = get_view_definition
+AthenaDialect.get_table_options = get_table_options
+
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+
 logger = ingestion_logger()
 
 ATHENA_TAG = "ATHENA TAG"
@@ -72,165 +84,7 @@ ATHENA_INTERVAL_TYPE_MAP = {
 }
 
 
-def _get_column_type(self, type_):
-    """
-    Function overwritten from AthenaDialect
-    to add custom SQA typing.
-    """
-    type_ = type_.replace(" ", "").lower()
-    match = self._pattern_column_type.match(type_)  # pylint: disable=protected-access
-    if match:
-        name = match.group(1).lower()
-        length = match.group(2)
-    else:
-        name = type_.lower()
-        length = None
-
-    args = []
-    col_map = {
-        "boolean": types.BOOLEAN,
-        "float": types.FLOAT,
-        "double": types.FLOAT,
-        "real": types.FLOAT,
-        "tinyint": types.INTEGER,
-        "smallint": types.INTEGER,
-        "integer": types.INTEGER,
-        "int": types.INTEGER,
-        "bigint": types.BIGINT,
-        "string": types.String,
-        "date": types.DATE,
-        "timestamp": types.TIMESTAMP,
-        "binary": types.BINARY,
-        "varbinary": types.BINARY,
-        "array": types.ARRAY,
-        "json": types.JSON,
-        "struct": sqa_types.SQAStruct,
-        "row": sqa_types.SQAStruct,
-        "map": sqa_types.SQAMap,
-        "decimal": types.DECIMAL,
-        "varchar": types.VARCHAR,
-        "char": types.CHAR,
-    }
-    if name in ["decimal", "char", "varchar"]:
-        col_type = col_map[name]
-        if length:
-            args = [int(l) for l in length.split(",")]
-    elif type_.startswith("array"):
-        parsed_type = (
-            ColumnTypeParser._parse_datatype_string(  # pylint: disable=protected-access
-                type_
-            )
-        )
-        col_type = col_map["array"]
-        if parsed_type["arrayDataType"].lower().startswith("array"):
-            # as OpenMetadata doesn't store any details on children of array, we put
-            # in type as string as default to avoid Array item_type required issue
-            # from sqlalchemy types
-            args = [types.String]
-        else:
-            args = [col_map.get(parsed_type.get("arrayDataType").lower(), types.String)]
-    elif col_map.get(name):
-        col_type = col_map.get(name)
-    else:
-        logger.warning(f"Did not recognize type '{type_}'")
-        col_type = types.NullType
-    return col_type(*args)
-
-
-def _get_projection_details(
-    columns: List[Dict], projection_parameters: Dict
-) -> List[Dict]:
-    """Get the projection details for the columns
-
-    Args:
-        columns (List[Dict]): list of columns
-        projection_parameters (Dict): projection parameters
-    """
-    if not projection_parameters:
-        return columns
-
-    columns = deepcopy(columns)
-    for col in columns:
-        projection_details = next(
-            ({k: v} for k, v in projection_parameters.items() if k == col["name"]), None
-        )
-        if projection_details:
-            col["projection_type"] = projection_details[col["name"]]
-
-    return columns
-
-
-@reflection.cache
-def get_columns(self, connection, table_name, schema=None, **kw):
-    """
-    Method to handle table columns
-    """
-    metadata = self._get_table(  # pylint: disable=protected-access
-        connection, table_name, schema=schema, **kw
-    )
-    columns = [
-        {
-            "name": c.name,
-            "type": self._get_column_type(c.type),  # pylint: disable=protected-access
-            "nullable": True,
-            "default": None,
-            "autoincrement": False,
-            "comment": c.comment,
-            "system_data_type": c.type,
-            "is_complex": is_complex_type(c.type),
-            "dialect_options": {"awsathena_partition": True},
-        }
-        for c in metadata.partition_keys
-    ]
-
-    if kw.get("only_partition_columns"):
-        # Return projected partition information to set partition type in `get_table_partition_details`
-        # projected partition fields are stored in the form of `projection.<field_name>.type` as a table parameter
-        projection_parameters = {
-            key_.split(".")[1]: value_
-            for key_, value_ in metadata.parameters.items()
-            if key_.startswith("projection") and key_.endswith("type")
-        }
-        columns = _get_projection_details(columns, projection_parameters)
-        return columns
-
-    columns += [
-        {
-            "name": c.name,
-            "type": self._get_column_type(c.type),  # pylint: disable=protected-access
-            "nullable": True,
-            "default": None,
-            "autoincrement": False,
-            "comment": c.comment,
-            "system_data_type": c.type,
-            "is_complex": is_complex_type(c.type),
-            "dialect_options": {"awsathena_partition": None},
-        }
-        for c in metadata.columns
-    ]
-
-    return columns
-
-
-# pylint: disable=unused-argument
-@reflection.cache
-def get_view_definition(self, connection, view_name, schema=None, **kw):
-    """
-    Gets the view definition
-    """
-    full_view_name = f'"{view_name}"' if not schema else f'"{schema}"."{view_name}"'
-    res = connection.execute(f"SHOW CREATE VIEW {full_view_name}").fetchall()
-    if res:
-        return "\n".join(i[0] for i in res)
-    return None
-
-
-AthenaDialect._get_column_type = _get_column_type  # pylint: disable=protected-access
-AthenaDialect.get_columns = get_columns
-AthenaDialect.get_view_definition = get_view_definition
-
-
-class AthenaSource(CommonDbSourceService):
+class AthenaSource(ExternalTableLineageMixin, CommonDbSourceService):
     """
     Implements the necessary methods to extract
     Database metadata from Athena Source
@@ -240,8 +94,8 @@ class AthenaSource(CommonDbSourceService):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: AthenaConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: AthenaConnection = config.serviceConnection.root.config
         if not isinstance(connection, AthenaConnection):
             raise InvalidSourceException(
                 f"Expected AthenaConnection, but got {connection}"
@@ -257,6 +111,7 @@ class AthenaSource(CommonDbSourceService):
         self.athena_lake_formation_client = AthenaLakeFormationClient(
             connection=self.service_connection
         )
+        self.external_location_map = {}
 
     def query_table_names_and_types(
         self, schema_name: str
@@ -395,3 +250,23 @@ class AthenaSource(CommonDbSourceService):
                         stackTrace=traceback.format_exc(),
                     )
                 )
+
+    def get_table_description(
+        self, schema_name: str, table_name: str, inspector: Inspector
+    ) -> str:
+        description = None
+        try:
+            table_info: dict = inspector.get_table_comment(table_name, schema_name)
+            table_option = inspector.get_table_options(table_name, schema_name)
+            self.external_location_map[
+                (self.context.get().database, schema_name, table_name)
+            ] = table_option.get("awsathena_location")
+        # Catch any exception without breaking the ingestion
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Table description error for table [{schema_name}.{table_name}]: {exc}"
+            )
+        else:
+            description = table_info.get("text")
+        return description

metadata/ingestion/source/database/athena/models.py

@@ -19,33 +19,33 @@ from pydantic import BaseModel
 
 
 class QueryExecutionIdsResponse(BaseModel):
-    QueryExecutionIds: Optional[List[str]]
+    QueryExecutionIds: Optional[List[str]] = None
 
 
 class Status(BaseModel):
     State: Optional[str] = "FAILED"  # Default value
-    SubmissionDateTime: Optional[datetime]
+    SubmissionDateTime: Optional[datetime] = None
 
 
 class Statistics(BaseModel):
-    TotalExecutionTimeInMillis: Optional[int]
+    TotalExecutionTimeInMillis: Optional[int] = None
 
 
 class AthenaQueryExecution(BaseModel):
-    Query: Optional[str]
-    Statistics: Optional[Statistics]
-    Status: Optional[Status]
+    Query: Optional[str] = None
+    Statistics: Optional[Statistics] = None
+    Status: Optional[Status] = None
 
 
 class AthenaQueryExecutionList(BaseModel):
-    QueryExecutions: Optional[List[AthenaQueryExecution]]
+    QueryExecutions: Optional[List[AthenaQueryExecution]] = None
 
 
 class WorkGroup(BaseModel):
-    Name: Optional[str]
-    State: Optional[str]
+    Name: Optional[str] = None
+    State: Optional[str] = None
 
 
 class WorkGroupsList(BaseModel):
     WorkGroups: Optional[List[WorkGroup]] = []
-    NextToken: Optional[str]
+    NextToken: Optional[str] = None

metadata/ingestion/source/database/athena/query_parser.py

@@ -60,8 +60,8 @@ class AthenaQueryParserSource(QueryParserSource, ABC):
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
         """Create class instance"""
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: AthenaConnection = config.serviceConnection.__root__.config
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: AthenaConnection = config.serviceConnection.root.config
         if not isinstance(connection, AthenaConnection):
             raise InvalidSourceException(
                 f"Expected AthenaConnection, but got {connection}"