import functools
import logging
from collections import defaultdict
from typing import Dict, Iterable, List, Optional, Type, Union

import humanfriendly

# These imports verify that the dependencies are available.
import pydantic
import redshift_connector

from datahub.configuration.common import AllowDenyPattern
from datahub.configuration.pattern_utils import is_schema_allowed
from datahub.emitter.mce_builder import (
    make_data_platform_urn,
    make_dataset_urn_with_platform_instance,
    make_tag_urn,
)
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import (
    SourceCapability,
    SupportStatus,
    capability,
    config_class,
    platform_name,
    support_status,
)
from datahub.ingestion.api.incremental_lineage_helper import auto_incremental_lineage
from datahub.ingestion.api.source import (
    CapabilityReport,
    MetadataWorkUnitProcessor,
    TestableSource,
    TestConnectionReport,
)
from datahub.ingestion.api.source_helpers import (
    auto_workunit,
    create_dataset_props_patch_builder,
)
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.glossary.classification_mixin import (
    ClassificationHandler,
    classification_workunit_processor,
)
from datahub.ingestion.source.common.data_reader import DataReader
from datahub.ingestion.source.common.subtypes import (
    DatasetContainerSubTypes,
    DatasetSubTypes,
    SourceCapabilityModifier,
)
from datahub.ingestion.source.redshift.config import RedshiftConfig
from datahub.ingestion.source.redshift.datashares import RedshiftDatasharesHelper
from datahub.ingestion.source.redshift.exception import handle_redshift_exceptions_yield
from datahub.ingestion.source.redshift.lineage import RedshiftSqlLineage
from datahub.ingestion.source.redshift.profile import RedshiftProfiler
from datahub.ingestion.source.redshift.redshift_data_reader import RedshiftDataReader
from datahub.ingestion.source.redshift.redshift_schema import (
    RedshiftColumn,
    RedshiftDatabase,
    RedshiftDataDictionary,
    RedshiftSchema,
    RedshiftTable,
    RedshiftView,
)
from datahub.ingestion.source.redshift.report import RedshiftReport
from datahub.ingestion.source.redshift.usage import RedshiftUsageExtractor
from datahub.ingestion.source.sql.sql_common import SqlWorkUnit
from datahub.ingestion.source.sql.sql_types import resolve_postgres_modified_type
from datahub.ingestion.source.sql.sql_utils import (
    add_table_to_schema_container,
    gen_database_container,
    gen_database_key,
    gen_schema_container,
    gen_schema_key,
    get_dataplatform_instance_aspect,
    get_domain_wu,
)
from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler
from datahub.ingestion.source.state.redundant_run_skip_handler import (
    RedundantLineageRunSkipHandler,
    RedundantUsageRunSkipHandler,
)
from datahub.ingestion.source.state.stale_entity_removal_handler import (
    StaleEntityRemovalHandler,
)
from datahub.ingestion.source.state.stateful_ingestion_base import (
    StatefulIngestionSourceBase,
)
from datahub.ingestion.source_report.ingestion_stage import (
    LINEAGE_EXTRACTION,
    METADATA_EXTRACTION,
    PROFILING,
    USAGE_EXTRACTION_INGESTION,
)
from datahub.metadata.com.linkedin.pegasus2avro.common import SubTypes, TimeStamp
from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
    DatasetProperties,
    ViewProperties,
)
from datahub.metadata.com.linkedin.pegasus2avro.schema import (
    ArrayType,
    BooleanType,
    BytesType,
    MySqlDDL,
    NullType,
    NumberType,
    RecordType,
    SchemaField,
    SchemaFieldDataType,
    SchemaMetadata,
    StringType,
    TimeType,
)
from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
from datahub.utilities import memory_footprint
from datahub.utilities.mapping import Constants
from datahub.utilities.perf_timer import PerfTimer
from datahub.utilities.registries.domain_registry import DomainRegistry

logger: logging.Logger = logging.getLogger(__name__)


