from prefect import task, flow
from prefect import get_run_logger

from prefect import flow
from prefect_dbt.cli.commands import trigger_dbt_cli_command

from prefect_snowflake.credentials import SnowflakeCredentials
from prefect_snowflake.database import SnowflakeConnector

from prefect_dbt.cli.commands import trigger_dbt_cli_command
from prefect_dbt.cli.configs import SnowflakeTargetConfigs
from prefect_dbt.cli import DbtCliProfile
from prefect.blocks.system import Secret

@flow
def build_dbt_docs():
    logger = get_run_logger()
    <http://logger.info|logger.info>("Calling dbt deps")
    
    # Access the stored secret
    secret_block = Secret.load("XXX")

    connector = SnowflakeConnector(
        schema="XXX",
        database="XXX",
        warehouse="COMPUTE_WH",
        credentials=SnowflakeCredentials(
            user=""XXX"",
            password=secret_block.get(),
            account="XXX",
            role="XXX",
        ),
    )
    target_configs = SnowflakeTargetConfigs(
        connector=connector
    )
    dbt_cli_profile = DbtCliProfile(
        name="XXX",
        target="XXX",
        target_configs=target_configs,
    )
    trigger_dbt_cli_command(
        "dbt deps",
        overwrite_profiles=True,
        dbt_cli_profile=dbt_cli_profile
    )
    <http://logger.info|logger.info>("Called dbt deps, calling docs generate.")
    trigger_dbt_cli_command(
        "dbt docs generate",
        overwrite_profiles=True,
        dbt_cli_profile=dbt_cli_profile
    )
    <http://logger.info|logger.info>("Done calling dbt run.")
    
if __name__ == "__main__":
    build_dbt_docs()