diff --git a/polar_bare/bigquery/bigquery.py b/polar_bare/bigquery/bigquery.py index 7390bcf..7e105de 100644 --- a/polar_bare/bigquery/bigquery.py +++ b/polar_bare/bigquery/bigquery.py @@ -1,3 +1,6 @@ +import json +import os +import tempfile from contextlib import contextmanager from typing import Optional, Union @@ -5,9 +8,6 @@ from google.cloud.bigquery import dbapi from polar_bare.pbear.generic_db import PolarBareDB -import json -import os -import tempfile ########## diff --git a/polar_bare/pbear/generic_db.py b/polar_bare/pbear/generic_db.py index c641df4..abc4e9d 100644 --- a/polar_bare/pbear/generic_db.py +++ b/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 ### diff --git a/polar_bare/redshift/redshift.py b/polar_bare/redshift/redshift.py new file mode 100644 index 0000000..8086aa3 --- /dev/null +++ b/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" diff --git a/polar_bare/snowflake/snowflake.py b/polar_bare/snowflake/snowflake.py new file mode 100644 index 0000000..25e045d --- /dev/null +++ b/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()