Merge pull request #2 from IanRFerguson/ian/bigquery

BigQuery Connector

.github/workflows/ci_cd.yml → .github/workflows/automated_tests.yml

@@ -1,4 +1,4 @@
-name: checks
+name: formatting_checks
 
 on:
  pull_request:
@@ -7,7 +7,7 @@ on:
  branches: ["master", "dev"]
 
 jobs:
- run_checks: 
+ check_linting_and_formatting: 
  runs-on: macos-latest
  steps:
 
@@ -32,9 +32,26 @@ jobs:
  run: |
  python -m pip install black flake8 isort
 
- - name: Run Python checks
+ - name: Run Python formatting and linting
  if: steps.changed-python-files.outputs.any_changed == 'true'
  run: |
  flake8 ${{ steps.changed-python-files.outputs.all_changed_files }} --extend-ignore=E203,W503 --max-line-length=120
  black --check ${{ steps.changed-python-files.outputs.all_changed_files }}
- isort --profile black ${{ steps.changed-python-files.outputs.all_changed_files }}
+ isort --profile black ${{ steps.changed-python-files.outputs.all_changed_files }}
+ 
+ run_unit_tests:
+ runs-on: macos-latest
+ steps:
+ - name: Checkout branch
+ uses: actions/checkout@v3
+ with:
+ fetch-depth: 0
+
+ - name: Setup Python environment
+ uses: actions/setup-python@v3
+
+ - name: Run Unit Tests
+ run: |
+ pip install -r requirements.txt
+ pip install pytest
+ pytest -rf tests/*

.gitignore

@@ -4,4 +4,5 @@
 
 ### DIRECTORIES
 venv/
-.vscode/
+.vscode/
+dev/

README.md

@@ -2,12 +2,39 @@
 
 <img src="https://upload.wikimedia.org/wikipedia/en/d/d5/Klondike_logo.svg">
 
-Klondike offers a lightweight API to read and write data to Google BigQuery using rust-optimized Polars DataFrames.
+Klondike offers a lightweight API to read and write data to Google BigQuery using Polars DataFrames.
 
 ## Installation
 
+### Installing Klondike
 Install at the command line
 
 ```
 pip install klondike==0.1.0
+```
+
+### Installing Rust
+Since Polars leverages Rust speedups, you need to have Rust installed in your environment as well. See the Rust installation guide [here](https://www.rust-lang.org/tools/install).
+
+
+## Usage
+
+```
+# Instantiate a connection to BigQuery
+from klondike import BigQueryConnector
+
+bq = BigQueryConnector(
+ app_creds="/path/to/your/service_account.json"
+)
+
+# Read data from BigQuery
+sql = "SELECT * FROM nba_dbt.staging__nyk_players"
+df = bq.read_dataframe_from_bigquery(sql=sql)
+
+# Write data to BigQuery
+bq.write_dataframe_to_bigquery(
+ df=df,
+ table_name="nba_dbt.my_new_table",
+ if_eixsts="truncate"
+)
 ```

klondike/__init__.py

@@ -24,7 +24,7 @@
 
 __all__ = []
 POLAR_OBJECTS = [
- ("klondike.bigquery.bigquery", "PolarBigQuery"),
+ ("klondike.bigquery.bigquery", "BigQueryConnector"),
 ]
 
 for module_, object_ in POLAR_OBJECTS:

klondike/bigquery/bigquery.py

@@ -1,11 +1,14 @@
+import io
 import json
 import os
 import tempfile
-from contextlib import contextmanager
 from typing import Optional, Union
 
+import polars as pl
 from google.cloud import bigquery
-from google.cloud.bigquery import dbapi
+from google.cloud.bigquery import LoadJobConfig
+
+from klondike import logger
 
 ##########
 
@@ -15,39 +18,54 @@
  "https://www.googleapis.com/auth/bigquery",
  ]
  },
-)
+)[0]
 
 
-class PolarBigQuery:
+class BigQueryConnector:
  """
  Establish and authenticate a connection to a BigQuery warehouse
+
+ Args:
+ app_creds: Google service account, either as a relative path or a dictionary instance
+ project: Name of Google Project
+ location: Location of Google Project
+ timeout: Temporal threshold to kill a stalled job, defaults to 60s
+ client_options: API scopes
+ google_environment_variable: Provided for flexibility, defaults to `GOOGLE_APPLICATION_CREDENTIALS`
  """
 
  def __init__(
  self,
  app_creds: Optional[Union[str, dict]] = None,
- gcp_project: Optional[str] = None,
+ project: Optional[str] = None,
+ location: Optional[str] = None,
  timeout: int = 60,
  client_options: dict = SCOPES,
  google_environment_variable: str = "GOOGLE_APPLICATION_CREDENTIALS",
  ):
  self.app_creds = app_creds
- self.gcp_project = gcp_project
+ self.project = project
+ self.location = location
  self.timeout = timeout
  self.client_options = client_options
 
  self._client = None
  self.dialect = "bigquery"
- self._dbapi = dbapi
 
- self.__setup_google_app_creds(
- app_creds=app_creds, env_variable=google_environment_variable
- )
+ if not self.app_creds:
+ if not os.environ.get(google_environment_variable):
+ raise OSError("No app_creds provided")
+ else:
+ logger.info(
+ f"Using `{google_environment_variable}` variable defined in environment"
+ )
+ else:
+ self.__setup_google_app_creds(
+ app_creds=self.app_creds, env_variable=google_environment_variable
+ )
 
  def __setup_google_app_creds(self, app_creds: Union[str, dict], env_variable: str):
- """
- Sets runtime environment variablefor Google SDK
- """
+ "Sets runtime environment variable for Google SDK"
 
  if isinstance(app_creds, dict):
  creds = json.dumps(app_creds)
@@ -61,6 +79,70 @@ def __setup_google_app_creds(self, app_creds: Union[str, dict], env_variable: st
 
  os.environ[env_variable] = creds
 
+ def __set_load_job_config(
+ self,
+ base_job_config: Optional[LoadJobConfig] = None,
+ max_bad_records: int = 0,
+ table_schema: list = None,
+ if_exists: str = "fail",
+ **kwargs,
+ ):
+ "Defines `LoadConfigJob` when writing to BigQuery"
+
+ if not base_job_config:
+ logger.debug("No job config provided, starting fresh")
+ base_job_config = LoadJobConfig()
+
+ # This is default behavior for this class
+ base_job_config.source_format = bigquery.SourceFormat.PARQUET
+
+ # Create table schema mapping if provided
+ if table_schema:
+ base_job_config.schema = self.__set_table_schema(table_schema=table_schema)
+ else:
+ base_job_config.schema = None
+
+ base_job_config.max_bad_records = max_bad_records
+
+ base_job_config.write_disposition = self.__set_write_disposition(
+ if_exists=if_exists
+ )
+
+ # List of LoadJobConfig attributes
+ _attributes = [x for x in dict(vars(LoadJobConfig)).keys()]
+
+ # Attributes that will not be overwritten
+ _reserved_attributes = [
+ "source_format",
+ "schema",
+ "max_bad_records",
+ "write_disposition",
+ ]
+
+ # Loop through keyword arguments and update `LoadJobConfig` object
+ for k, v in kwargs.items():
+ if k in _attributes and k not in _reserved_attributes:
+ logger.debug(f"Updating {k} parameter in job config [value={v}]")
+ base_job_config[k] = v
+
+ return base_job_config
+
+ def __set_table_schema(self, table_schema: list):
+ "TODO - Write about me"
+
+ return [bigquery.SchemaField(**x) for x in table_schema]
+
+ def __set_write_disposition(self, if_exists: str):
+ "TODO - Write about me"
+
+ DISPOSITION_MAP = {
+ "fail": bigquery.WriteDisposition.WRITE_EMPTY,
+ "append": bigquery.WriteDisposition.WRITE_APPEND,
+ "truncate": bigquery.WriteDisposition.WRITE_TRUNCATE,
+ }
+
+ return DISPOSITION_MAP[if_exists]
+
  @property
  def client(self):
  """
@@ -73,17 +155,81 @@ def client(self):
  location=self.location,
  client_options=self.client_options,
  )
+
  return self._client
 
- @contextmanager
- def connection(self):
+ def read_dataframe_from_bigquery(self, sql: str) -> pl.DataFrame:
+ """
+ Executes a SQL query and returns a Polars DataFrame.
+ TODO - Make this more flexible and incorporate query params
+
+ Args:
+ sql: String representation of SQL query
+
+ Returns:
+ Polars DataFrame object
+ """
+
+ # Define query job
+ logger.debug("Running SQL...", sql)
+ query_job = self.client.query(query=sql, timeout=self.timeout)
+
+ # Execute and wait for results
+ result = query_job.result()
+
+ # Populate DataFrame using PyArrow
+ df = pl.from_arrow(result.to_arrow())
+
+ if df.is_empty():
+ logger.info("No results returned from SQL call")
+ return
+
+ logger.info(f"Successfully read {len(df)} rows from BigQuery")
+ return df
+
+ def write_dataframe_to_bigquery(
+ self,
+ df: pl.DataFrame,
+ table_name: str,
+ load_job_config: Optional[LoadJobConfig] = None,
+ max_bad_records: int = 0,
+ table_schema: list = None,
+ if_exists: str = "fail",
+ **load_kwargs,
+ ) -> None:
  """
- TODO - Fill me in
+ Writes a Polars DataFrame to BigQuery
+
+ Args:
+ df: Polars DataFrame
+ table_name: Destination table name to write to - `dataset.table` convention
+ load_job_config: `LoadJobConfig` object. If none is supplied, several defaults are applied
+ max_bad_records: Tolerance for bad records in the load job, defaults to 0
+ table_schema: List of column names, types, and optional flags to include
+ if_exists: One of `fail`, `drop`, `append`, `truncate`
+ load_kwargs: See here for list of accepted values
+ https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.LoadJobConfig
  """
 
- conn_ = self._dbapi.connect(self.client)
+ if if_exists == "drop":
+ self.client.delete_table(table=table_name)
+ if_exists = "fail"
+
+ load_job_config = self.__set_load_job_config(
+ base_load_conifg=load_job_config,
+ max_bad_records=max_bad_records,
+ table_schema=table_schema,
+ if_exists=if_exists,
+ **load_kwargs,
+ )
+
+ logger.info(f"Writing to {table_name}...")
+ with io.BytesIO() as stream_:
+ df.write_parquet(stream_)
+ stream_.seek(0)
+ load_job = self.client.load_table_from_file(
+ stream_, destination=table_name, job_config=load_job_config
+ )
 
- try:
- yield conn_
- finally:
- conn_.close()
+ load_job.result()
+ logger.info(f"Successfuly wrote {len(df)} rows to {table_name}")

requirements.txt

@@ -1,4 +1,3 @@
 google-cloud-bigquery==3.19.0
 polars==0.20.16
-petl==1.7.15
-psycopg==3.1.18
+pyarrow==15.0.2

run_checks.sh

@@ -0,0 +1,3 @@
+python -m flake8 --extend-ignore=E203,W503 --max-line-length=120 klondike/ tests/
+python -m black klondike/ tests/
+python -m isort --profile black klondike/ tests/