Snowflake Extractor hangs jobs it is a part of, connection must be closed.

[davcamer]

Expected Behavior

With appropriate configuration, the Snowflake Extractor should run to completion, as other extractors do.

Current Behavior

When running on our CI/CD server, where I'm running all of our jobs, jobs that use the SnowflakeMetadataExtractor were hanging indefinitely. Here is a snippet of the end of the log:

INFO:databuilder.loader.file_system_neo4j_csv_loader:Deleting directory /var/tmp/amundsen/snowflake_production/relationships
INFO:databuilder.loader.file_system_neo4j_csv_loader:Deleting directory /var/tmp/amundsen/snowflake_production/nodes
INFO:root:Job completed
INFO:snowflake.connector.cursor:query: [ROLLBACK]
[..] On Cancel Task: Kills child processes

The process hung a the second to last line, before I manually canceled it, resulting in the final line.

Possible Solution

A "monkey patch" like the following resolves the issue for me:

    class ClosingSnowflakeMetadataExtractor(SnowflakeMetadataExtractor):
        def close(self) -> None:
            logging.info("Closing the SnowflakeMetadataExtractor's _alchemy_extractor's connection.")
            self._alchemy_extractor.connection.close()

    snowflake_extractor = ClosingSnowflakeMetadataExtractor()
    noop_transformer = NoopTransformer()
    csv_loader = FsNeo4jCSVLoader()

    snowflake_structure_task = DefaultTask(extractor=snowflake_extractor, loader=csv_loader, transformer=noop_transformer)

This subclasses the existing SnowflakeMetadataExtractor, and reaches in to the SQLAlchemyExtractor that it delegates to, and closes the connection. This resolves the problem, and the [Rollback] line where the process was hanging before is now seen at the end of the extraction process, like I would expect:

INFO:root:Closing the SnowflakeMetadataExtractor's _alchemy_extractor's connection.
INFO:snowflake.connector.cursor:query: [ROLLBACK]
INFO:snowflake.connector.cursor:query execution done
INFO:databuilder.publisher.neo4j_csv_publisher:Publishing Node csv files [

Snowflake's documentation warns about not explicitly closing connections, and I think this is the consequence of failing to close the connection.

As the code stands today, the Task class already calls the extractor's close method. But the SQLAlchemyExtractor and SnowflakeMetadataExtractor don't implement close, so their method is a simple pass statement implemented on the Scoped class. I think a good approach to address this would be to implement close behavior for both SQLAlchemyExtractor and SnowflakeMetadataExtractor.

SQLAlchemy engines can also be disposed, and SQLAlchemy's documentation there suggests that some applications might benefit from not even using the connection pool, but I think that is getting further in to uncommon approaches.

I would be happy to submit a PR to add a concrete close method to both SQLAlchemyExtractor and SnowflakeMetadataExtractor, if that approach seems correct to others.

Steps to Reproduce

1. Run a SnowflakeMetadataExtractor job...
2. In a docker container...
3. Observe that it hangs after publishing with a [Rollback] message from the Snowflake driver.

I am not sure how much of this is specific to our environment, or if others were working around this by killing the tasks?

Context

I had to avoid fully automating our Snowflake ingests before fixing this problem.

Your Environment

• Amundsen version used: amundsen-databuilder==4.0.3
• Data warehouse stores: Snowflake, MySQL
• Deployment (k8s or native): Amundsen Databuilder is running in a python:3.7 Docker image under GoCD.

[raels0]

I've also faced this issue, another possible solution is adapting the base extractor for SQLAlchemy to use engine-level transactions, documented here and here.

[shuichiro-makigaki]

It is simple and good to implement a method SQLAlchemyExtractor.close(), I think. (and, diff size becomes small.)
As a basic of programming, opened connections always should be closed.

Anyway, I'm waiting for the Pull Request because I've also faced this issue.

[davcamer]

I've opened the PR here: amundsen-io/amundsendatabuilder#453

[feng-tao]

change is merged

[davcamer]

I realized that PR #453 would cause AttributeErrors if there were problems during init for any of the extractors: close would still be called but the attribute for the inner extractor or connection might be undefined or a None value.

I've raised a second PR to fix this, along with a small refactoring: amundsen-io/amundsendatabuilder#454