# Buggy code
def execute_archive(profile, node, context):
    adapter = get_adapter(profile)

    node_cfg = node.get("config", {})

    source_columns = adapter.get_columns_in_table(
        profile, node_cfg.get("source_schema"), node_cfg.get("source_table")
    )

    if len(source_columns) == 0:
        raise RuntimeError(
            'Source table "{}"."{}" does not exist'.format(source_schema, source_table)
        )

    dest_columns = source_columns + [
        dbt.schema.Column("valid_from", "timestamp", None),
        dbt.schema.Column("valid_to", "timestamp", None),
        dbt.schema.Column("scd_id", "text", None),
        dbt.schema.Column("dbt_updated_at", "timestamp", None),
    ]

    adapter.create_table(
        profile,
        schema=node_cfg.get("target_schema"),
        table=node_cfg.get("target_table"),
        columns=dest_columns,
        sort=node_cfg.get("updated_at"),
        dist=node_cfg.get("unique_key"),
    )

    # TODO move this to inject_runtime_config, generate archive SQL
    # in wrap step. can't do this right now because we actually need
    # to inspect status of the schema at runtime and archive requires
    # a lot of information about the schema to generate queries.
    template_ctx = context.copy()
    template_ctx.update(node_cfg)

    select = dbt.clients.jinja.get_rendered(
        dbt.templates.SCDArchiveTemplate, template_ctx
    )

    insert_stmt = dbt.templates.ArchiveInsertTemplate().wrap(
        schema=node_cfg.get("target_schema"),
        table=node_cfg.get("target_table"),
        query=select,
        unique_key=node_cfg.get("unique_key"),
    )

    node["wrapped_sql"] = dbt.clients.jinja.get_rendered(insert_stmt, template_ctx)

    result = adapter.execute_model(profile=profile, model=node)

    return result


# Traceback (psycopg2.NotSupportedError)
(venv)Tristans-MacBook-Pro:analytics tristan$ dbt archive
Traceback (most recent call last):
File "/Users/tristan/dev/venv/bin/dbt", line 6, in <module>
exec(compile(open(__file__).read(), __file__, 'exec'))
File "/Users/tristan/dev/dbt/scripts/dbt", line 8, in <module>
dbt.main.main(sys.argv[1:])
File "/Users/tristan/dev/dbt/dbt/main.py", line 38, in main
return handle(args)
File "/Users/tristan/dev/dbt/dbt/main.py", line 52, in handle
res = run_from_args(parsed)
File "/Users/tristan/dev/dbt/dbt/main.py", line 85, in run_from_args
return task.run()
File "/Users/tristan/dev/dbt/dbt/task/archive.py", line 19, in run
self.compile()
File "/Users/tristan/dev/dbt/dbt/task/archive.py", line 15, in compile
compiled = compiler.compile_archives()
File "/Users/tristan/dev/dbt/dbt/compilation.py", line 412, in compile_archives
sql = archive.compile()
File "/Users/tristan/dev/dbt/dbt/model.py", line 705, in compile
query = archival.compile()
File "/Users/tristan/dev/dbt/dbt/archival.py", line 41, in compile
self.schema.create_table(target_schema, target_table, dest_columns, sort=updated_at, dist=unique_key)
File "/Users/tristan/dev/dbt/dbt/schema.py", line 257, in create_table
self.execute_and_handle_permissions(sql, table)
File "/Users/tristan/dev/dbt/dbt/schema.py", line 159, in execute_and_handle_permissions
return self.execute(query)
File "/Users/tristan/dev/dbt/dbt/schema.py", line 137, in execute
raise e
File "/Users/tristan/dev/dbt/dbt/schema.py", line 129, in execute
cursor.execute(sql)
psycopg2.NotSupportedError: column "_sdc_source_key__id || '|' || _sdc_level_0_id" specified as distkey/sortkey is not in the table "mongo.accounts__duplicate_authors"

# Fixed code
def execute_archive(profile, node, context):
    adapter = get_adapter(profile)

    node_cfg = node.get("config", {})

    source_columns = adapter.get_columns_in_table(
        profile, node_cfg.get("source_schema"), node_cfg.get("source_table")
    )

    if len(source_columns) == 0:
        source_schema = node_cfg.get("source_schema")
        source_table = node_cfg.get("source_table")
        raise RuntimeError(
            'Source table "{}"."{}" does not exist'.format(source_schema, source_table)
        )

    dest_columns = source_columns + [
        dbt.schema.Column("valid_from", "timestamp", None),
        dbt.schema.Column("valid_to", "timestamp", None),
        dbt.schema.Column("scd_id", "text", None),
        dbt.schema.Column("dbt_updated_at", "timestamp", None),
    ]

    adapter.create_table(
        profile,
        schema=node_cfg.get("target_schema"),
        table=node_cfg.get("target_table"),
        columns=dest_columns,
        sort="dbt_updated_at",
        dist="scd_id",
    )

    # TODO move this to inject_runtime_config, generate archive SQL
    # in wrap step. can't do this right now because we actually need
    # to inspect status of the schema at runtime and archive requires
    # a lot of information about the schema to generate queries.
    template_ctx = context.copy()
    template_ctx.update(node_cfg)

    select = dbt.clients.jinja.get_rendered(
        dbt.templates.SCDArchiveTemplate, template_ctx
    )

    insert_stmt = dbt.templates.ArchiveInsertTemplate().wrap(
        schema=node_cfg.get("target_schema"),
        table=node_cfg.get("target_table"),
        query=select,
        unique_key=node_cfg.get("unique_key"),
    )

    node["wrapped_sql"] = dbt.clients.jinja.get_rendered(insert_stmt, template_ctx)

    result = adapter.execute_model(profile=profile, model=node)

    return result