Add snowflake client

Author: aaronfriedman6

CHANGELOG.md

@@ -1,4 +1,9 @@
 # Changelog
+## v1.10.0 3/17/26
+- Add Snowflake client
+- Update config helper to allow loading config files without PLAINTEXT/ENCRYPTED structure
+- Update structured log helper to include name of the logger by default
+
 ## v1.9.1 3/11/26
 - Add merge_contextvars to default structlog configuration
 

README.md

@@ -12,6 +12,7 @@ This package contains common Python utility classes and functions.
 * Connecting to and querying a MySQL database
 * Connecting to and querying a PostgreSQL database
 * Connecting to and querying Redshift
+* Connecting to and querying Snowflake
 * Making requests to the Oauth2 authenticated APIs such as NYPL Platform API and Sierra
 * Interacting with vendor APIs such as cloudLibrary
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nypl_py_utils"
-version = "1.9.1"
+version = "1.10.0"
 authors = [
   { name="Aaron Friedman", email="aaronfriedman@nypl.org" },
 ]
@@ -70,6 +70,10 @@ secrets-manager-client = [
     "boto3>=1.26.5",
     "botocore>=1.29.5"
 ]
+snowflake-client = [
+    "nypl_py_utils[log-helper]",
+    "snowflake-connector-python>=4.3.0"
+]
 sftp-client = [
     "nypl_py_utils[log-helper]",
     "paramiko>=3.4.1"

src/nypl_py_utils/classes/snowflake_client.py

@@ -0,0 +1,99 @@
+import snowflake.connector as sc
+
+from nypl_py_utils.functions.log_helper import create_log
+
+
+class SnowflakeClient:
+    """Client for managing connections to Snowflake"""
+
+    def __init__(self, account, user, password, warehouse=None):
+        self.logger = create_log('snowflake_client')
+        self.conn = None
+        self.account = account
+        self.user = user
+        self.password = password
+        self.warehouse = warehouse
+
+    def connect(self, **kwargs):
+        """
+        Connects to a Snowflake database using the given credentials. If
+        warehouse parameter is None, uses the default warehouse for the user.
+
+        Parameters
+        ----------
+        kwargs:
+            All possible arguments (such as timeouts) can be found here:
+            https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#connect
+        """
+        self.logger.info('Connecting to Snowflake')
+        try:
+            self.conn = sc.connect(
+                account=self.account,
+                user=self.user,
+                password=self.password,
+                warehouse=self.warehouse,
+                **kwargs)
+        except Exception as e:
+            raise SnowflakeClientError(
+                f'Error connecting to Snowflake: {e}') from None
+
+    def execute_query(self, query, **kwargs):
+        """
+        Executes an arbitrary query against the given database connection.
+
+        Note that:
+            1) All results will be fetched by default, so this method is not
+                suitable if you do not want to load all rows into memory
+            2) AUTOCOMMIT is on by default, so this method is not suitable if
+               you want to execute multiple queries in a single transaction
+            3) This method can be used for both read and write queries, but
+               it's not optimized for writing -- there is no parameter binding
+               or executemany support, and the return value for write queries
+               can be unpredictable.
+
+        Parameters
+        ----------
+        kwargs:
+            All possible arguments (such as timeouts) can be found here:
+            https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#execute
+
+        Returns
+        -------
+        sequence
+            A list of tuples
+        """
+        self.logger.info('Querying database')
+        cursor = self.conn.cursor()
+        try:
+            try:
+                cursor.execute(query, **kwargs)
+                return cursor.fetchall()
+            except Exception:
+                raise
+            finally:
+                cursor.close()
+        except Exception as e:
+            # If there was an error, also close the database connection
+            self.close_connection()
+
+            short_q = str(query)
+            if len(short_q) > 2500:
+                short_q = short_q[:2497] + "..."
+            raise SnowflakeClientError(
+                f'Error executing Snowflake query {short_q}: {e}', self.logger
+            ) from None
+
+    def close_connection(self):
+        """Closes the connection"""
+        self.logger.info('Closing Snowflake connection')
+        self.conn.close()
+
+
+class SnowflakeClientError(Exception):
+    def __init__(self, message='', logger=None):
+        self.message = message
+        if logger is not None:
+            logger.error(message)
+
+    def __str__(self):
+        return self.message

src/nypl_py_utils/functions/config_helper.py

@@ -10,14 +10,13 @@
 
 def load_env_file(run_type, file_string):
     """
-    This method loads a YAML config file containing environment variables,
-    decrypts whichever are encrypted, and puts them all into os.environ as
-    strings. For a YAML variable containing a list of values, the list is
-    exported into os.environ as a json string and should be loaded as such.
+    This method reads a YAML config file containing environment variables and
+    loads them all into os.environ as strings. See _parse_yaml_dict for more.
 
-    It requires the YAML file to be split into a 'PLAINTEXT_VARIABLES' section
-    and an 'ENCRYPTED_VARIABLES' section. See config/sample.yaml for an example
-    config file.
+    If the config file is divided into 'PLAINTEXT_VARIABLES' and
+    'ENCRYPTED_VARIABLES' sections (see config/sample.yaml for an exmaple), the
+    'ENCRYPTED_VARIABLES' variables will be decrypted first. Otherwise, all
+    variables will be loaded as is.
 
     Parameters
         ----------
@@ -36,31 +35,53 @@ def load_env_file(run_type, file_string):
             try:
                 env_dict = yaml.safe_load(env_stream)
             except yaml.YAMLError:
-                logger.error('Invalid YAML file: {}'.format(open_file))
                 raise ConfigHelperError(
                     'Invalid YAML file: {}'.format(open_file)) from None
     except FileNotFoundError:
-        logger.error('Could not find config file {}'.format(open_file))
         raise ConfigHelperError(
             'Could not find config file {}'.format(open_file)) from None
 
     if env_dict:
-        for key, value in env_dict.get('PLAINTEXT_VARIABLES', {}).items():
-            if type(value) is list:
-                os.environ[key] = json.dumps(value)
-            else:
-                os.environ[key] = str(value)
-
-        kms_client = KmsClient()
-        for key, value in env_dict.get('ENCRYPTED_VARIABLES', {}).items():
-            if type(value) is list:
-                decrypted_list = [kms_client.decrypt(v) for v in value]
-                os.environ[key] = json.dumps(decrypted_list)
-            else:
-                os.environ[key] = kms_client.decrypt(value)
-        kms_client.close()
+        if ('PLAINTEXT_VARIABLES' in env_dict
+                or 'ENCRYPTED_VARIABLES' in env_dict):
+            _parse_yaml_dict(env_dict.get('PLAINTEXT_VARIABLES', {}))
+
+            kms_client = KmsClient()
+            _parse_yaml_dict(env_dict.get(
+                'ENCRYPTED_VARIABLES', {}), kms_client)
+            kms_client.close()
+        else:
+            _parse_yaml_dict(env_dict)
+
+
+def _parse_yaml_dict(yaml_dict, kms_client=None):
+    """
+    Loads YAML dict into os.environ. All values are stored as strings to match
+    how AWS Lambda environment variables are stored. For list variables, the
+    list is exported into os.environ as a json string.
+
+    If kms_client is not empty, decrypts the variables first.
+
+    Does not allow for sub-dictionaries.
+    """
+    for key, value in yaml_dict.items():
+        if type(value) is dict:
+            raise ConfigHelperError(
+                'Found sub-dictionary in YAML config') from None
+        elif type(value) is list:
+            val = [kms_client.decrypt(v)
+                   for v in value] if kms_client else value
+            os.environ[key] = json.dumps(val)
+        else:
+            val = kms_client.decrypt(value) if kms_client else value
+            os.environ[key] = str(val)
 
 
 class ConfigHelperError(Exception):
-    def __init__(self, message=None):
+    def __init__(self, message='', logger=None):
         self.message = message
+        if logger is not None:
+            logger.error(message)
+
+    def __str__(self):
+        return self.message

src/nypl_py_utils/functions/log_helper.py

@@ -12,30 +12,37 @@
 }
 
 
-# Configure structlog to be machine-readable first and foremost
-# while still making it easy for humans to parse
-# End result (without additional bindings) is JSON like this:
-# {
-#     "logger": "module param",
-#     "message": "this is a test log event",
-#     "level": "info",
-#     "timestamp": "2023-11-01 18:50:47"
-# }
 def get_structlog(module):
-    structlog.configure(
+    """
+    Standard logging without additional bindings looks as follows:
+    {
+        "level": "info",
+        "timestamp": "2026-01-01T12:00:00.613719Z",
+        "logger": "module param",
+        "message": "this is a test log event"
+    }
+
+    Note that: 1) using bind_contextvars will bind variables to *all* loggers
+    that have been created, and 2) you cannot use the same module name for a
+    structlog and for a standard logger
+    """
+    logger = logging.getLogger(module)
+    logger.addHandler(logging.StreamHandler(sys.stdout))
+    logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO').upper())
+    logger.propagate = False  # Prevents double logging
+
+    return structlog.wrap_logger(
+        logger,
         processors=[
             structlog.contextvars.merge_contextvars,
             structlog.processors.add_log_level,
             structlog.processors.TimeStamper(fmt='iso'),
+            structlog.stdlib.add_logger_name,
             structlog.processors.EventRenamer('message'),
             structlog.processors.JSONRenderer(),
-        ],
-        context_class=dict,
-        logger_factory=structlog.PrintLoggerFactory(),
+        ]
     )
 
-    return structlog.get_logger(module)
-
 
 def standard_logger(module):
     logger = logging.getLogger(module)
@@ -58,7 +65,7 @@ def standard_logger(module):
 
 
 def create_log(module, json=False):
-    if (json):
+    if json:
         return get_structlog(module)
     else:
         return standard_logger(module)