Polar Redshift + Polar Snowflake

polar_bare/bigquery/bigquery.py

@@ -1,13 +1,13 @@
+import json
+import os
+import tempfile
 from contextlib import contextmanager
 from typing import Optional, Union
 
 from google.cloud import bigquery
 from google.cloud.bigquery import dbapi
 
 from polar_bare.pbear.generic_db import PolarBareDB
-import json
-import os
-import tempfile
 
 ##########
 

polar_bare/pbear/generic_db.py

@@ -17,7 +17,11 @@ def __init__(self):
         self.__tempfiles = []
 
     def query(
-        self, sql: str, parameters: Optional[list] = None, batch_size: int = 250_000
+        self,
+        sql: str,
+        parameters: Optional[list] = None,
+        batch_size: int = 250_000,
+        return_values: bool = True,
     ) -> Optional[DataFrame]:
         """
         Leverages internal connection to query against Postgres instance
@@ -36,7 +40,7 @@ def query(
                 cursor.execute(query=sql, params=parameters)
 
                 # If no results, return without raising exeption
-                if not cursor.description:
+                if not cursor.description or not return_values:
                     return None
 
                 ###

polar_bare/redshift/redshift.py

@@ -0,0 +1,29 @@
+from polar_bare.postgres.postgres import PolarPostgres
+
+##########
+
+
+class PolarRedshift(PolarPostgres):
+    """
+    Establish and authenticate a connection to a Redshift cluster
+
+    Args:
+        host: Required if `PG_HOST` absent in runtime environment
+        db: Required if `PG_DBNAME` absent in runtime environment
+        port: Required if `PG_PORT` absent in runtime environment
+        username: Required if `PG_USERNAME` absent in runtime environment
+        password: Required if `PG_PASSWORD` absent in runtime environment
+        timeout: Defines connection timeout tolerance (default=60s)
+    """
+
+    def __init__(self, host, port, db, username, password, timeout):
+        super().__init__(
+            host=host,
+            db=db,
+            port=port,
+            username=username,
+            password=password,
+            timeout=timeout,
+        )
+
+        self.dialect = "redshift"

polar_bare/snowflake/snowflake.py

@@ -0,0 +1,49 @@
+from contextlib import contextmanager
+from typing import Optional
+
+import snowflake.connector
+
+from polar_bare.pbear.generic_db import PolarBareDB
+
+##########
+
+
+class PolarSnowflake(PolarBareDB):
+    """
+    Establish and authenticate a connection to a Snowflake warehouse
+    """
+
+    def __init__(
+        self,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        account: Optional[str] = None,
+        session_parameters: dict = {},
+    ):
+        self.__username = username
+        self.__password = password
+        self.account = account
+        self.session_parameters = session_parameters
+
+    @contextmanager
+    def connection(self):
+        """
+        TODO - Fill this in
+        """
+
+        connection = snowflake.connector.connect(
+            user=self.__username,
+            password=self.__password,
+            account=self.account,
+            session_parameters=self.session_parameters,
+        )
+
+        try:
+            yield connection
+        except snowflake.connector.Error:
+            connection.rollback()
+            raise
+        else:
+            connection.commit()
+        finally:
+            connection.close()