@platform_name("Redshift")
@config_class(RedshiftConfig)
@support_status(SupportStatus.CERTIFIED)
@capability(
    SourceCapability.CONTAINERS,
    "Enabled by default",
    subtype_modifier=[
        SourceCapabilityModifier.DATABASE,
        SourceCapabilityModifier.SCHEMA,
    ],
)
@capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
@capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
@capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
@capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
@capability(
    SourceCapability.LINEAGE_FINE,
    "Optionally enabled via configuration (`mixed` or `sql_based` lineage needs to be enabled)",
)
@capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
@capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
@capability(
    SourceCapability.USAGE_STATS,
    "Optionally enabled via `include_usage_statistics`",
)
@capability(
    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
)
@capability(
    SourceCapability.CLASSIFICATION,
    "Optionally enabled via `classification.enabled`",
    supported=True,
)
@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
    """
    This plugin extracts the following:

    - Metadata for databases, schemas, views and tables
    - Column types associated with each table
    - Table, row, and column statistics via optional SQL profiling
    - Table lineage
    - Usage statistics
    """

    # TODO: Replace with standardized types in sql_types.py
    REDSHIFT_FIELD_TYPE_MAPPINGS: Dict[
        str,
        Type[
            Union[
                ArrayType,
                BytesType,
                BooleanType,
                NumberType,
                RecordType,
                StringType,
                TimeType,
                NullType,
            ]
        ],
    ] = {
        "BYTES": BytesType,
        "BOOL": BooleanType,
        "BOOLEAN": BooleanType,
        "DOUBLE": NumberType,
        "DOUBLE PRECISION": NumberType,
        "DECIMAL": NumberType,
        "NUMERIC": NumberType,
        "BIGNUMERIC": NumberType,
        "BIGDECIMAL": NumberType,
        "FLOAT64": NumberType,
        "INT": NumberType,
        "INT64": NumberType,
        "SMALLINT": NumberType,
        "INTEGER": NumberType,
        "BIGINT": NumberType,
        "TINYINT": NumberType,
        "BYTEINT": NumberType,
        "STRING": StringType,
        "TIME": TimeType,
        "TIMESTAMP": TimeType,
        "DATE": TimeType,
        "DATETIME": TimeType,
        "GEOGRAPHY": NullType,
        "JSON": NullType,
        "INTERVAL": NullType,
        "ARRAY": ArrayType,
        "STRUCT": RecordType,
        "CHARACTER VARYING": StringType,
        "CHARACTER": StringType,
        "CHAR": StringType,
        "TIMESTAMP WITHOUT TIME ZONE": TimeType,
        "REAL": NumberType,
        "VARCHAR": StringType,
        "TIMESTAMPTZ": TimeType,
        "GEOMETRY": NullType,
        "HLLSKETCH": NullType,
        "TIMETZ": TimeType,
        "VARBYTE": StringType,
        "SUPER": NullType,
    }

    def get_platform_instance_id(self) -> str:
        """
        The source identifier such as the specific source host address required for stateful ingestion.
        Individual subclasses need to override this method appropriately.
        """
        return str(self.platform)

    @staticmethod
    def test_connection(config_dict: dict) -> TestConnectionReport:
        test_report = TestConnectionReport()
        try:
            RedshiftConfig.Config.extra = (
                pydantic.Extra.allow
            )  # we are okay with extra fields during this stage
            config = RedshiftConfig.parse_obj(config_dict)
            # source = RedshiftSource(config, report)
            connection: redshift_connector.Connection = (
                RedshiftSource.get_redshift_connection(config)
            )
            cur = connection.cursor()
            cur.execute("select 1")
            test_report.basic_connectivity = CapabilityReport(capable=True)

            test_report.capability_report = {}
            try:
                RedshiftDataDictionary.get_schemas(connection, database=config.database)
                test_report.capability_report[SourceCapability.SCHEMA_METADATA] = (
                    CapabilityReport(capable=True)
                )
            except Exception as e:
                test_report.capability_report[SourceCapability.SCHEMA_METADATA] = (
                    CapabilityReport(capable=False, failure_reason=str(e))
                )

        except Exception as e:
            test_report.basic_connectivity = CapabilityReport(
                capable=False, failure_reason=str(e)
            )
        return test_report

    def get_report(self) -> RedshiftReport:
        return self.report

    eskind_to_platform = {1: "glue", 2: "hive", 3: "postgres", 4: "redshift"}

    def __init__(self, config: RedshiftConfig, ctx: PipelineContext):
        super().__init__(config, ctx)
        self.catalog_metadata: Dict = {}
        self.config: RedshiftConfig = config
        self.report: RedshiftReport = RedshiftReport()
        self.classification_handler = ClassificationHandler(self.config, self.report)
        self.datashares_helper = RedshiftDatasharesHelper(
            self.config, self.report, self.ctx.graph
        )
        self.platform = "redshift"
        self.domain_registry = None
        if self.config.domain:
            self.domain_registry = DomainRegistry(
                cached_domains=list(self.config.domain.keys()), graph=self.ctx.graph
            )

        self.redundant_lineage_run_skip_handler: Optional[
            RedundantLineageRunSkipHandler
        ] = None
        if self.config.enable_stateful_lineage_ingestion:
            self.redundant_lineage_run_skip_handler = RedundantLineageRunSkipHandler(
                source=self,
                config=self.config,
                pipeline_name=self.ctx.pipeline_name,
                run_id=self.ctx.run_id,
            )

        self.profiling_state_handler: Optional[ProfilingHandler] = None
        if self.config.enable_stateful_profiling:
            self.profiling_state_handler = ProfilingHandler(
                source=self,
                config=self.config,
                pipeline_name=self.ctx.pipeline_name,
                run_id=self.ctx.run_id,
            )

        self.data_dictionary = RedshiftDataDictionary(
            is_serverless=self.config.is_serverless
        )

        self.db: Optional[RedshiftDatabase] = None
        self.db_tables: Dict[str, Dict[str, List[RedshiftTable]]] = {}
        self.db_views: Dict[str, Dict[str, List[RedshiftView]]] = {}
        self.db_schemas: Dict[str, Dict[str, RedshiftSchema]] = {}

        self.add_config_to_report()

    @classmethod
    def create(cls, config_dict, ctx):
        config = RedshiftConfig.parse_obj(config_dict)
        return cls(config, ctx)

    @staticmethod
    def get_redshift_connection(
        config: RedshiftConfig,
    ) -> redshift_connector.Connection:
        host, port = config.host_port.split(":")
        conn = redshift_connector.connect(
            host=host,
            port=int(port),
            user=config.username,
            database=config.database,
            password=config.password.get_secret_value() if config.password else None,
            **config.extra_client_options,
        )

        conn.autocommit = True

        return conn

    def gen_database_container(self, database: str) -> Iterable[MetadataWorkUnit]:
        database_container_key = gen_database_key(
            database=database,
            platform=self.platform,
            platform_instance=self.config.platform_instance,
            env=self.config.env,
        )

        yield from gen_database_container(
            database=database,
            database_container_key=database_container_key,
            sub_types=[DatasetContainerSubTypes.DATABASE],
        )

    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
        return [
            *super().get_workunit_processors(),
            functools.partial(
                auto_incremental_lineage, self.config.incremental_lineage
            ),
            StaleEntityRemovalHandler.create(
                self, self.config, self.ctx
            ).workunit_processor,
        ]

    def _warn_deprecated_configs(self):
        if (
            self.config.match_fully_qualified_names is not None
            and not self.config.match_fully_qualified_names
            and self.config.schema_pattern is not None
            and self.config.schema_pattern != AllowDenyPattern.allow_all()
        ):
            self.report.report_warning(
                message="Please update `schema_pattern` to match against fully qualified schema name `<database_name>.<schema_name>` and set config `match_fully_qualified_names : True`."
                "Current default `match_fully_qualified_names: False` is only to maintain backward compatibility. "
                "The config option `match_fully_qualified_names` will be removed in future and the default behavior will be like `match_fully_qualified_names: True`.",
                context="Config option deprecation warning",
                title="Config option deprecation warning",
            )

    def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
        self._warn_deprecated_configs()
        connection = self._try_get_redshift_connection(self.config)

        if connection is None:
            # If we failed to establish a connection, short circuit the connector.
            return

        database = self.config.database
        logger.info(f"Processing db {database}")

        self.db = self.data_dictionary.get_database_details(connection, database)
        self.report.is_shared_database = (
            self.db is not None and self.db.is_shared_database()
        )
        with self.report.new_stage(METADATA_EXTRACTION):
            self.db_tables[database] = defaultdict()
            self.db_views[database] = defaultdict()
            self.db_schemas.setdefault(database, {})

        # TODO: Ideally, we'd push down exception handling to the place where the connection is used, as opposed to keeping
        # this fallback. For now, this gets us broad coverage quickly.
        yield from handle_redshift_exceptions_yield(
            self.report, self._extract_metadata, connection, database
        )

    def _extract_metadata(
        self, connection: redshift_connector.Connection, database: str
    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
        yield from self.gen_database_container(
            database=database,
        )

        self.cache_tables_and_views(connection, database)

        self.report.tables_in_mem_size[database] = humanfriendly.format_size(
            memory_footprint.total_size(self.db_tables)
        )
        self.report.views_in_mem_size[database] = humanfriendly.format_size(
            memory_footprint.total_size(self.db_views)
        )

        with RedshiftSqlLineage(
            config=self.config,
            report=self.report,
            context=self.ctx,
            database=database,
            redundant_run_skip_handler=self.redundant_lineage_run_skip_handler,
        ) as lineage_extractor:
            yield from lineage_extractor.aggregator.register_schemas_from_stream(
                self.process_schemas(connection, database)
            )

            with self.report.new_stage(LINEAGE_EXTRACTION):
                yield from self.extract_lineage_v2(
                    connection=connection,
                    database=database,
                    lineage_extractor=lineage_extractor,
                )

        all_tables = self.get_all_tables()

        if self.config.include_usage_statistics:
            with self.report.new_stage(USAGE_EXTRACTION_INGESTION):
                yield from self.extract_usage(
                    connection=connection, all_tables=all_tables, database=database
                )

        if self.config.is_profiling_enabled():
            with self.report.new_stage(PROFILING):
                profiler = RedshiftProfiler(
                    config=self.config,
                    report=self.report,
                    state_handler=self.profiling_state_handler,
                )
                yield from profiler.get_workunits(self.db_tables)

    def process_schemas(self, connection, database):
        for schema in self.data_dictionary.get_schemas(
            conn=connection, database=database
        ):
            if not is_schema_allowed(
                self.config.schema_pattern,
                schema.name,
                database,
                self.config.match_fully_qualified_names,
            ):
                self.report.report_dropped(f"{database}.{schema.name}")
                continue

            logger.info(f"Processing schema: {database}.{schema.name}")

            self.db_schemas[database][schema.name] = schema
            yield from self.process_schema(connection, database, schema)

    def make_data_reader(
        self,
        connection: redshift_connector.Connection,
    ) -> Optional[DataReader]:
        if self.classification_handler.is_classification_enabled():
            return RedshiftDataReader.create(connection)

        return None

    def process_schema(
        self,
        connection: redshift_connector.Connection,
        database: str,
        schema: RedshiftSchema,
    ) -> Iterable[MetadataWorkUnit]:
        report_key = f"{database}.{schema.name}"
        with PerfTimer() as timer:
            schema_container_key = gen_schema_key(
                db_name=database,
                schema=schema.name,
                platform=self.platform,
                platform_instance=self.config.platform_instance,
                env=self.config.env,
            )

            database_container_key = gen_database_key(
                database=database,
                platform=self.platform,
                platform_instance=self.config.platform_instance,
                env=self.config.env,
            )

            yield from gen_schema_container(
                schema=schema.name,
                database=database,
                schema_container_key=schema_container_key,
                database_container_key=database_container_key,
                domain_config=self.config.domain,
                domain_registry=self.domain_registry,
                sub_types=[DatasetSubTypes.SCHEMA],
            )

            schema_columns: Dict[str, Dict[str, List[RedshiftColumn]]] = {}
            schema_columns[schema.name] = self.data_dictionary.get_columns_for_schema(
                conn=connection,
                database=database,
                schema=schema,
                is_shared_database=self.report.is_shared_database,
            )

            if self.config.include_tables:
                data_reader = self.make_data_reader(connection)
                logger.info(f"Process tables in schema {database}.{schema.name}")
                if (
                    self.db_tables[schema.database]
                    and schema.name in self.db_tables[schema.database]
                ):
                    for table in self.db_tables[schema.database][schema.name]:
                        table.columns = schema_columns[schema.name].get(table.name, [])
                        table.column_count = len(table.columns)
                        table_wu_generator = self._process_table(
                            table, database=database
                        )
                        yield from classification_workunit_processor(
                            table_wu_generator,
                            self.classification_handler,
                            data_reader,
                            [schema.database, schema.name, table.name],
                        )
                        self.report.table_processed[report_key] = (
                            self.report.table_processed.get(
                                f"{database}.{schema.name}", 0
                            )
                            + 1
                        )
                        logger.debug(
                            f"Table processed: {schema.database}.{schema.name}.{table.name}"
                        )
                else:
                    self.report.info(
                        title="No tables found in some schemas",
                        message="No tables found in some schemas. This may be due to insufficient privileges for the provided user.",
                        context=f"Schema: {schema.database}.{schema.name}",
                    )
            else:
                logger.info("Table processing disabled, skipping")

            if self.config.include_views:
                logger.info(f"Process views in schema {schema.database}.{schema.name}")
                if (
                    self.db_views[schema.database]
                    and schema.name in self.db_views[schema.database]
                ):
                    for view in self.db_views[schema.database][schema.name]:
                        view.columns = schema_columns[schema.name].get(view.name, [])
                        view.column_count = len(view.columns)
                        yield from self._process_view(
                            table=view, database=database, schema=schema
                        )

                        self.report.view_processed[report_key] = (
                            self.report.view_processed.get(
                                f"{database}.{schema.name}", 0
                            )
                            + 1
                        )
                        logger.debug(
                            f"Table processed: {schema.database}.{schema.name}.{view.name}"
                        )
                else:
                    self.report.info(
                        title="No views found in some schemas",
                        message="No views found in some schemas. This may be due to insufficient privileges for the provided user.",
                        context=f"Schema: {schema.database}.{schema.name}",
                    )
            else:
                logger.info("View processing disabled, skipping")

            self.report.metadata_extraction_sec[report_key] = timer.elapsed_seconds(
                digits=2
            )

    def _process_table(
        self,
        table: RedshiftTable,
        database: str,
    ) -> Iterable[MetadataWorkUnit]:
        datahub_dataset_name = f"{database}.{table.schema}.{table.name}"

        self.report.report_entity_scanned(datahub_dataset_name)

        if not self.config.table_pattern.allowed(datahub_dataset_name):
            self.report.report_dropped(datahub_dataset_name)
            return

        yield from self.gen_table_dataset_workunits(
            table, database=database, dataset_name=datahub_dataset_name
        )

    def _process_view(
        self, table: RedshiftView, database: str, schema: RedshiftSchema
    ) -> Iterable[MetadataWorkUnit]:
        datahub_dataset_name = f"{database}.{schema.name}.{table.name}"

        self.report.report_entity_scanned(datahub_dataset_name)

        if not self.config.table_pattern.allowed(datahub_dataset_name):
            self.report.report_dropped(datahub_dataset_name)
            return

        yield from self.gen_view_dataset_workunits(
            view=table,
            database=database,
            schema=schema.name,
        )

    def gen_table_dataset_workunits(
        self,
        table: RedshiftTable,
        database: str,
        dataset_name: str,
    ) -> Iterable[MetadataWorkUnit]:
        custom_properties = {}

        if table.type:
            custom_properties["table_type"] = table.type

        if table.location:
            custom_properties["location"] = table.location

        if table.input_parameters:
            custom_properties["input_parameters"] = table.input_parameters

        if table.output_parameters:
            custom_properties["output_parameters"] = table.output_parameters

        if table.dist_style:
            custom_properties["dist_style"] = table.dist_style

        if table.parameters:
            custom_properties["parameters"] = table.parameters

        if table.serde_parameters:
            custom_properties["serde_parameters"] = table.serde_parameters

        assert table.schema
        assert table.type
        yield from self.gen_dataset_workunits(
            table=table,
            database=database,
            schema=table.schema,
            sub_type=DatasetSubTypes.TABLE,
            custom_properties=custom_properties,
        )

    # TODO: Remove to common?
    def gen_view_dataset_workunits(
        self,
        view: RedshiftView,
        database: str,
        schema: str,
    ) -> Iterable[MetadataWorkUnit]:
        yield from self.gen_dataset_workunits(
            table=view,
            database=self.config.database,
            schema=schema,
            sub_type=DatasetSubTypes.VIEW,
            custom_properties={},
        )

        datahub_dataset_name = f"{database}.{schema}.{view.name}"
        dataset_urn = self.gen_dataset_urn(datahub_dataset_name)
        if view.ddl:
            view_properties_aspect = ViewProperties(
                materialized=view.materialized,
                viewLanguage="SQL",
                viewLogic=view.ddl,
            )
            yield MetadataChangeProposalWrapper(
                entityUrn=dataset_urn, aspect=view_properties_aspect
            ).as_workunit()

    # TODO: Remove to common?
    def gen_schema_fields(self, columns: List[RedshiftColumn]) -> List[SchemaField]:
        schema_fields: List[SchemaField] = []

        for col in columns:
            tags: List[TagAssociationClass] = []
            if col.dist_key:
                tags.append(TagAssociationClass(make_tag_urn(Constants.TAG_DIST_KEY)))

            if col.sort_key:
                tags.append(TagAssociationClass(make_tag_urn(Constants.TAG_SORT_KEY)))

            data_type = self.REDSHIFT_FIELD_TYPE_MAPPINGS.get(col.data_type)
            # We have to remove the precision part to properly parse it
            if data_type is None:
                # attempt Postgres modified type
                data_type = resolve_postgres_modified_type(col.data_type.lower())

            if any(type in col.data_type.lower() for type in ["struct", "array"]):
                fields = self.data_dictionary.get_schema_fields_for_column(col)
                schema_fields.extend(fields)
            else:
                field = SchemaField(
                    fieldPath=col.name,
                    type=SchemaFieldDataType(data_type() if data_type else NullType()),
                    # NOTE: nativeDataType will not be in sync with older connector
                    nativeDataType=col.data_type,
                    description=col.comment,
                    nullable=col.is_nullable,
                    globalTags=GlobalTagsClass(tags=tags),
                )
                schema_fields.append(field)
        return schema_fields

    # TODO: Move to common?
    def gen_schema_metadata(
        self,
        dataset_urn: str,
        table: Union[RedshiftTable, RedshiftView],
        dataset_name: str,
    ) -> Iterable[MetadataWorkUnit]:
        schema_metadata = SchemaMetadata(
            schemaName=dataset_name,
            platform=make_data_platform_urn(self.platform),
            version=0,
            hash="",
            platformSchema=MySqlDDL(tableSchema=""),
            fields=self.gen_schema_fields(table.columns),
        )
        yield MetadataChangeProposalWrapper(
            entityUrn=dataset_urn, aspect=schema_metadata
        ).as_workunit()

    def gen_dataset_urn(self, datahub_dataset_name: str) -> str:
        return make_dataset_urn_with_platform_instance(
            platform=self.platform,
            name=datahub_dataset_name,
            platform_instance=self.config.platform_instance,
            env=self.config.env,
        )

    # TODO: Move to common
    def gen_dataset_workunits(
        self,
        table: Union[RedshiftTable, RedshiftView],
        database: str,
        schema: str,
        sub_type: str,
        custom_properties: Optional[Dict[str, str]] = None,
    ) -> Iterable[MetadataWorkUnit]:
        datahub_dataset_name = f"{database}.{schema}.{table.name}"
        dataset_urn = self.gen_dataset_urn(datahub_dataset_name)

        yield from self.gen_schema_metadata(
            dataset_urn, table, str(datahub_dataset_name)
        )

        dataset_properties = DatasetProperties(
            name=table.name,
            created=(
                TimeStamp(time=int(table.created.timestamp() * 1000))
                if table.created
                else None
            ),
            lastModified=(
                TimeStamp(time=int(table.last_altered.timestamp() * 1000))
                if table.last_altered
                else None
            ),
            description=table.comment,
            qualifiedName=str(datahub_dataset_name),
            customProperties=custom_properties,
        )
        if self.config.patch_custom_properties:
            # TODO: use auto_incremental_properties workunit processor instead
            # Deprecate use of patch_custom_properties
            patch_builder = create_dataset_props_patch_builder(
                dataset_urn, dataset_properties
            )
            for patch_mcp in patch_builder.build():
                yield MetadataWorkUnit(
                    id=f"{dataset_urn}-{patch_mcp.aspectName}", mcp_raw=patch_mcp
                )
        else:
            yield MetadataChangeProposalWrapper(
                entityUrn=dataset_urn, aspect=dataset_properties
            ).as_workunit()

        # TODO: Check if needed
        # if tags_to_add:
        #    yield gen_tags_aspect_workunit(dataset_urn, tags_to_add)

        schema_container_key = gen_schema_key(
            db_name=database,
            schema=schema,
            platform=self.platform,
            platform_instance=self.config.platform_instance,
            env=self.config.env,
        )

        yield from add_table_to_schema_container(
            dataset_urn,
            parent_container_key=schema_container_key,
        )
        dpi_aspect = get_dataplatform_instance_aspect(
            dataset_urn=dataset_urn,
            platform=self.platform,
            platform_instance=self.config.platform_instance,
        )
        if dpi_aspect:
            yield dpi_aspect

        subTypes = SubTypes(typeNames=[sub_type])
        yield MetadataChangeProposalWrapper(
            entityUrn=dataset_urn, aspect=subTypes
        ).as_workunit()

        if self.domain_registry:
            yield from get_domain_wu(
                dataset_name=str(datahub_dataset_name),
                entity_urn=dataset_urn,
                domain_registry=self.domain_registry,
                domain_config=self.config.domain,
            )

    def cache_tables_and_views(
        self, connection: redshift_connector.Connection, database: str
    ) -> None:
        tables, views = self.data_dictionary.get_tables_and_views(
            conn=connection,
            database=database,
            skip_external_tables=self.config.skip_external_tables,
            is_shared_database=self.report.is_shared_database,
        )
        for schema in tables:
            if not is_schema_allowed(
                self.config.schema_pattern,
                schema,
                database,
                self.config.match_fully_qualified_names,
            ):
                logger.debug(
                    f"Not caching table for schema {database}.{schema} which is not allowed by schema_pattern"
                )
                continue

            self.db_tables[database][schema] = []
            for table in tables[schema]:
                if self.config.table_pattern.allowed(
                    f"{database}.{schema}.{table.name}"
                ):
                    self.db_tables[database][schema].append(table)
                    self.report.table_cached[f"{database}.{schema}"] = (
                        self.report.table_cached.get(f"{database}.{schema}", 0) + 1
                    )
                else:
                    logger.debug(
                        f"Table {database}.{schema}.{table.name} is filtered by table_pattern"
                    )
                    self.report.table_filtered[f"{database}.{schema}"] = (
                        self.report.table_filtered.get(f"{database}.{schema}", 0) + 1
                    )

        for schema in views:
            if not is_schema_allowed(
                self.config.schema_pattern,
                schema,
                database,
                self.config.match_fully_qualified_names,
            ):
                logger.debug(
                    f"Not caching views for schema {database}.{schema} which is not allowed by schema_pattern"
                )
                continue

            self.db_views[database][schema] = []
            for view in views[schema]:
                if self.config.view_pattern.allowed(f"{database}.{schema}.{view.name}"):
                    self.db_views[database][schema].append(view)
                    self.report.view_cached[f"{database}.{schema}"] = (
                        self.report.view_cached.get(f"{database}.{schema}", 0) + 1
                    )
                else:
                    logger.debug(
                        f"View {database}.{schema}.{view.name} is filtered by view_pattern"
                    )
                    self.report.view_filtered[f"{database}.{schema}"] = (
                        self.report.view_filtered.get(f"{database}.{schema}", 0) + 1
                    )

    def get_all_tables(
        self,
    ) -> Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]]:
        all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]] = (
            defaultdict(dict)
        )
        for db in set().union(self.db_tables, self.db_views):
            tables = self.db_tables.get(db, {})
            views = self.db_views.get(db, {})
            for schema in set().union(tables, views):
                all_tables[db][schema] = [
                    *tables.get(schema, []),
                    *views.get(schema, []),
                ]
        return all_tables

    def extract_usage(
        self,
        connection: redshift_connector.Connection,
        database: str,
        all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]],
    ) -> Iterable[MetadataWorkUnit]:
        with PerfTimer() as timer:
            redundant_usage_run_skip_handler: Optional[RedundantUsageRunSkipHandler] = (
                None
            )
            if self.config.enable_stateful_usage_ingestion:
                redundant_usage_run_skip_handler = RedundantUsageRunSkipHandler(
                    source=self,
                    config=self.config,
                    pipeline_name=self.ctx.pipeline_name,
                    run_id=self.ctx.run_id,
                )
            usage_extractor = RedshiftUsageExtractor(
                config=self.config,
                connection=connection,
                report=self.report,
                dataset_urn_builder=self.gen_dataset_urn,
                redundant_run_skip_handler=redundant_usage_run_skip_handler,
            )

            yield from usage_extractor.get_usage_workunits(all_tables=all_tables)

            self.report.usage_extraction_sec[database] = timer.elapsed_seconds(digits=2)

    def extract_lineage_v2(
        self,
        connection: redshift_connector.Connection,
        database: str,
        lineage_extractor: RedshiftSqlLineage,
    ) -> Iterable[MetadataWorkUnit]:
        if self.config.include_share_lineage:
            outbound_shares = self.data_dictionary.get_outbound_datashares(connection)
            yield from auto_workunit(
                self.datashares_helper.to_platform_resource(list(outbound_shares))
            )

            if self.db and self.db.is_shared_database():
                inbound_share = self.db.get_inbound_share()
                if inbound_share is None:
                    self.report.warning(
                        title="Upstream lineage of inbound datashare will be missing",
                        message="Database options do not contain sufficient information",
                        context=f"Database: {database}, Options {self.db.options}",
                    )
                else:
                    for known_lineage in self.datashares_helper.generate_lineage(
                        inbound_share, self.get_all_tables()[database]
                    ):
                        lineage_extractor.aggregator.add(known_lineage)

        # TODO: distinguish between definition level lineage and audit log based lineage.
        # Definition level lineage should never be skipped
        if not self._should_ingest_lineage():
            return

        with PerfTimer() as timer:
            all_tables = self.get_all_tables()

            lineage_extractor.build(
                connection=connection, all_tables=all_tables, db_schemas=self.db_schemas
            )

            yield from lineage_extractor.generate()

            self.report.lineage_extraction_sec[f"{database}"] = timer.elapsed_seconds(
                digits=2
            )

        if self.redundant_lineage_run_skip_handler:
            # Update the checkpoint state for this run.
            self.redundant_lineage_run_skip_handler.update_state(
                lineage_extractor.start_time, lineage_extractor.end_time
            )

    def _should_ingest_lineage(self) -> bool:
        if (
            self.redundant_lineage_run_skip_handler
            and self.redundant_lineage_run_skip_handler.should_skip_this_run(
                cur_start_time=self.config.start_time,
                cur_end_time=self.config.end_time,
            )
        ):
            # Skip this run
            self.report.report_warning(
                "lineage-extraction",
                "Skip this run as there was already a run for current ingestion window.",
            )
            return False

        return True

    def add_config_to_report(self):
        self.report.stateful_lineage_ingestion_enabled = (
            self.config.enable_stateful_lineage_ingestion
        )
        self.report.stateful_usage_ingestion_enabled = (
            self.config.enable_stateful_usage_ingestion
        )
        self.report.window_start_time, self.report.window_end_time = (
            self.config.start_time,
            self.config.end_time,
        )

    def _try_get_redshift_connection(
        self,
        config: RedshiftConfig,
    ) -> Optional[redshift_connector.Connection]:
        try:
            return RedshiftSource.get_redshift_connection(config)
        except redshift_connector.Error as e:
            error_message = str(e).lower()
            if "password authentication failed" in error_message:
                self.report.report_failure(
                    title="Invalid credentials",
                    message="Failed to connect to Redshift. Please verify your username, password, and database.",
                    exc=e,
                )
            elif "timeout" in error_message:
                self.report.report_failure(
                    title="Unable to connect",
                    message="Failed to connect to Redshift. Please verify your host name and port number.",
                    exc=e,
                )
            elif "communication error" in error_message:
                self.report.report_failure(
                    title="Unable to connect",
                    message="Failed to connect to Redshift. Please verify that the host name is valid and reachable.",
                    exc=e,
                )
            elif "database" in error_message and "does not exist" in error_message:
                self.report.report_failure(
                    title="Database does not exist",
                    message="Failed to connect to Redshift. Please verify that the provided database exists and the provided user has access to it.",
                    exc=e,
                )
            else:
                self.report.report_failure(
                    title="Unable to connect",
                    message="Failed to connect to Redshift. Please verify your connection details.",
                    exc=e,
                )
            return None
