HRQB 11 - Data Warehouse client and SQLQueryExtractTask base task

[ghukill]

Purpose and background context

This PR introduces class DWClient to assist with making queries to the Oracle Data Warehouse. The goal was to keep this as simple as possible, assisting with setting up a connection, supporting testing to a local SQLite database, and returning the results of a query as a pandas DataFrame.

Additionally, a new base task class was created SQLQueryExtractTask. This base class is designed to assist with Extract tasks that query for data from the warehouse, using the DWClient, and return a DataFrame for downstream tasks to work with. This leans into the requirement of PandasPickleTask to define a get_dataframe() method, which in turn this defines as an opinionated method to perform a SQL query.

The end result is a task that can look as simple as this:

class ExtractTable1AndTable2Data(SQLQueryExtractTask):
    pipeline = luigi.Parameter()
    stage = luigi.Parameter("Extract")

    @property
    def sql_query(self) -> str:
        return """
        select
           t1.column_a,
           t1.column_b,
           t2.column_c
        from my_table t1
        inner join my_other_table t2 on t1.column_a = t2.column_a
        where t2.column_d = 'foo'
        """

When this task runs, it will use the sql_query property defined, perform a query, and return the result as a DataFrame. The goal is to focus on writing the SQL queries themselves, and not the orchestration of executing queries or preparing the results for downstream tasks.

How can a reviewer manually see the effects of these changes?

An ipython shell is likely the easiest way to get a feel for the new base task type SQLQueryExtractTask and DWClient working together.

Set env vars:

WORKSPACE=dev
SENTRY_DSN=None
DYLD_LIBRARY_PATH=/usr/local/lib:$DYLD_LIBRARY_PATH
LUIGI_CONFIG_PATH=hrqb/luigi.cfg
QUICKBASE_API_TOKEN=abc123
QUICKBASE_APP_ID=def456

Ipython shell:

# import a testing Task class from fixtures
from tests.fixtures.tasks.extract import SQLExtractAnimalColors

# instantiate the Task with pipeline name and stage
task = SQLExtractAnimalColors(pipeline="Testing", stage="Extract")

# run get_dataframe() method
# this relies on the base class SQLQueryExtractTask to perform a SQL query
# (configured here to query a local SQLite database that has data)
# and returns a dataframe only, no files are written
task.get_dataframe()
# Out[3]: 
#   animal_id  color
# 0         42  green
# 1        101    red

# run actual run() method, and note the file created:
# output/Testing__Extract__SQLExtractAnimalColors.pickle
task.run()

# other confirmations that data was written identical
# to the PandasPickleTask this task class extends

task.path
# Out[4]: 'output/Testing__Extract__SQLExtractAnimalColors.pickle'

task.target
# Out[5]: <hrqb.base.target.PandasPickleTarget at 0x12de27e10>

task.target.exists()
# Out[6]: True

task.complete()
# Out[7]: True

Includes new or updated dependencies?

YES

Changes expectations for external applications?

YES; hypothetical connection to data warehouse, though credentials are not deployed yet

What are the relevant tickets?

• https://mitlibraries.atlassian.net/browse/HRQB-11

Developer

• All new ENV is documented in README
• All new ENV has been added to staging and production environments
• All related Jira tickets are linked in commit message(s)
• Stakeholder approval has been confirmed (or is not needed)

Code Reviewer(s)

• The commit message is clear and follows our guidelines (not just this PR message)
• There are appropriate tests covering any new functionality
• The provided documentation is sufficient for understanding any new functionality introduced
• Any manual tests have been performed or provided examples verified
• New dependencies are appropriate or there were no changes

[ghukill]

Unsure if I commented on this in tests, but using a factory=lambda: ... pattern here is helpful for testing. This ensures that only when the DWClient is instantiated, does it pull the env var DATA_WAREHOUSE_CONNECTION_STRING if an explicit connection string is not passed.

If this had been default=Config().DATA_WAREHOUSE_CONNECTION_STRING, then monkey patching the env var at pytest time would not be sufficient, as this class is already defined with the default value for any future instances being the env var before that monkey patching.

[jonavellecuerdo]

[Non-blocking] Maybe pull some of this information into the doc string or as a comment above the declaration of connection_string?

[ehanson8]

Seems solid, a few suggestions and questions

[ehanson8]

Looks good but still recommend updating that docstring. Edit: don't know why this didn't link properly 16ad2a7#r141964334

[jonavellecuerdo]

Looking good! Just have a few small suggestions / questions. 🤓

[jonavellecuerdo]

Should this be updated: self.default_engine_parameters -> self.engine_parameters?

[ghukill]

Good catch! Updated, and even better, can remove that altogether; I think setting engine_parameters explicitly is pretty transparent in how it works now.

[jonavellecuerdo]

Would field(default={"thick_mode": True}) be sufficient in this case as it doesn't call any methods to set the value? 🤔

[ghukill]

Yep! Given the DATA_WAREHOUSE_CONNECTIONS_STRING env var, these are the only SQLAlchemy engine parameters required (at least so far in testing).

[ghukill]

Related to this comment and another, updated the docstring for this class to explain the fields:

"""Client to provide Oracle Data Warehouse connection and querying.

Fields:
    connection_string: str
        - full SQLAlchemy connection string, e.g.
            - oracle: oracle+oracledb://user1:pass1@example.org:1521/ABCDE
            - sqlite: sqlite:///:memory:
        - defaults to env var DATA_WAREHOUSE_CONNECTION_STRING, loaded from env vars
        at time of DWClient initialization
    engine_parameters: dict
        - optional dictionary of SQLAlchemy engine parameters
    engine: Engine
        - set via self.init_engine()
"""

[jonavellecuerdo]

[Non-blocking] Maybe pull some of this information into the doc string or as a comment above the declaration of connection_string?

[jonavellecuerdo]

Where would this SQL files live and in what cases is a SQL file preferred over a string? 🤔

Hmm, what do you think of an update like:

class SQLQueryExtractTask(PandasPickleTask):
    """Base class for Tasks that make SQL queries for data."""

   ...

    @property
    def sql_query(self) -> str | None:
        """SQL query from string or loaded from file"""
        if self.sql_query:
            query = self.sql_query
        elif self.sql_file:
            with open(self.sql_file) as f:
                query = f.read()
        else:
            message = "Property sql_query or sql_file must be set."
            raise AttributeError(message)
        return query

    @property
    def sql_query_string(self) -> str | None:
        """SQL query from string to execute."""
        return None

    @property
    def sql_file(self) -> str | None:
        """SQL query loaded from file to execute."""
        return None

    @property
    def sql_query_parameters(self) -> dict:
        """Optional parameters to include with SQL query."""
        return {}

    def get_dataframe(self) -> pd.DataFrame:
        return self.dwclient.execute_query(
            self.sql_query,
            params=self.sql_query_parameters,
        )

[ghukill]

The question and refactor suggestion are well received. In fact, given some work since this PR on actually scoping out some extract tasks, it seems pretty clear that most (if not all) SQL extract tasks will use a file.

However, being able to easily override that with an explicit string is handy for testing. I think there is a middleground here:

• self.sql_query returns a string, defaulting to reading self.sql_file
• if neither are overridden, throws an exception that sql_file must be defined OR sql_query overridden to provide a string vs reading a file

Going to make that change and push a commit.

[ghukill]

@jonavellecuerdo - added a couple of commits addressing your feedback.

[jonavellecuerdo]

@ghukill Thank you for taking the time to make this update. Just one follow-up question!