<filename>pygyver/etl/dw.py
""" Module containing BigQueryExecutor for Python """
import os
import re
import logging
import time
import json
import pandas as pd
from pandas._testing import assert_frame_equal
from google.cloud import bigquery
from google.cloud.exceptions import NotFound
from google.oauth2 import service_account
from google.api_core import exceptions
from pygyver.etl.lib import bq_token_file_valid
from pygyver.etl.lib import bq_token_file_path
from pygyver.etl.lib import bq_default_project
from pygyver.etl.lib import bq_default_dataset
from pygyver.etl.lib import gcs_default_bucket
from pygyver.etl.lib import read_table_schema_from_file
from pygyver.etl.lib import bq_start_date
from pygyver.etl.lib import bq_end_date
from pygyver.etl.lib import set_write_disposition, set_priority
from pygyver.etl.toolkit import (
    date_lister,
    flatten_df_columns,
    validate_date
)
from pygyver.etl.gs import load_gs_to_dataframe
from pygyver.etl.storage import GCSExecutor
from pygyver.etl.db import DBExecutor


class BigQueryExecutorError(Exception):
    pass


def print_kwargs_params(func):
    def inner(*args, **kwargs):
        logging.debug("Keyword args applied to the template:")
        for key, value in kwargs.items():
            if key in forbiden_kwargs():
                raise KeyError("{} is a forbidden keyword argument.".format(key))
        for key, value in kwargs.items():
            logging.debug("%s = %s" % (key, value))
        return func(*args, **kwargs)
    return inner


def forbiden_kwargs():
    return ['partition_date']


class SafeDict(dict):
    def __missing__(self, key):
        return '{' + key + '}'

@print_kwargs_params
def read_sql(file, *args, **kwargs):
    ''' Read SQL file and apply arguments/keyword arguments.

    If dataset_prefix is a kwarg then this function is being used in a dry-run test.
    In this case, the dataset_prefix is added to the dataset name and project names
    are removed from the SQL.

    Args:
        file (string): path to the SQL file from PROJECT_ROOT environment variable.
        *kwargs can be passed if some parameters are to be passed.

    Returns:
        a SQL formated with **kwargs if applicable.

    Example:
        With SQL as:
            "select .. {param2} .. {param1} .. {paramN}"
        *kwargs as:
            param1=value1
            param2=value2
            paranN=valueN
        The functions returns:
            "select .. value2 .. value1 .. valueN"
    '''
    path_to_file = os.path.join(os.getenv("PROJECT_ROOT"), file)
    file = open(path_to_file, 'r')
    sql = file.read()
    file.close()
    if kwargs.get('dataset_prefix', None) is not None:
        sql_split = sql.split("`")
        for index, table in enumerate(sql_split):
            if index%2==1 and "." in table:
                split_table = table.split(".")
                split_table[-2] = kwargs.get('dataset_prefix', '') + split_table[-2]
                sql_split[index] = '.'.join(split_table[-2:])
        sql = '`'.join(sql_split)

    if len(kwargs) > 0:
        sql = sql.format_map(SafeDict(**kwargs))
    return sql


class BigQueryExecutor:
    """ BigQuery handler

    Parameters:
        project_id (sql_file): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.

    Required:
        GOOGLE_APPLICATION_CREDENTIALS (env variable).

    Attributes:
        client (Client object)
        credentials (Credentials object)
        project_id (string): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.

    Returns:
        a BigQueryExecutor object.
    """
    def __init__(self):
        """ Resets client and credentials.
        """
        self.client = None
        self.credentials = None
        self.auth()

    def auth(self):
        """ Sets BigQuery client.
        """
        bq_token_file_valid()
        self.credentials = service_account.Credentials.from_service_account_file(
            bq_token_file_path(),
            scopes=[
                "https://www.googleapis.com/auth/drive",
                "https://www.googleapis.com/auth/bigquery"
            ]
        )
        self.client = bigquery.Client(
            project = bq_default_project(),
            credentials=self.credentials
        )


    def get_dataset_ref(self,  dataset_id, project_id=bq_default_project()):
        """ Returns BigQuery DatasetReference object.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            bigquery.dataset.DatasetReference object
        """
        return bigquery.dataset.DatasetReference(
            project = project_id,
            dataset_id = dataset_id
        )

    def get_table_ref(self, dataset_id, table_id, project_id=bq_default_project()):
        """ Returns BigQuery Table reference object.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            Returns BigQuery Table reference object.
        """
        dataset_ref = self.get_dataset_ref(dataset_id, project_id = project_id)
        return dataset_ref.table(table_id)

    def dataset_exists(self, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Checks if a BigQuery dataset exists.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            True is the dataset exists, False otherwise
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        try:
            self.client.get_dataset(dataset_ref)
            return True
        except NotFound:
            return False

    def delete_dataset(self, dataset_id=bq_default_dataset(),project_id=bq_default_project(),delete_contents=False):
        """ Deletes a BigQuery dataset.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            delete_contents (bool): If True, delete all the tables in the dataset.
            If False and the dataset contains tables, the request will fail. Default is False.
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        try:
            self.client.delete_dataset(
                dataset_ref,
                delete_contents=delete_contents
            )
            logging.info(
                "Dataset %s:%s deleted",
                project_id,
                dataset_id
            )
            time.sleep(1)
        except exceptions.Conflict as error:
            logging.error(error)


    def create_dataset(self, dataset_id=bq_default_dataset(), project_id=bq_default_project(), **kwargs):
        """ Creates a BigQuery dataset if the dataset does not exists. Otherwise pass.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        if self.dataset_exists(dataset_id,project_id=project_id):
            logging.info(
                "Dataset %s already exists in project %s",
                dataset_id,
                project_id
            )
        else:
            try:
                self.client.create_dataset(dataset_ref)
                logging.info(
                    "Created dataset %s in in project %s",
                    dataset_id,
                    project_id
                )
            except exceptions.Conflict as error:
                logging.error(error)

    def table_exists(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Checks if a BigQuery table exists.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.


        Returns:
            True is the table exists, False otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        try:
            self.client.get_table(table_ref)
            return True
        except NotFound:
            return False

    def delete_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Delete a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        """
        try:
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            self.client.delete_table(table_ref)
            logging.info(
                'Table %s:%s.%s deleted.',
                project_id,
                dataset_id,
                table_id
            )
            time.sleep(1)
        except NotFound as error:
            logging.error(error)

    def initiate_table(self,
                       table_id,
                       schema_path,
                       dataset_id=bq_default_dataset(),
                       project_id=bq_default_project(),
                       partition=False,
                       partition_field='_PARTITIONTIME',
                       clustering=None):
        """ Initiate a BigQuery table. If the table already exists, compares the schema_path and apply a patch if there is a schema change.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string: Path to the BigQuery table schema from the PROJECT_ROOT environement variable.
            partition (bool): Specify whether the BigQuery table is partioned. Default to False.
            partition_field (string): Specify partition_field if table partioned. Defaults to "_PARTITIONTIME".
            clustering (list): List of clustering fields. Defaults to None.
        """
        if self.table_exists(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id):
            logging.info("Table %s.%s already exists in project %s",
                         dataset_id,
                         table_id,
                         project_id)
            self.apply_patch(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path
            )
        else:
            schema = read_table_schema_from_file(schema_path)
            table = bigquery.Table(
                self.get_table_ref(dataset_id, table_id, project_id = project_id),
                schema=schema
                )
            if partition:
                table.time_partitioning = bigquery.table.TimePartitioning(type_='DAY')
                if partition_field and partition_field != '_PARTITIONTIME':
                    if isinstance(partition_field, str):
                        table.time_partitioning = bigquery.table.TimePartitioning(
                            field=partition_field
                        )
                    else:
                        raise ValueError("partition_field should be a string")
            table.clustering_fields = clustering
            try:
                table = self.client.create_table(table)
                logging.info(
                    'Created table %s.%s in in project %s',
                    dataset_id,
                    table_id,
                    project_id
                )
            except exceptions.Conflict as error:
                logging.error(error)

    def create_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), sql=None, file=None,
                     write_disposition='WRITE_TRUNCATE', use_legacy_sql=False,
                     location='US', schema_path='',
                     partition=False,
                     partition_field='_PARTITIONTIME',
                     clustering=None,
                     priority='INTERACTIVE',
                     description=None,
                     **kwargs):
        """ create a bigquery table from a sql query """

        if sql is None and file is None:
            raise BigQueryExecutorError("Either SQL or file containing the SQL must be provided")

        if sql is None:
            sql = read_sql(file, **kwargs)

        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                partition=partition,
                partition_field=partition_field,
                dataset_id=dataset_id,
                project_id=project_id
            )
            if write_disposition == "WRITE_TRUNCATE":
                self.truncate_table(dataset_id=dataset_id, table_id=table_id,project_id=project_id)
                write_disposition = "WRITE_EMPTY"
        else:
            pass

        job_config = bigquery.QueryJobConfig()
        job_config.destination = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        job_config.write_disposition = set_write_disposition(write_disposition)
        job_config.use_legacy_sql = use_legacy_sql
        job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED
        job_config.priority = set_priority(priority)
        if partition:
            if not partition_field or partition_field == '_PARTITIONTIME':
                job_config.time_partitioning = bigquery.table.TimePartitioning(
                    type_='DAY'
                )
                job_config.clustering_fields = clustering
            elif isinstance(partition_field, str):
                job_config.time_partitioning = bigquery.table.TimePartitioning(
                    field=partition_field
                )
            else:
                raise ValueError("partition_field should be a string")

        query_job = self.client.query(
            sql,
            location=location,
            job_config=job_config
        )
        if priority == 'INTERACTIVE':
            query_job.result()
            logging.info(
            'Query results loaded to table %s:%s.%s',
                project_id,
                dataset_id,
                table_id
            )
        if description:
            self.update_table_description(table_id=table_id, description=description,
                                          project_id=project_id, dataset_id=dataset_id)

        return query_job


    def delete_backfill_partitions(self,
                                   backfill_dataset_id,
                                   table_id,
                                   partition_field,
                                   from_partition,
                                   to_partition
                                   ):
        """
        Delete partitions from table in backfill dataset between from_partition and to_partition

        Args:
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): table in backfill dataset that the partitions will be deleted from
            partition_field (string): the time-unit column that the table is partitioned by
            from_partition (date): date of the first partition to be deleted
            to_partition (date): date of the latest partition to be deleted

        Returns:
            None
        """
        client = bigquery.Client()

        sql = f"""
               DELETE
                 `{backfill_dataset_id}.{table_id}`
               WHERE
                 CAST({partition_field} AS DATE) BETWEEN
                   CAST('{from_partition}' AS DATE) AND CAST('{to_partition}' AS DATE)
               """

        query_job = client.query(
            sql
        )

        query_job.result() # wait for job to complete


    def run_full_backfill(self,
                          dataset_id,
                          backfill_dataset_id,
                          table_id,
                          schema_path,
                          sql,
                          partition_field
                          ):
        """
        Run a full backfill for a time-unit column partitioned table

        Args:
            dataset_id (string): dataset of the table to be backfilled
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by

        Returns:
            None
        """
        logging.info("Starting full backfill")

        logging.info("Creating backfill dataset if it doesn't exist")
        self.create_dataset(
            dataset_id=backfill_dataset_id
        )

        logging.info("Creating table in backfill dataset with all partitions")
        self.create_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            sql=sql,
            partition=True,
            partition_field=partition_field,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Replacing prod table with backfill table")
        self.copy_table(
            source_dataset_id=backfill_dataset_id,
            source_table_id=table_id,
            dest_dataset_id=dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting backfill table")
        self.delete_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id
        )


    def run_partial_backfill(self,
                             dataset_id,
                             backfill_dataset_id,
                             table_id,
                             schema_path,
                             sql,
                             partition_field,
                             from_partition,
                             to_partition
                             ):
        """
        Run a partial backfill for a time-unit column partitioned table between the dates
        from_partition and to_partition

        Args:
            dataset_id (string): dataset of the table to be backfilled
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by
            from_partition (date): date of the first partition to be backfilled
            to_partition (date): date of the latest partition to be backfilled

        Returns:
            None
        """
        logging.info("Starting partial backfill")

        logging.info("Creating backfill dataset if it doesn't exist")
        self.create_dataset(
            dataset_id=backfill_dataset_id
        )

        logging.info("Copying table to backfill dataset")
        self.copy_table(
            source_dataset_id=dataset_id,
            source_table_id=table_id,
            dest_dataset_id=backfill_dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting partitions to be backfilled")
        self.delete_backfill_partitions(
            backfill_dataset_id,
            table_id,
            partition_field,
            from_partition,
            to_partition
        )

        logging.info("Appending new backfill partitions")
        self.create_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            sql=sql,
            partition=True,
            partition_field=partition_field,
            write_disposition="WRITE_APPEND"
        )

        logging.info("Replacing prod table with backfill table")
        self.copy_table(
            source_dataset_id=backfill_dataset_id,
            source_table_id=table_id,
            dest_dataset_id=dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting backfill table")
        self.delete_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id
        )


    def run_backfill(self,
                     backfill,
                     dataset_id,
                     table_id,
                     schema_path,
                     sql,
                     partition_field,
                     from_partition,
                     to_partition
                     ):
        """
        Run a full or partial backfill for a time-unit column partitioned table

        Args:
            backfill (string): specifies full or partial backfill (if running a backfill)
            dataset_id (string): dataset of the table to be backfilled
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by

        Returns:
            None
        """
        logging.info("Starting backfill")

        backfill_dataset = "backfill"

        if backfill not in ("full", "partial"):
            raise BigQueryExecutorError("BACKFILL env var needs to be 'full' or 'partial'")

        if backfill == "full":
            self.run_full_backfill(
                dataset_id=dataset_id,
                backfill_dataset_id=backfill_dataset,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field
            )

        elif backfill == "partial":
            self.run_partial_backfill(
                dataset_id=dataset_id,
                backfill_dataset_id=backfill_dataset,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field,
                from_partition=from_partition,
                to_partition=to_partition
            )

    def create_partitions_legacy(self,
                                 sql,
                                 file,
                                 partition_dates,
                                 table_id,
                                 project_id,
                                 dataset_id,
                                 bq_start_date,
                                 bq_end_date,
                                 write_disposition,
                                 use_legacy_sql,
                                 partition_field,
                                 clustering,
                                 priority,
                                 **kwargs
                                 ):
        if sql is None:
            sql = read_sql(file, **kwargs)

        if partition_dates is None:
            existing_dates = self.get_existing_partition_dates(
                table_id=table_id,
                project_id=project_id,
                dataset_id=dataset_id
            )
            dates = self.get_partition_dates(
                start_date=bq_start_date,
                end_date=bq_end_date,
                existing_dates=existing_dates
            )
        else:
            self.validate_partition_dates(
                partition_dates=partition_dates
            )
            dates = partition_dates

        jobs = []
        for date in dates:
            partition_name = self.set_partition_name(table=table_id, date=date)
            logging.info("Updating partition: %s", partition_name)
            job = self.create_table(
                sql=self.apply_partition_filter(
                    sql=sql,
                    date=date
                ),
                dataset_id=dataset_id,
                project_id=project_id,
                table_id=partition_name,
                write_disposition=write_disposition,
                use_legacy_sql=use_legacy_sql,
                partition=True,
                partition_field=partition_field,
                clustering=clustering,
                priority=priority
            )
            jobs.append(job)

        if priority == 'BATCH':
            for job in jobs:
                job.result()


    def create_partition_table(self,
                               table_id,
                               dataset_id=bq_default_dataset(),
                               project_id=bq_default_project(),
                               bq_start_date=bq_start_date(),
                               bq_end_date=bq_end_date(),
                               sql=None,
                               file=None,
                               use_legacy_sql=False,
                               write_disposition='WRITE_TRUNCATE',
                               partition_dates=None,
                               pipeline_partition_date=None,
                               partition_field='_PARTITIONTIME',
                               clustering=None,
                               priority="INTERACTIVE",
                               schema_path="",
                               description=None,
                               **kwargs
                              ):
        """
        New and legacy ways of creating partition tables:

        Legacy:
        Partition to be generated are either passed through partition_dates or automatically generated using existing partitions.
        To filter on a specific partition, the filter DATE(_PARTITIONTIME) = {partition_date} can be used in your sql query.

        New:
        pipeline_partition_date is required for a normal pipeline run and is a single date.
        To run a backfill, the BACKFILL env var needs to be set to 'full' or 'partial' and an optional date range 
        can be set for a partial backfill by setting the FROM_PARTITION and TO_PARTITION env vars.
        """
        if sql is None and file is None:
            raise BigQueryExecutorError("Either SQL or file containing the SQL must be provided")

        if schema_path != "":
            self.initiate_table(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                partition=True,
                partition_field=partition_field,
                clustering=clustering
            )

        if not self.table_exists(dataset_id=dataset_id, table_id=table_id,project_id=project_id):
            raise BigQueryExecutorError("To create a partition, please initiate the table first using initiate_table.")

        if "BACKFILL" not in os.environ:
            if pipeline_partition_date is None: # legacy approach
                self.create_partitions_legacy(
                    sql=sql,
                    file=file,
                    partition_dates=partition_dates,
                    table_id=table_id,
                    project_id=project_id,
                    dataset_id=dataset_id,
                    bq_start_date=bq_start_date,
                    bq_end_date=bq_end_date,
                    write_disposition=write_disposition,
                    use_legacy_sql=use_legacy_sql,
                    partition_field=partition_field,
                    clustering=clustering,
                    priority=priority,
                    **kwargs
                )

            else: # new approach, no backfill
                from_partition = pipeline_partition_date
                to_partition = pipeline_partition_date
                sql = read_sql(
                    file=file,
                    from_partition=from_partition,
                    to_partition=to_partition,
                    **kwargs
                )

                # change date format from YYYY-MM-DD to YYYYMMDD for partition_id 
                partition_id = pipeline_partition_date.replace("-", "")
                partition_name = self.set_partition_name(
                    table=table_id,
                    date=partition_id
                )

                logging.info("Updating partition: %s", partition_name)
                self.create_table(
                    project_id=project_id,
                    dataset_id=dataset_id,
                    table_id=partition_name,
                    sql=sql,
                    partition=True,
                    partition_field=partition_field,
                    write_disposition=write_disposition,
                    use_legacy_sql=use_legacy_sql,
                    clustering=clustering,
                    priority=priority
                )

        else: # new approach, run backfill
            from_partition = os.environ.get("FROM_PARTITION", bq_start_date)
            to_partition = os.environ.get("TO_PARTITION", bq_end_date)
            sql = read_sql(
                file=file,
                from_partition=from_partition,
                to_partition=to_partition,
                **kwargs
            )
            self.run_backfill(
                backfill=os.environ["BACKFILL"],
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field,
                from_partition=from_partition,
                to_partition=to_partition
            )

        if description:
            self.update_table_description(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                description=description
            )


    def apply_partition_filter(self, sql, date):
        """ Apply partition_date to the SQL query.

        Parameters:
            sql (string): the SQL query
            partition_date (string): the partiton date

        Returns:
            Formatted SQL query.
        """
        return sql.format(
            partition_date=date
        )

    def validate_partition_dates(self, partition_dates):
        """ Validates the partition_dates parameter. Checks whether it is a list and the elements are in the right format.

        Parameters:
            partition_dates (list of string): list of partition_dates
            partition_date (string): the partiton date

        Raises:
            BigQueryExecutorError if the partition_dates parameter in not a list.
            ValueError if incorrect date format.
        """
        if not isinstance(partition_dates, list):
            raise BigQueryExecutorError("Partition dates need to be a list of date eg ['YYYYmmdd']")
        else:
            for date in partition_dates:
                validate_date(date=date, format='%Y%m%d')

    def set_partition_name(self, table, date):
        """ Validates the date and sets the partition_name associated.

        Parameters:
            table (string): table name
            date (string): date with format '%Y%m%d'

        Returns:
            Partition name

        Raises:
            BigQueryExecutorError if the partition_dates parameter in not a list.
            ValueError if incorrect date format.
        """
        validate_date(date=date, format='%Y%m%d')
        return table + "$" + date.replace("-", "")

    def get_partition_dates(self, start_date, end_date, existing_dates):
        """ Returns the partitions dates required based on start_date, end_date and existing_dates.

        Parameters:
            start_date (string): First partition required.
            end_date (string): Last partition required.
            existing_dates (list of string): List of existing partition.

        Returns:
            List of partition dates required.
        """
        partition_dates = []
        required_dates = date_lister(start_date=start_date, end_date=end_date)
        if existing_dates == []:
            for date in required_dates:
                partition_date = date.replace("-", "")
                partition_dates.append(partition_date)
        else:
            for date in required_dates:
                partition_date = date.replace("-", "")
                if partition_date not in existing_dates:
                    partition_dates.append(partition_date)
        return partition_dates

    def get_existing_partition_query(self, dataset_id, table_id,project_id = bq_default_project()):
        """ Gets existing partitions from BigQuery as dataframe.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            DataFrame of existing partitions
        """
        sql = """ SELECT
                    FORMAT_DATE('%Y%m%d', DATE(_PARTITIONTIME)) AS partition_id
                  FROM
                    `{dataset_id}.{table_id}`
                  GROUP BY
                    1 """.format(
                        dataset_id=dataset_id,
                        table_id=table_id
                        )
        return self.execute_sql(
            sql=sql,
            project_id=project_id
        )

    def get_existing_partition_dates(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Gets existing partitions.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            List of existing partition dates.
        """
        if not self.table_exists(dataset_id=dataset_id, table_id=table_id,project_id=project_id):
            existing_partition_dates = []
        else:
            res = self.get_existing_partition_query(dataset_id=dataset_id, table_id=table_id)
            # checks that res has number of rows > 0
            if res.shape[0] > 0:
                existing_partition_dates = res['partition_id'].to_list()
            # if not, no existing partitions
            else:
                existing_partition_dates = []
        return existing_partition_dates

    def get_table_schema(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table schema object

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.


        Returns:
            Table schema.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        table_schema = self.client.get_table(table_ref).schema
        return table_schema

    def get_table_partitioning_type(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table partitioning_type

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            partitioning_type is 'DAY' if the table is partitioned, None otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        table_properties = self.client.get_table(table_ref)._properties
        partitioning_type = table_properties.get('timePartitioning', {}).get('type')
        return partitioning_type

    def get_table_attributes(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        dict_of_attributes = {}
        attributes = ['clustering_fields', 'description', 'encryption_configuration', 'expires',
        'external_data_configuration', 'friendly_name', 'labels', 'range_partitioning',
        'require_partition_filter', 'schema', 'time_partitioning']
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        for attribute in attributes:
            my_attribute = getattr(self.client.get_table(table_ref), attribute)
            dict_of_attributes.update({attribute: my_attribute})
        return dict_of_attributes


    def get_table_clustering_fields(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table clustering_fields

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            list of clustering fields is the table has cluster fields, None otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        clustering_fields = self.client.get_table(table_ref).clustering_fields
        return clustering_fields

    def identify_new_fields(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Identify new fields in based on a schema file.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to the schema to compare to.

        Returns:
            List of new fields.
        """
        list_field = []
        schema_a = self.get_table_schema(
            table_id=table_id,
            dataset_id=dataset_id,
            project_id=project_id
        )
        schema_b = read_table_schema_from_file(schema_path)
        field_list_a = [schema_field.name for schema_field in schema_a]
        for schema_field in schema_b:
            if schema_field.name not in field_list_a:
                list_field.append(schema_field)
        return list_field

    def append_field(self, table_id, field, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Append fields to a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            field (string): Schema field object.

        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        table = self.client.get_table(table_ref)  # API request

        original_schema = table.schema
        new_schema = original_schema[:]  # creates a copy of the schema
        new_schema.append(field)

        table.schema = new_schema
        table = self.client.update_table(table, ["schema"])  # API request
        assert len(table.schema) == len(original_schema) + 1 == len(new_schema)
        return 0

    def apply_patch(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Apply a patch to a BigQuery Table if required.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file to compare to.

        Returns:
            Lenght of new schema
        """
        logging.info("Attempting patch")
        logging.info("Checking for new fields...")
        new_fields = self.identify_new_fields(
            table_id=table_id,
            schema_path=schema_path,
            project_id=project_id,
            dataset_id=dataset_id
        )
        if new_fields != []:
            logging.info("New fields to be added:")
            logging.info(new_fields)
            for field in new_fields:
                self.append_field(
                    field=field,
                    table_id=table_id,
                    dataset_id=dataset_id,
                    project_id=project_id
                )
            logging.info("Done!")
        else:
            logging.info("No field to be added")

        logging.info("Checking for schema update...")
        self.update_schema(
            table_id=table_id,
            schema_path=schema_path,
            dataset_id=dataset_id,
            project_id=project_id
        )
        return len(
            self.get_table_schema(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
            )
            )

    def update_schema(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Perform a schema update. Used to update descriptions.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file to compare to.

        Raises:
            BadRequest if the update fails.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)  # API request
        new_schema = read_table_schema_from_file(schema_path)
        if table.schema == new_schema:
            logging.info("No changes needed")
        else:
            assert len(table.schema) == len(new_schema)
            table.schema = new_schema
            try:
                table = self.client.update_table(table, ["schema"])  # API request
                return 0
            except exceptions.BadRequest as error:
                raise error

    def update_table_description(
            self,
            table_id,
            description,
            project_id=bq_default_project(),
            dataset_id=bq_default_dataset()
            ):
        """ Performs a table update to fill in description for the table.

        Parameters:
            table_id (string): BigQuery table ID.
            description (string): The descriptive text to describe the content of the table.
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.

        Raises:
            BadRequest if the update fails.
        """

        table_ref = self.get_table_ref(dataset_id=dataset_id, table_id=table_id, project_id=project_id)
        table = self.client.get_table(table_ref)  # API request

        if table.description == description:
            logging.info("No changes to table description required")
        else:
            try:
                table.description = description
                self.client.update_table(table, ["description"])  # API request
            except exceptions.BadRequest as error:
                raise error

    def execute_sql(self, sql, project_id=bq_default_project(), dialect='standard'):
        """ Executes a SQL query and loads it as a DataFrame.

        Parameters:
            sql (string): SQL Query.
            project_id (string): BigQuery Project ID.
            dialect (string): BigQuery dialect. Defaults to standard.

        Returns:
            Query result as a DataFrame.
        """
        data = pd.read_gbq(
            sql,
            project_id=project_id,
            credentials=self.credentials,
            dialect=dialect
        )

        return data

    def execute_dml(self, sql=None, file=None, location='US', project=bq_default_project()):

        job_config = bigquery.QueryJobConfig()

        try:

            sql = sql if sql else read_sql(file)
        
            job = self.client.query(sql, location=location, job_config=job_config)
            
            results = job.result()
            
            logging.info("%s - %s - %s rows effected (%s elapsed, %s B processed)",
                job.statement_type,
                job.state,
                job.num_dml_affected_rows,
                job.ended - job.started,
                job.total_bytes_processed
            )
        
        except Exception as e:
            logging.error(e)
            raise e

        return

    def execute_file(self, file, project_id=bq_default_project(),
                     dialect='standard', *args, **kwargs):
        """ Executes a SQL file and loads it as a DataFrame.

        Parameters:
            file (string): Path to SQL file.
            project_id (string): BigQuery Project ID.
            dialect (string): BigQuery dialect. Defaults to standard.

        **kwargs can be passed if the SQL file contains arguments formatted with {}.

        Forbidden kwargs:
            partition_date

        Returns:
            Query result as a DataFrame.
        """
        sql = read_sql(file, *args, **kwargs)
        data = self.execute_sql(
            sql=sql,
            project_id=project_id,
            dialect=dialect
        )
        return data

    def load_dataframe(self, df, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads DataFrame to BigQuery table.

        Parameters:
            df (pd.DataFrame): Pandas DataFrame
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id,
                schema_path=schema_path
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        data = df.rename(columns=lambda cname: cname.replace('.', '_'))
        table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        job_config = bigquery.LoadJobConfig(schema=schema)
        job_config.write_disposition = set_write_disposition(write_disposition)
        job = self.client.load_table_from_dataframe(
            data,
            table_ref,
            job_config=job_config
        )
        job.result()


    def create_gs_table(self,
                        table_id,
                        dataset_id=bq_default_dataset(),
                        project_id=bq_default_project(),
                        schema_path='',
                        googlesheet_uri=None,
                        googlesheet_key=None,
                        sheet_name=None,
                        header=True):
        """ Creates BigQuery Table with live connection to Google Sheets

        Args:
            table_id (str): BigQuery table ID
            dataset_id (str): BigQuery dataset ID
            project_id (str): BigQuery project ID
            schema_path (str): Path to schema file, if not set then BQ will auto-detect
            googlesheet_uri (str): Google Sheet URI
            googlesheet_key (str): Google Sheet Key, an alternate option instead of URI
            sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0
            header (bool): Defaults to True
        """
        if googlesheet_uri is None:
            if googlesheet_key is None:
                raise BigQueryExecutorError("A googlesheet_uri or googlesheet_key must be provided")
            else:
                googlesheet_uri = f"https://docs.google.com/spreadsheets/d/{googlesheet_key}"

        external_config = bigquery.ExternalConfig("GOOGLE_SHEETS")
        external_config.source_uris = [googlesheet_uri]
        if sheet_name:
            external_config.options.range = (sheet_name)

        if header:
            external_config.options.skip_leading_rows = 1

        if schema_path != '':
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None
            external_config.autodetect = True

        gs_table = bigquery.Table(
            self.get_table_ref(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id
            ),
            schema=schema
        )
        gs_table.external_data_configuration=external_config

        try:
            self.client.delete_table(gs_table, not_found_ok=True)
            self.client.create_table(gs_table)
            logging.info(
                f"Created table {project_id}:{dataset_id}.{table_id} with live connection to {googlesheet_uri}"
            )
        except exceptions.Conflict as error:
            logging.error(error)
            raise error


    def load_google_sheet(self,
                        table_id,
                        dataset_id=bq_default_dataset(),
                        project_id=bq_default_project(),
                        schema_path='',
                        googlesheet_key=None,
                        googlesheet_uri=None,
                        sheet_name=None,
                        description=None,
                        header=True,
                        write_disposition='WRITE_TRUNCATE'):
        """ Loads Google Sheets data into a normal BigQuery Table

        Args:
            table_id (str): BigQuery table ID
            dataset_id (str): BigQuery dataset ID
            project_id (str): BigQuery project ID
            schema_path (str): Path to schema file, if not set then BQ will auto-detect
            googlesheet_uri (str): Google Sheet URI
            googlesheet_key (str): Google Sheet Key, an alternate option instead of URI
            sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0
            description (str): The descriptive text to describe the content of the table
            header (bool): Defaults to True
            write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY
        """

        temp_table_id=f"temp_gs__{table_id}"

        try:
            self.create_gs_table(
                table_id=temp_table_id,
                dataset_id=dataset_id,
                project_id=project_id,
                schema_path=schema_path,
                googlesheet_key=googlesheet_key,
                googlesheet_uri=googlesheet_uri,
                sheet_name=sheet_name,
                header=header
            )
            self.create_table(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                write_disposition=write_disposition,
                description=description,
                sql=f"SELECT * FROM `{project_id}.{dataset_id}.{temp_table_id}`"
            )
        except Exception as error:
            logging.error(error)
            raise error
        finally:
            if self.table_exists(dataset_id=dataset_id, table_id=temp_table_id):
                self.delete_table(dataset_id=dataset_id, table_id=temp_table_id)


    def load_json_file(self, file, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads JSON file to BigQuery table.

        Parameters:
            file (string): Path to JSON file.
            project_id (string): BigQuery Project ID.
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        if self.table_exists(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
        ):
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            job_config = bigquery.LoadJobConfig(schema=schema)
            job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            job_config.write_disposition = set_write_disposition(write_disposition)
            with open(file, mode='rb') as data:
                job = self.client.load_table_from_file(
                    file_obj=data,
                    destination=table_ref,
                    location='US',
                    job_config=job_config
                )
                job.result()
        else:
            raise Exception("Please initiate %s:%s.%s or pass the schema file",project_id ,dataset_id, table_id)

    def load_json_data(self, json, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads JSON data to BigQuery table.

        Parameters:
            json (string): JSON data.
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        if self.table_exists(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
        ):
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            job_config = bigquery.LoadJobConfig(schema=schema)
            job_config.write_disposition = set_write_disposition(write_disposition)
            job = self.client.load_table_from_json(
                json_rows=json,
                destination=table_ref,
                location='US',
                job_config=job_config
            )
            job.result()
        else:
            raise Exception("Please initiate %s:%s.%s or pass the schema file",project_id,dataset_id, table_id)

    def load_gcs(self, dataset_id, table_id, gcs_path, gcs_bucket=gcs_default_bucket(),project_id=bq_default_project(), location='US', schema_path='', header=True, write_disposition='WRITE_TRUNCATE'):
        """ Loads Google Cloud Storage CSV file into a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            gcs_bucket (string): Google Cloud Storage Bucket.
            gcs_path (string): Google Cloud Storage Path.
            location (string): Defaults to 'US'.
            schema NBD
            header (bool): Defaults to True.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """

        if isinstance(gcs_path, list):
            uri = ["gs://{}/{}".format(gcs_bucket, p) for p in gcs_path]
            uri_out = uri[0]
        else:
            uri = "gs://{}/{}".format(gcs_bucket, gcs_path)
            uri_out = uri

        job_config = bigquery.LoadJobConfig()

        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            job_config.schema = read_table_schema_from_file(schema_path)
            if header:
                job_config.skip_leading_rows = 1
        else:
            job_config.schema = None
            job_config.autodetect = True

        job_config.source_format = bigquery.SourceFormat.CSV
        job_config.write_disposition = write_disposition

        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )

        load_job = self.client.load_table_from_uri(
            uri,
            table_ref,
            job_config=job_config,
            location=location
        )

        load_job.result()  # Waits for job to complete.

        logging.info(
            'Loaded %s to %s:%s.%s',
            uri,
            project_id,
            dataset_id,
            table_id
        )

    def load_from_db(
        self,
        dataset_id,
        table_id,
        project_id=bq_default_project(),
        schema_path="",
        write_disposition="WRITE_TRUNCATE",
        source=None,
        source_url=None,
        sql=None,
        file=None,
        **kwargs
    ):
        """ Load data from another db into a BigQuery table
        Note:
            Sqlalchemy's built-in drivers should be supported
            Additional drivers can be added using requirements.txt
            Postgres(pg8000) and MySQL(pymysql) are included in tests

        Args:
            dataset_id (str): BigQuery dataset ID.
            table_id (str): BigQuery table ID.
            project_id (str): BigQuery project ID.
            schema_path (str): Path to schema file, if not set then BQ will auto-detect when creating a new table.
            write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
            source (str): Defines connection using env vars with this prefix, e.g. "ERP", "MAGENTO", etc.
            source_url (str): Defines connection using SQLAlchemy database URL format, e.g.
            sql (str): SQL query (or table name).
            file (str): path to the SQL file relative to PROJECT_ROOT env var.
            **kwargs: an be applied to pass parameters into SQL file
        """
        db = DBExecutor(name=source, url=source_url)
        df = db.execute_query(sql=sql, file=file, **kwargs)
        db.clean_up()

        # By default, convert_dtypes will attempt to convert a Series
        # (or each Series in a DataFrame) to dtypes that support pd.NA
        # This prevents, for example, int64 being force to float in columns that contain nulls
        df = df.convert_dtypes()

        gcs = GCSExecutor()
        gcs_path = f"temp_load_from_db/{dataset_id}/{table_id}.csv"
        gcs.df_to_gcs(
            gcs_path=gcs_path,
            df=df
        )

        self.load_gcs(
            project_id=project_id,
            dataset_id=dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            gcs_path=gcs_path,
            header=True,
            write_disposition=write_disposition
        )

        gcs.delete_directory(gcs_directory="temp_load_from_db")


    def insert_rows_json(self, dataset_id, table_id, rows,project_id=bq_default_project()):
        """ Insert rows into a table via the streaming API.
        Requires the table to exists in BigQuery.

        Arguments:
            - project_id (string): BigQuery project ID
            - dataset_id (string): BigQuery dataset ID
            - table_id (string): BigQuery table ID
            - rows (one of: list of tuples, list of dictonaries): Row data to be inserted.
            If a list of tuples is given, each tuple should contain data for each schema
            field on the current table and in the same order as the schema fields.
            If a list of dictionaries is given, the keys must include all required fields
            in the schema. Keys which do not correspond to a field in the schema are ignored.
        """
        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )
        try:
            error_response = self.client.insert_rows_json(
                table_ref,
                rows,
                ignore_unknown_values=True
            )
            if error_response == []:
                logging.info(
                    'Loaded %s row(s) into %s:%s.%s',
                    len(rows),
                    project_id,
                    dataset_id,
                    table_id
                )
            else:
                if hasattr(error_response, 'errors'):
                    for error in error_response.errors:
                        logging.warning(' - ' + error)
                        logging.warning(' - ' + rows)
                else:
                    logging.warning(json.dumps(error_response))

        except Exception:
            logging.exception('BigQuery Exception', exc_info=True)

    def extract_table_to_gcs(self, dataset_id, table_id, gcs_path,project_id=bq_default_project(), gcs_bucket=gcs_default_bucket(), location='US', shard=False):
        """ Extract BigQuery table into Google Cloud Storage.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            gcs_bucket (string): Google Cloud Storage Bucket.
            gcs_path (string): Google Cloud Storage Path.
            location (string): Defaults to 'US'.
            shard (bool)
        """

        uri = "gs://{}/{}".format(gcs_bucket, gcs_path)
        if shard:
            uri = re.sub('\\.', '*.', uri)

        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )

        job = self.client.extract_table(
            table_ref,
            uri,
            location=location
        )

        job.result()  # Waits for job to complete.

        logging.info(
            'Table %s:%s.%s loaded to %s',
            project_id,
            dataset_id,
            table_id,
            uri
        )

    def truncate_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Delete all records from table but preserve structure e.g. schema, partitioning, clustering, description, labels, etc.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
        """
        query=f"DELETE FROM `{project_id}.{dataset_id}.{table_id}` WHERE TRUE"
        self.execute_sql(sql=query, project_id=project_id)
        logging.info(
            'Table %s:%s.%s has been truncated',
            project_id,
            dataset_id,
            table_id
        )

    def copy_table_structure(
        self,
        source_table_id,
        dest_table_id,
        source_dataset_id=bq_default_dataset(),
        dest_dataset_id=bq_default_dataset(),
        source_project_id=bq_default_project(),
        dest_project_id=bq_default_project()
    ):


        if not self.dataset_exists(dest_dataset_id):
            self.create_dataset(dest_dataset_id)

        if self.table_exists(table_id=dest_table_id, dataset_id=dest_dataset_id):
            self.delete_table(table_id=dest_table_id, dataset_id=dest_dataset_id)

        if self.table_exists(table_id=source_table_id, dataset_id=source_dataset_id, project_id=source_project_id):

            attributes = self.get_table_attributes(
                table_id=source_table_id,
                dataset_id=source_dataset_id,
                project_id=source_project_id
            )
            table_def = "{}.{}.{}".format(dest_project_id, dest_dataset_id, dest_table_id)
            table = bigquery.Table(table_def)

            for attribute_key, attribute_value in attributes.items():
                if not attribute_value is None:
                    setattr(table, attribute_key, attribute_value)

            self.client.create_table(table) # Make an API request.
            logging.info(
                "Created table {}.{}.{}".format(table.project, table.dataset_id, table.table_id)
            )



    def copy_table(self, source_table_id, dest_table_id,
                   source_dataset_id=bq_default_dataset(), dest_dataset_id=bq_default_dataset(),
                   source_project_id=bq_default_project(),dest_project_id = bq_default_project(), write_disposition='WRITE_TRUNCATE'):
        """ Copy a BigQuery table.

        Parameters:
            source_project_id (string): Source BigQuery Project ID.
            source_dataset_id (string): Source BigQuery Dataset ID.
            source_table_id (string): Source BigQuery table ID.
            dest_project_id (string): Source BigQuery Project ID.
            dest_dataset_id (string): Destination BigQuery Dataset ID.
            dest_table_id (string): Destination BigQuery table ID.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.

        Limitations:
            Destination project is the default, i.e. specified in `local.env`
            Destination dataset must reside in the same location as source (US, EU, etc.)

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the source
        """
        source_table_ref = self.get_table_ref(source_dataset_id, source_table_id,project_id=source_project_id)
        dest_table_ref = self.get_table_ref(dest_dataset_id, dest_table_id,project_id=dest_project_id)

        if not self.dataset_exists(dest_dataset_id):
            self.create_dataset(dest_dataset_id)

        job_config = bigquery.CopyJobConfig()
        job_config.write_disposition = set_write_disposition(write_disposition)
        job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED
        job = self.client.copy_table(
            source_table_ref,
            dest_table_ref,
            project=source_project_id,
            job_config=job_config
        )
        job.result()
        logging.info(
            'Table %s:%s.%s copied to %s:%s.%s',
            source_project_id,
            source_dataset_id,
            source_table_id,
            dest_project_id,
            dest_dataset_id,
            dest_table_id
        )

    def list_datasets(self, project_id=bq_default_project(), include_all=False):
        """ Lists datasets in project

        Parameters:
            project_id (string): BigQuery project ID
            include_all (boolean): True if results include hidden datasets. Defaults to False.

        Returns:
            A sorted list of dataset ids

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the project
        """
        try:
            datasets = list(self.client.list_datasets(project=project_id, include_all=include_all))
            dataset_ids = [dataset.dataset_id for dataset in datasets]
            return dataset_ids
        except NotFound as error:
            logging.error(error)

    def list_tables(self, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Lists tables in dataset

        Parameters:
            project_id (string): BigQuery project ID
            dataset_id (string): BigQuery dataset ID

        Returns:
            A sorted list of table ids

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the project
        """
        try:
            dataset_ref = self.get_dataset_ref(dataset_id, project_id = project_id)
            tables = self.client.list_tables(dataset_ref)
            table_ids = [table.table_id for table in tables]
            return table_ids
        except NotFound as error:
            logging.error(error)

    def count_rows(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count rows in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.

        Returns:
            Number of rows in a BigQuery table.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)
        return table.num_rows

    def count_columns(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count columns in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.

        Returns:
            Number of columns in a BigQuery table.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)
        return len(table.schema)

    def count_duplicates(self, table_id, primary_key: list, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count duplicate rows in primary key

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.
            primary_key: a list of one or more column names, e.g. ['col1', 'col2']

        Returns:
            non-negative integer
        """
        if not primary_key:
            logging.warning("No primary key supplied, skipping duplicates count.")
            return 0
        data = self.execute_sql(
            f"""
            SELECT
                COALESCE(SUM(dup_count), 0) AS dup_total
            FROM (
                SELECT
                    {', '.join(primary_key)},
                    (COUNT(*) - 1) AS dup_count
                FROM
                    `{dataset_id}.{table_id}`
                GROUP BY
                    {', '.join(primary_key)}
            )
            """,
            project_id = project_id
        )
        return data['dup_total'].values[0]

    def assert_unique(self, table_id, primary_key: list, dataset_id=bq_default_dataset(),project_id=bq_default_project(), ignore_error=False, **kwargs):
        """ Assert uniqueness of primary key in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.
            primary_key: a list of one or more column names, e.g. ['col1', 'col2']
            ignore_error: boolean flag to prevent error being raised, useful for debugging

        Returns:
            Nothing if there are no duplicate rows in primary key
            Raise and log AssertionError if there are duplicate rows and ignore_error=False (default)
            Log a warning if there are dupicate rows and ignore_error=True (debugging)
        """
        if self.count_duplicates(table_id, primary_key, dataset_id, project_id=project_id) != 0:
            msg = "Table %s:%s.%s is not unique on %s" % (
                project_id,
                dataset_id,
                table_id,
                ', '.join(primary_key)
            )
            if ignore_error:
                logging.warning(msg)
            else:
                raise AssertionError(msg)

    def assert_acceptance(self, sql, cte, output_table_name='expected_output', **kwargs):
        """
        Compares the target cte defined in the mock file with the output of the sql.

        Notes:
        - order of rows and columns doesn't matter
        - only handles 1 layer of nesting
        """
        try:
            sql_extract_output_table = "WITH {} ( SELECT * FROM `{}` )".format(cte, output_table_name)

            df_result = self.execute_sql(sql_extract_output_table)
        except:
            raise AssertionError("Bad SQL or CTE")
        composite_sql = "WITH {} ( {} )".format(cte, sql)
        df = self.execute_sql(composite_sql)

        # flatten and sort - ignore column and row order
        df = flatten_df_columns(df)
        df = df[df.columns.sort_values()]
        df = df.sort_values(list(df.columns)).reset_index(drop=True)
        df_result = flatten_df_columns(df_result)
        df_result = df_result[df_result.columns.sort_values()]
        df_result = df_result.sort_values(list(df_result.columns)).reset_index(drop=True)
        assert_frame_equal(df, df_result)

<filename>pygyver/etl/dw.py
""" Module containing BigQueryExecutor for Python """
import os
import re
import logging
import time
import json
import pandas as pd
from pandas._testing import assert_frame_equal
from google.cloud import bigquery
from google.cloud.exceptions import NotFound
from google.oauth2 import service_account
from google.api_core import exceptions
from pygyver.etl.lib import bq_token_file_valid
from pygyver.etl.lib import bq_token_file_path
from pygyver.etl.lib import bq_default_project
from pygyver.etl.lib import bq_default_dataset
from pygyver.etl.lib import gcs_default_bucket
from pygyver.etl.lib import read_table_schema_from_file
from pygyver.etl.lib import bq_start_date
from pygyver.etl.lib import bq_end_date
from pygyver.etl.lib import set_write_disposition, set_priority
from pygyver.etl.toolkit import (
    date_lister,
    flatten_df_columns,
    validate_date
)
from pygyver.etl.gs import load_gs_to_dataframe
from pygyver.etl.storage import GCSExecutor
from pygyver.etl.db import DBExecutor


class BigQueryExecutorError(Exception):
    pass


def print_kwargs_params(func):
    def inner(*args, **kwargs):
        logging.debug("Keyword args applied to the template:")
        for key, value in kwargs.items():
            if key in forbiden_kwargs():
                raise KeyError("{} is a forbidden keyword argument.".format(key))
        for key, value in kwargs.items():
            logging.debug("%s = %s" % (key, value))
        return func(*args, **kwargs)
    return inner


def forbiden_kwargs():
    return ['partition_date']


class SafeDict(dict):
    def __missing__(self, key):
        return '{' + key + '}'

@print_kwargs_params
def read_sql(file, *args, **kwargs):
    ''' Read SQL file and apply arguments/keyword arguments.

    If dataset_prefix is a kwarg then this function is being used in a dry-run test.
    In this case, the dataset_prefix is added to the dataset name and project names
    are removed from the SQL.

    Args:
        file (string): path to the SQL file from PROJECT_ROOT environment variable.
        *kwargs can be passed if some parameters are to be passed.

    Returns:
        a SQL formated with **kwargs if applicable.

    Example:
        With SQL as:
            "select .. {param2} .. {param1} .. {paramN}"
        *kwargs as:
            param1=value1
            param2=value2
            paranN=valueN
        The functions returns:
            "select .. value2 .. value1 .. valueN"
    '''
    path_to_file = os.path.join(os.getenv("PROJECT_ROOT"), file)
    file = open(path_to_file, 'r')
    sql = file.read()
    file.close()
    if kwargs.get('dataset_prefix', None) is not None:
        sql_split = sql.split("`")
        for index, table in enumerate(sql_split):
            if index%2==1 and "." in table:
                split_table = table.split(".")
                split_table[-2] = kwargs.get('dataset_prefix', '') + split_table[-2]
                sql_split[index] = '.'.join(split_table[-2:])
        sql = '`'.join(sql_split)

    if len(kwargs) > 0:
        sql = sql.format_map(SafeDict(**kwargs))
    return sql


class BigQueryExecutor:
    """ BigQuery handler

    Parameters:
        project_id (sql_file): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.

    Required:
        GOOGLE_APPLICATION_CREDENTIALS (env variable).

    Attributes:
        client (Client object)
        credentials (Credentials object)
        project_id (string): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.

    Returns:
        a BigQueryExecutor object.
    """
    def __init__(self):
        """ Resets client and credentials.
        """
        self.client = None
        self.credentials = None
        self.auth()

    def auth(self):
        """ Sets BigQuery client.
        """
        bq_token_file_valid()
        self.credentials = service_account.Credentials.from_service_account_file(
            bq_token_file_path(),
            scopes=[
                "https://www.googleapis.com/auth/drive",
                "https://www.googleapis.com/auth/bigquery"
            ]
        )
        self.client = bigquery.Client(
            project = bq_default_project(),
            credentials=self.credentials
        )


    def get_dataset_ref(self,  dataset_id, project_id=bq_default_project()):
        """ Returns BigQuery DatasetReference object.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            bigquery.dataset.DatasetReference object
        """
        return bigquery.dataset.DatasetReference(
            project = project_id,
            dataset_id = dataset_id
        )

    def get_table_ref(self, dataset_id, table_id, project_id=bq_default_project()):
        """ Returns BigQuery Table reference object.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            Returns BigQuery Table reference object.
        """
        dataset_ref = self.get_dataset_ref(dataset_id, project_id = project_id)
        return dataset_ref.table(table_id)

    def dataset_exists(self, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Checks if a BigQuery dataset exists.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            True is the dataset exists, False otherwise
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        try:
            self.client.get_dataset(dataset_ref)
            return True
        except NotFound:
            return False

    def delete_dataset(self, dataset_id=bq_default_dataset(),project_id=bq_default_project(),delete_contents=False):
        """ Deletes a BigQuery dataset.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            delete_contents (bool): If True, delete all the tables in the dataset.
            If False and the dataset contains tables, the request will fail. Default is False.
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        try:
            self.client.delete_dataset(
                dataset_ref,
                delete_contents=delete_contents
            )
            logging.info(
                "Dataset %s:%s deleted",
                project_id,
                dataset_id
            )
            time.sleep(1)
        except exceptions.Conflict as error:
            logging.error(error)


    def create_dataset(self, dataset_id=bq_default_dataset(), project_id=bq_default_project(), **kwargs):
        """ Creates a BigQuery dataset if the dataset does not exists. Otherwise pass.

        Parameters:
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
        """
        dataset_ref = self.get_dataset_ref(dataset_id,project_id=project_id)
        if self.dataset_exists(dataset_id,project_id=project_id):
            logging.info(
                "Dataset %s already exists in project %s",
                dataset_id,
                project_id
            )
        else:
            try:
                self.client.create_dataset(dataset_ref)
                logging.info(
                    "Created dataset %s in in project %s",
                    dataset_id,
                    project_id
                )
            except exceptions.Conflict as error:
                logging.error(error)

    def table_exists(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Checks if a BigQuery table exists.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.


        Returns:
            True is the table exists, False otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        try:
            self.client.get_table(table_ref)
            return True
        except NotFound:
            return False

    def delete_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Delete a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        """
        try:
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            self.client.delete_table(table_ref)
            logging.info(
                'Table %s:%s.%s deleted.',
                project_id,
                dataset_id,
                table_id
            )
            time.sleep(1)
        except NotFound as error:
            logging.error(error)

    def initiate_table(self,
                       table_id,
                       schema_path,
                       dataset_id=bq_default_dataset(),
                       project_id=bq_default_project(),
                       partition=False,
                       partition_field='_PARTITIONTIME',
                       clustering=None):
        """ Initiate a BigQuery table. If the table already exists, compares the schema_path and apply a patch if there is a schema change.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string: Path to the BigQuery table schema from the PROJECT_ROOT environement variable.
            partition (bool): Specify whether the BigQuery table is partioned. Default to False.
            partition_field (string): Specify partition_field if table partioned. Defaults to "_PARTITIONTIME".
            clustering (list): List of clustering fields. Defaults to None.
        """
        if self.table_exists(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id):
            logging.info("Table %s.%s already exists in project %s",
                         dataset_id,
                         table_id,
                         project_id)
            self.apply_patch(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path
            )
        else:
            schema = read_table_schema_from_file(schema_path)
            table = bigquery.Table(
                self.get_table_ref(dataset_id, table_id, project_id = project_id),
                schema=schema
                )
            if partition:
                table.time_partitioning = bigquery.table.TimePartitioning(type_='DAY')
                if partition_field and partition_field != '_PARTITIONTIME':
                    if isinstance(partition_field, str):
                        table.time_partitioning = bigquery.table.TimePartitioning(
                            field=partition_field
                        )
                    else:
                        raise ValueError("partition_field should be a string")
            table.clustering_fields = clustering
            try:
                table = self.client.create_table(table)
                logging.info(
                    'Created table %s.%s in in project %s',
                    dataset_id,
                    table_id,
                    project_id
                )
            except exceptions.Conflict as error:
                logging.error(error)

    def create_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), sql=None, file=None,
                     write_disposition='WRITE_TRUNCATE', use_legacy_sql=False,
                     location='US', schema_path='',
                     partition=False,
                     partition_field='_PARTITIONTIME',
                     clustering=None,
                     priority='INTERACTIVE',
                     description=None,
                     **kwargs):
        """ create a bigquery table from a sql query """

        if sql is None and file is None:
            raise BigQueryExecutorError("Either SQL or file containing the SQL must be provided")

        if sql is None:
            sql = read_sql(file, **kwargs)

        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                partition=partition,
                partition_field=partition_field,
                dataset_id=dataset_id,
                project_id=project_id
            )
            if write_disposition == "WRITE_TRUNCATE":
                self.truncate_table(dataset_id=dataset_id, table_id=table_id,project_id=project_id)
                write_disposition = "WRITE_EMPTY"
        else:
            pass

        job_config = bigquery.QueryJobConfig()
        job_config.destination = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        job_config.write_disposition = set_write_disposition(write_disposition)
        job_config.use_legacy_sql = use_legacy_sql
        job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED
        job_config.priority = set_priority(priority)
        if partition:
            if not partition_field or partition_field == '_PARTITIONTIME':
                job_config.time_partitioning = bigquery.table.TimePartitioning(
                    type_='DAY'
                )
                job_config.clustering_fields = clustering
            elif isinstance(partition_field, str):
                job_config.time_partitioning = bigquery.table.TimePartitioning(
                    field=partition_field
                )
            else:
                raise ValueError("partition_field should be a string")

        query_job = self.client.query(
            sql,
            location=location,
            job_config=job_config
        )
        if priority == 'INTERACTIVE':
            query_job.result()
            logging.info(
            'Query results loaded to table %s:%s.%s',
                project_id,
                dataset_id,
                table_id
            )
        if description:
            self.update_table_description(table_id=table_id, description=description,
                                          project_id=project_id, dataset_id=dataset_id)

        return query_job


    def delete_backfill_partitions(self,
                                   backfill_dataset_id,
                                   table_id,
                                   partition_field,
                                   from_partition,
                                   to_partition
                                   ):
        """
        Delete partitions from table in backfill dataset between from_partition and to_partition

        Args:
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): table in backfill dataset that the partitions will be deleted from
            partition_field (string): the time-unit column that the table is partitioned by
            from_partition (date): date of the first partition to be deleted
            to_partition (date): date of the latest partition to be deleted

        Returns:
            None
        """
        client = bigquery.Client()

        sql = f"""
               DELETE
                 `{backfill_dataset_id}.{table_id}`
               WHERE
                 CAST({partition_field} AS DATE) BETWEEN
                   CAST('{from_partition}' AS DATE) AND CAST('{to_partition}' AS DATE)
               """

        query_job = client.query(
            sql
        )

        query_job.result() # wait for job to complete


    def run_full_backfill(self,
                          dataset_id,
                          backfill_dataset_id,
                          table_id,
                          schema_path,
                          sql,
                          partition_field
                          ):
        """
        Run a full backfill for a time-unit column partitioned table

        Args:
            dataset_id (string): dataset of the table to be backfilled
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by

        Returns:
            None
        """
        logging.info("Starting full backfill")

        logging.info("Creating backfill dataset if it doesn't exist")
        self.create_dataset(
            dataset_id=backfill_dataset_id
        )

        logging.info("Creating table in backfill dataset with all partitions")
        self.create_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            sql=sql,
            partition=True,
            partition_field=partition_field,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Replacing prod table with backfill table")
        self.copy_table(
            source_dataset_id=backfill_dataset_id,
            source_table_id=table_id,
            dest_dataset_id=dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting backfill table")
        self.delete_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id
        )


    def run_partial_backfill(self,
                             dataset_id,
                             backfill_dataset_id,
                             table_id,
                             schema_path,
                             sql,
                             partition_field,
                             from_partition,
                             to_partition
                             ):
        """
        Run a partial backfill for a time-unit column partitioned table between the dates
        from_partition and to_partition

        Args:
            dataset_id (string): dataset of the table to be backfilled
            backfill_dataset_id (string): dataset where the backfill tables are created
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by
            from_partition (date): date of the first partition to be backfilled
            to_partition (date): date of the latest partition to be backfilled

        Returns:
            None
        """
        logging.info("Starting partial backfill")

        logging.info("Creating backfill dataset if it doesn't exist")
        self.create_dataset(
            dataset_id=backfill_dataset_id
        )

        logging.info("Copying table to backfill dataset")
        self.copy_table(
            source_dataset_id=dataset_id,
            source_table_id=table_id,
            dest_dataset_id=backfill_dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting partitions to be backfilled")
        self.delete_backfill_partitions(
            backfill_dataset_id,
            table_id,
            partition_field,
            from_partition,
            to_partition
        )

        logging.info("Appending new backfill partitions")
        self.create_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            sql=sql,
            partition=True,
            partition_field=partition_field,
            write_disposition="WRITE_APPEND"
        )

        logging.info("Replacing prod table with backfill table")
        self.copy_table(
            source_dataset_id=backfill_dataset_id,
            source_table_id=table_id,
            dest_dataset_id=dataset_id,
            dest_table_id=table_id,
            write_disposition="WRITE_TRUNCATE"
        )

        logging.info("Deleting backfill table")
        self.delete_table(
            dataset_id=backfill_dataset_id,
            table_id=table_id
        )


    def run_backfill(self,
                     backfill,
                     dataset_id,
                     table_id,
                     schema_path,
                     sql,
                     partition_field,
                     from_partition,
                     to_partition
                     ):
        """
        Run a full or partial backfill for a time-unit column partitioned table

        Args:
            backfill (string): specifies full or partial backfill (if running a backfill)
            dataset_id (string): dataset of the table to be backfilled
            table_id (string): name of the table to be backfilled
            schema_path (string): file path to the schema of the table to be backfilled
            sql (string): query that's used in the pipeline to create this table
                Note: it has to contain the clause:
                WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)
            partition_field (string): the time-unit column that the table is partitioned by

        Returns:
            None
        """
        logging.info("Starting backfill")

        backfill_dataset = "backfill"

        if backfill not in ("full", "partial"):
            raise BigQueryExecutorError("BACKFILL env var needs to be 'full' or 'partial'")

        if backfill == "full":
            self.run_full_backfill(
                dataset_id=dataset_id,
                backfill_dataset_id=backfill_dataset,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field
            )

        elif backfill == "partial":
            self.run_partial_backfill(
                dataset_id=dataset_id,
                backfill_dataset_id=backfill_dataset,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field,
                from_partition=from_partition,
                to_partition=to_partition
            )

    def create_partitions_legacy(self,
                                 sql,
                                 file,
                                 partition_dates,
                                 table_id,
                                 project_id,
                                 dataset_id,
                                 bq_start_date,
                                 bq_end_date,
                                 write_disposition,
                                 use_legacy_sql,
                                 partition_field,
                                 clustering,
                                 priority,
                                 **kwargs
                                 ):
        if sql is None:
            sql = read_sql(file, **kwargs)

        if partition_dates is None:
            existing_dates = self.get_existing_partition_dates(
                table_id=table_id,
                project_id=project_id,
                dataset_id=dataset_id
            )
            dates = self.get_partition_dates(
                start_date=bq_start_date,
                end_date=bq_end_date,
                existing_dates=existing_dates
            )
        else:
            self.validate_partition_dates(
                partition_dates=partition_dates
            )
            dates = partition_dates

        jobs = []
        for date in dates:
            partition_name = self.set_partition_name(table=table_id, date=date)
            logging.info("Updating partition: %s", partition_name)
            job = self.create_table(
                sql=self.apply_partition_filter(
                    sql=sql,
                    date=date
                ),
                dataset_id=dataset_id,
                project_id=project_id,
                table_id=partition_name,
                write_disposition=write_disposition,
                use_legacy_sql=use_legacy_sql,
                partition=True,
                partition_field=partition_field,
                clustering=clustering,
                priority=priority
            )
            jobs.append(job)

        if priority == 'BATCH':
            for job in jobs:
                job.result()


    def create_partition_table(self,
                               table_id,
                               dataset_id=bq_default_dataset(),
                               project_id=bq_default_project(),
                               bq_start_date=bq_start_date(),
                               bq_end_date=bq_end_date(),
                               sql=None,
                               file=None,
                               use_legacy_sql=False,
                               write_disposition='WRITE_TRUNCATE',
                               partition_dates=None,
                               pipeline_partition_date=None,
                               partition_field='_PARTITIONTIME',
                               clustering=None,
                               priority="INTERACTIVE",
                               schema_path="",
                               description=None,
                               **kwargs
                              ):
        """
        New and legacy ways of creating partition tables:

        Legacy:
        Partition to be generated are either passed through partition_dates or automatically generated using existing partitions.
        To filter on a specific partition, the filter DATE(_PARTITIONTIME) = {partition_date} can be used in your sql query.

        New:
        pipeline_partition_date is required for a normal pipeline run and is a single date.
        To run a backfill, the BACKFILL env var needs to be set to 'full' or 'partial' and an optional date range 
        can be set for a partial backfill by setting the FROM_PARTITION and TO_PARTITION env vars.
        """
        if sql is None and file is None:
            raise BigQueryExecutorError("Either SQL or file containing the SQL must be provided")

        if schema_path != "":
            self.initiate_table(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                partition=True,
                partition_field=partition_field,
                clustering=clustering
            )

        if not self.table_exists(dataset_id=dataset_id, table_id=table_id,project_id=project_id):
            raise BigQueryExecutorError("To create a partition, please initiate the table first using initiate_table.")

        if "BACKFILL" not in os.environ:
            if pipeline_partition_date is None: # legacy approach
                self.create_partitions_legacy(
                    sql=sql,
                    file=file,
                    partition_dates=partition_dates,
                    table_id=table_id,
                    project_id=project_id,
                    dataset_id=dataset_id,
                    bq_start_date=bq_start_date,
                    bq_end_date=bq_end_date,
                    write_disposition=write_disposition,
                    use_legacy_sql=use_legacy_sql,
                    partition_field=partition_field,
                    clustering=clustering,
                    priority=priority,
                    **kwargs
                )

            else: # new approach, no backfill
                from_partition = pipeline_partition_date
                to_partition = pipeline_partition_date
                sql = read_sql(
                    file=file,
                    from_partition=from_partition,
                    to_partition=to_partition,
                    **kwargs
                )

                # change date format from YYYY-MM-DD to YYYYMMDD for partition_id 
                partition_id = pipeline_partition_date.replace("-", "")
                partition_name = self.set_partition_name(
                    table=table_id,
                    date=partition_id
                )

                logging.info("Updating partition: %s", partition_name)
                self.create_table(
                    project_id=project_id,
                    dataset_id=dataset_id,
                    table_id=partition_name,
                    sql=sql,
                    partition=True,
                    partition_field=partition_field,
                    write_disposition=write_disposition,
                    use_legacy_sql=use_legacy_sql,
                    clustering=clustering,
                    priority=priority
                )

        else: # new approach, run backfill
            from_partition = os.environ.get("FROM_PARTITION", bq_start_date)
            to_partition = os.environ.get("TO_PARTITION", bq_end_date)
            sql = read_sql(
                file=file,
                from_partition=from_partition,
                to_partition=to_partition,
                **kwargs
            )
            self.run_backfill(
                backfill=os.environ["BACKFILL"],
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                sql=sql,
                partition_field=partition_field,
                from_partition=from_partition,
                to_partition=to_partition
            )

        if description:
            self.update_table_description(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                description=description
            )


    def apply_partition_filter(self, sql, date):
        """ Apply partition_date to the SQL query.

        Parameters:
            sql (string): the SQL query
            partition_date (string): the partiton date

        Returns:
            Formatted SQL query.
        """
        return sql.format(
            partition_date=date
        )

    def validate_partition_dates(self, partition_dates):
        """ Validates the partition_dates parameter. Checks whether it is a list and the elements are in the right format.

        Parameters:
            partition_dates (list of string): list of partition_dates
            partition_date (string): the partiton date

        Raises:
            BigQueryExecutorError if the partition_dates parameter in not a list.
            ValueError if incorrect date format.
        """
        if not isinstance(partition_dates, list):
            raise BigQueryExecutorError("Partition dates need to be a list of date eg ['YYYYmmdd']")
        else:
            for date in partition_dates:
                validate_date(date=date, format='%Y%m%d')

    def set_partition_name(self, table, date):
        """ Validates the date and sets the partition_name associated.

        Parameters:
            table (string): table name
            date (string): date with format '%Y%m%d'

        Returns:
            Partition name

        Raises:
            BigQueryExecutorError if the partition_dates parameter in not a list.
            ValueError if incorrect date format.
        """
        validate_date(date=date, format='%Y%m%d')
        return table + "$" + date.replace("-", "")

    def get_partition_dates(self, start_date, end_date, existing_dates):
        """ Returns the partitions dates required based on start_date, end_date and existing_dates.

        Parameters:
            start_date (string): First partition required.
            end_date (string): Last partition required.
            existing_dates (list of string): List of existing partition.

        Returns:
            List of partition dates required.
        """
        partition_dates = []
        required_dates = date_lister(start_date=start_date, end_date=end_date)
        if existing_dates == []:
            for date in required_dates:
                partition_date = date.replace("-", "")
                partition_dates.append(partition_date)
        else:
            for date in required_dates:
                partition_date = date.replace("-", "")
                if partition_date not in existing_dates:
                    partition_dates.append(partition_date)
        return partition_dates

    def get_existing_partition_query(self, dataset_id, table_id,project_id = bq_default_project()):
        """ Gets existing partitions from BigQuery as dataframe.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            DataFrame of existing partitions
        """
        sql = """ SELECT
                    FORMAT_DATE('%Y%m%d', DATE(_PARTITIONTIME)) AS partition_id
                  FROM
                    `{dataset_id}.{table_id}`
                  GROUP BY
                    1 """.format(
                        dataset_id=dataset_id,
                        table_id=table_id
                        )
        return self.execute_sql(
            sql=sql,
            project_id=project_id
        )

    def get_existing_partition_dates(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Gets existing partitions.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            List of existing partition dates.
        """
        if not self.table_exists(dataset_id=dataset_id, table_id=table_id,project_id=project_id):
            existing_partition_dates = []
        else:
            res = self.get_existing_partition_query(dataset_id=dataset_id, table_id=table_id)
            # checks that res has number of rows > 0
            if res.shape[0] > 0:
                existing_partition_dates = res['partition_id'].to_list()
            # if not, no existing partitions
            else:
                existing_partition_dates = []
        return existing_partition_dates

    def get_table_schema(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table schema object

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.


        Returns:
            Table schema.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        table_schema = self.client.get_table(table_ref).schema
        return table_schema

    def get_table_partitioning_type(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table partitioning_type

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            partitioning_type is 'DAY' if the table is partitioned, None otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        table_properties = self.client.get_table(table_ref)._properties
        partitioning_type = table_properties.get('timePartitioning', {}).get('type')
        return partitioning_type

    def get_table_attributes(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        dict_of_attributes = {}
        attributes = ['clustering_fields', 'description', 'encryption_configuration', 'expires',
        'external_data_configuration', 'friendly_name', 'labels', 'range_partitioning',
        'require_partition_filter', 'schema', 'time_partitioning']
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        for attribute in attributes:
            my_attribute = getattr(self.client.get_table(table_ref), attribute)
            dict_of_attributes.update({attribute: my_attribute})
        return dict_of_attributes


    def get_table_clustering_fields(self, table_id, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Gets table clustering_fields

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.

        Returns:
            list of clustering fields is the table has cluster fields, None otherwise.
        """
        table_ref = self.get_table_ref(dataset_id, table_id, project_id=project_id)
        clustering_fields = self.client.get_table(table_ref).clustering_fields
        return clustering_fields

    def identify_new_fields(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Identify new fields in based on a schema file.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to the schema to compare to.

        Returns:
            List of new fields.
        """
        list_field = []
        schema_a = self.get_table_schema(
            table_id=table_id,
            dataset_id=dataset_id,
            project_id=project_id
        )
        schema_b = read_table_schema_from_file(schema_path)
        field_list_a = [schema_field.name for schema_field in schema_a]
        for schema_field in schema_b:
            if schema_field.name not in field_list_a:
                list_field.append(schema_field)
        return list_field

    def append_field(self, table_id, field, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Append fields to a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            field (string): Schema field object.

        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        table = self.client.get_table(table_ref)  # API request

        original_schema = table.schema
        new_schema = original_schema[:]  # creates a copy of the schema
        new_schema.append(field)

        table.schema = new_schema
        table = self.client.update_table(table, ["schema"])  # API request
        assert len(table.schema) == len(original_schema) + 1 == len(new_schema)
        return 0

    def apply_patch(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Apply a patch to a BigQuery Table if required.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file to compare to.

        Returns:
            Lenght of new schema
        """
        logging.info("Attempting patch")
        logging.info("Checking for new fields...")
        new_fields = self.identify_new_fields(
            table_id=table_id,
            schema_path=schema_path,
            project_id=project_id,
            dataset_id=dataset_id
        )
        if new_fields != []:
            logging.info("New fields to be added:")
            logging.info(new_fields)
            for field in new_fields:
                self.append_field(
                    field=field,
                    table_id=table_id,
                    dataset_id=dataset_id,
                    project_id=project_id
                )
            logging.info("Done!")
        else:
            logging.info("No field to be added")

        logging.info("Checking for schema update...")
        self.update_schema(
            table_id=table_id,
            schema_path=schema_path,
            dataset_id=dataset_id,
            project_id=project_id
        )
        return len(
            self.get_table_schema(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
            )
            )

    def update_schema(self, table_id, schema_path, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Perform a schema update. Used to update descriptions.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file to compare to.

        Raises:
            BadRequest if the update fails.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)  # API request
        new_schema = read_table_schema_from_file(schema_path)
        if table.schema == new_schema:
            logging.info("No changes needed")
        else:
            assert len(table.schema) == len(new_schema)
            table.schema = new_schema
            try:
                table = self.client.update_table(table, ["schema"])  # API request
                return 0
            except exceptions.BadRequest as error:
                raise error

    def update_table_description(
            self,
            table_id,
            description,
            project_id=bq_default_project(),
            dataset_id=bq_default_dataset()
            ):
        """ Performs a table update to fill in description for the table.

        Parameters:
            table_id (string): BigQuery table ID.
            description (string): The descriptive text to describe the content of the table.
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.

        Raises:
            BadRequest if the update fails.
        """

        table_ref = self.get_table_ref(dataset_id=dataset_id, table_id=table_id, project_id=project_id)
        table = self.client.get_table(table_ref)  # API request

        if table.description == description:
            logging.info("No changes to table description required")
        else:
            try:
                table.description = description
                self.client.update_table(table, ["description"])  # API request
            except exceptions.BadRequest as error:
                raise error

    def execute_sql(self, sql, project_id=bq_default_project(), dialect='standard'):
        """ Executes a SQL query and loads it as a DataFrame.

        Parameters:
            sql (string): SQL Query.
            project_id (string): BigQuery Project ID.
            dialect (string): BigQuery dialect. Defaults to standard.

        Returns:
            Query result as a DataFrame.
        """
        data = pd.read_gbq(
            sql,
            project_id=project_id,
            credentials=self.credentials,
            dialect=dialect
        )

        return data

    def execute_dml(self, sql=None, file=None, location='US', project=bq_default_project()):

        job_config = bigquery.QueryJobConfig()

        try:

            sql = sql if sql else read_sql(file)
        
            job = self.client.query(sql, location=location, job_config=job_config)
            
            results = job.result()
            
            logging.info("%s - %s - %s rows effected (%s elapsed, %s B processed)",
                job.statement_type,
                job.state,
                job.num_dml_affected_rows,
                job.ended - job.started,
                job.total_bytes_processed
            )
        
        except Exception as e:
            logging.error(e)
            raise e

        return

    def execute_file(self, file, project_id=bq_default_project(),
                     dialect='standard', *args, **kwargs):
        """ Executes a SQL file and loads it as a DataFrame.

        Parameters:
            file (string): Path to SQL file.
            project_id (string): BigQuery Project ID.
            dialect (string): BigQuery dialect. Defaults to standard.

        **kwargs can be passed if the SQL file contains arguments formatted with {}.

        Forbidden kwargs:
            partition_date

        Returns:
            Query result as a DataFrame.
        """
        sql = read_sql(file, *args, **kwargs)
        data = self.execute_sql(
            sql=sql,
            project_id=project_id,
            dialect=dialect
        )
        return data

    def load_dataframe(self, df, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads DataFrame to BigQuery table.

        Parameters:
            df (pd.DataFrame): Pandas DataFrame
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id,
                schema_path=schema_path
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        data = df.rename(columns=lambda cname: cname.replace('.', '_'))
        table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
        job_config = bigquery.LoadJobConfig(schema=schema)
        job_config.write_disposition = set_write_disposition(write_disposition)
        job = self.client.load_table_from_dataframe(
            data,
            table_ref,
            job_config=job_config
        )
        job.result()


    def create_gs_table(self,
                        table_id,
                        dataset_id=bq_default_dataset(),
                        project_id=bq_default_project(),
                        schema_path='',
                        googlesheet_uri=None,
                        googlesheet_key=None,
                        sheet_name=None,
                        header=True):
        """ Creates BigQuery Table with live connection to Google Sheets

        Args:
            table_id (str): BigQuery table ID
            dataset_id (str): BigQuery dataset ID
            project_id (str): BigQuery project ID
            schema_path (str): Path to schema file, if not set then BQ will auto-detect
            googlesheet_uri (str): Google Sheet URI
            googlesheet_key (str): Google Sheet Key, an alternate option instead of URI
            sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0
            header (bool): Defaults to True
        """
        if googlesheet_uri is None:
            if googlesheet_key is None:
                raise BigQueryExecutorError("A googlesheet_uri or googlesheet_key must be provided")
            else:
                googlesheet_uri = f"https://docs.google.com/spreadsheets/d/{googlesheet_key}"

        external_config = bigquery.ExternalConfig("GOOGLE_SHEETS")
        external_config.source_uris = [googlesheet_uri]
        if sheet_name:
            external_config.options.range = (sheet_name)

        if header:
            external_config.options.skip_leading_rows = 1

        if schema_path != '':
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None
            external_config.autodetect = True

        gs_table = bigquery.Table(
            self.get_table_ref(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id
            ),
            schema=schema
        )
        gs_table.external_data_configuration=external_config

        try:
            self.client.delete_table(gs_table, not_found_ok=True)
            self.client.create_table(gs_table)
            logging.info(
                f"Created table {project_id}:{dataset_id}.{table_id} with live connection to {googlesheet_uri}"
            )
        except exceptions.Conflict as error:
            logging.error(error)
            raise error


    def load_google_sheet(self,
                        table_id,
                        dataset_id=bq_default_dataset(),
                        project_id=bq_default_project(),
                        schema_path='',
                        googlesheet_key=None,
                        googlesheet_uri=None,
                        sheet_name=None,
                        description=None,
                        header=True,
                        write_disposition='WRITE_TRUNCATE'):
        """ Loads Google Sheets data into a normal BigQuery Table

        Args:
            table_id (str): BigQuery table ID
            dataset_id (str): BigQuery dataset ID
            project_id (str): BigQuery project ID
            schema_path (str): Path to schema file, if not set then BQ will auto-detect
            googlesheet_uri (str): Google Sheet URI
            googlesheet_key (str): Google Sheet Key, an alternate option instead of URI
            sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0
            description (str): The descriptive text to describe the content of the table
            header (bool): Defaults to True
            write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY
        """

        temp_table_id=f"temp_gs__{table_id}"

        try:
            self.create_gs_table(
                table_id=temp_table_id,
                dataset_id=dataset_id,
                project_id=project_id,
                schema_path=schema_path,
                googlesheet_key=googlesheet_key,
                googlesheet_uri=googlesheet_uri,
                sheet_name=sheet_name,
                header=header
            )
            self.create_table(
                project_id=project_id,
                dataset_id=dataset_id,
                table_id=table_id,
                schema_path=schema_path,
                write_disposition=write_disposition,
                description=description,
                sql=f"SELECT * FROM `{project_id}.{dataset_id}.{temp_table_id}`"
            )
        except Exception as error:
            logging.error(error)
            raise error
        finally:
            if self.table_exists(dataset_id=dataset_id, table_id=temp_table_id):
                self.delete_table(dataset_id=dataset_id, table_id=temp_table_id)


    def load_json_file(self, file, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads JSON file to BigQuery table.

        Parameters:
            file (string): Path to JSON file.
            project_id (string): BigQuery Project ID.
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        if self.table_exists(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
        ):
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            job_config = bigquery.LoadJobConfig(schema=schema)
            job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            job_config.write_disposition = set_write_disposition(write_disposition)
            with open(file, mode='rb') as data:
                job = self.client.load_table_from_file(
                    file_obj=data,
                    destination=table_ref,
                    location='US',
                    job_config=job_config
                )
                job.result()
        else:
            raise Exception("Please initiate %s:%s.%s or pass the schema file",project_id ,dataset_id, table_id)

    def load_json_data(self, json, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project(), schema_path='', write_disposition="WRITE_TRUNCATE"):
        """ Loads JSON data to BigQuery table.

        Parameters:
            json (string): JSON data.
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            schema_path (string): Path to schema file.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """
        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            schema = read_table_schema_from_file(schema_path)
        else:
            schema = None

        if self.table_exists(
                table_id=table_id,
                dataset_id=dataset_id,
                project_id=project_id
        ):
            table_ref = self.get_table_ref(dataset_id, table_id,project_id=project_id)
            job_config = bigquery.LoadJobConfig(schema=schema)
            job_config.write_disposition = set_write_disposition(write_disposition)
            job = self.client.load_table_from_json(
                json_rows=json,
                destination=table_ref,
                location='US',
                job_config=job_config
            )
            job.result()
        else:
            raise Exception("Please initiate %s:%s.%s or pass the schema file",project_id,dataset_id, table_id)

    def load_gcs(self, dataset_id, table_id, gcs_path, gcs_bucket=gcs_default_bucket(),project_id=bq_default_project(), location='US', schema_path='', header=True, write_disposition='WRITE_TRUNCATE'):
        """ Loads Google Cloud Storage CSV file into a BigQuery table.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            gcs_bucket (string): Google Cloud Storage Bucket.
            gcs_path (string): Google Cloud Storage Path.
            location (string): Defaults to 'US'.
            schema NBD
            header (bool): Defaults to True.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
        """

        if isinstance(gcs_path, list):
            uri = ["gs://{}/{}".format(gcs_bucket, p) for p in gcs_path]
            uri_out = uri[0]
        else:
            uri = "gs://{}/{}".format(gcs_bucket, gcs_path)
            uri_out = uri

        job_config = bigquery.LoadJobConfig()

        if schema_path != '':
            self.initiate_table(
                table_id=table_id,
                schema_path=schema_path,
                dataset_id=dataset_id,
                project_id=project_id
            )
            job_config.schema = read_table_schema_from_file(schema_path)
            if header:
                job_config.skip_leading_rows = 1
        else:
            job_config.schema = None
            job_config.autodetect = True

        job_config.source_format = bigquery.SourceFormat.CSV
        job_config.write_disposition = write_disposition

        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )

        load_job = self.client.load_table_from_uri(
            uri,
            table_ref,
            job_config=job_config,
            location=location
        )

        load_job.result()  # Waits for job to complete.

        logging.info(
            'Loaded %s to %s:%s.%s',
            uri,
            project_id,
            dataset_id,
            table_id
        )

    def load_from_db(
        self,
        dataset_id,
        table_id,
        project_id=bq_default_project(),
        schema_path="",
        write_disposition="WRITE_TRUNCATE",
        source=None,
        source_url=None,
        sql=None,
        file=None,
        **kwargs
    ):
        """ Load data from another db into a BigQuery table
        Note:
            Sqlalchemy's built-in drivers should be supported
            Additional drivers can be added using requirements.txt
            Postgres(pg8000) and MySQL(pymysql) are included in tests

        Args:
            dataset_id (str): BigQuery dataset ID.
            table_id (str): BigQuery table ID.
            project_id (str): BigQuery project ID.
            schema_path (str): Path to schema file, if not set then BQ will auto-detect when creating a new table.
            write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.
            source (str): Defines connection using env vars with this prefix, e.g. "ERP", "MAGENTO", etc.
            source_url (str): Defines connection using SQLAlchemy database URL format, e.g.
            sql (str): SQL query (or table name).
            file (str): path to the SQL file relative to PROJECT_ROOT env var.
            **kwargs: an be applied to pass parameters into SQL file
        """
        db = DBExecutor(name=source, url=source_url)
        df = db.execute_query(sql=sql, file=file, **kwargs)
        db.clean_up()

        # By default, convert_dtypes will attempt to convert a Series
        # (or each Series in a DataFrame) to dtypes that support pd.NA
        # This prevents, for example, int64 being force to float in columns that contain nulls
        df = df.convert_dtypes()

        gcs = GCSExecutor()
        gcs_path = f"temp_load_from_db/{dataset_id}/{table_id}.csv"
        gcs.df_to_gcs(
            gcs_path=gcs_path,
            df=df
        )

        self.load_gcs(
            project_id=project_id,
            dataset_id=dataset_id,
            table_id=table_id,
            schema_path=schema_path,
            gcs_path=gcs_path,
            header=True,
            write_disposition=write_disposition
        )

        gcs.delete_directory(gcs_directory="temp_load_from_db")


    def insert_rows_json(self, dataset_id, table_id, rows,project_id=bq_default_project()):
        """ Insert rows into a table via the streaming API.
        Requires the table to exists in BigQuery.

        Arguments:
            - project_id (string): BigQuery project ID
            - dataset_id (string): BigQuery dataset ID
            - table_id (string): BigQuery table ID
            - rows (one of: list of tuples, list of dictonaries): Row data to be inserted.
            If a list of tuples is given, each tuple should contain data for each schema
            field on the current table and in the same order as the schema fields.
            If a list of dictionaries is given, the keys must include all required fields
            in the schema. Keys which do not correspond to a field in the schema are ignored.
        """
        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )
        try:
            error_response = self.client.insert_rows_json(
                table_ref,
                rows,
                ignore_unknown_values=True
            )
            if error_response == []:
                logging.info(
                    'Loaded %s row(s) into %s:%s.%s',
                    len(rows),
                    project_id,
                    dataset_id,
                    table_id
                )
            else:
                if hasattr(error_response, 'errors'):
                    for error in error_response.errors:
                        logging.warning(' - ' + error)
                        logging.warning(' - ' + rows)
                else:
                    logging.warning(json.dumps(error_response))

        except Exception:
            logging.exception('BigQuery Exception', exc_info=True)

    def extract_table_to_gcs(self, dataset_id, table_id, gcs_path,project_id=bq_default_project(), gcs_bucket=gcs_default_bucket(), location='US', shard=False):
        """ Extract BigQuery table into Google Cloud Storage.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
            gcs_bucket (string): Google Cloud Storage Bucket.
            gcs_path (string): Google Cloud Storage Path.
            location (string): Defaults to 'US'.
            shard (bool)
        """

        uri = "gs://{}/{}".format(gcs_bucket, gcs_path)
        if shard:
            uri = re.sub('\\.', '*.', uri)

        table_ref = self.get_table_ref(
            dataset_id=dataset_id,
            table_id=table_id,
            project_id=project_id
        )

        job = self.client.extract_table(
            table_ref,
            uri,
            location=location
        )

        job.result()  # Waits for job to complete.

        logging.info(
            'Table %s:%s.%s loaded to %s',
            project_id,
            dataset_id,
            table_id,
            uri
        )

    def truncate_table(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Delete all records from table but preserve structure e.g. schema, partitioning, clustering, description, labels, etc.

        Parameters:
            table_id (string): BigQuery table ID.
            dataset_id (string): BigQuery dataset ID.
            project_id (string): BigQuery project ID.
        """
        query=f"DELETE FROM `{project_id}.{dataset_id}.{table_id}` WHERE TRUE"
        self.execute_sql(sql=query, project_id=project_id)
        logging.info(
            'Table %s:%s.%s has been truncated',
            project_id,
            dataset_id,
            table_id
        )

    def copy_table_structure(
        self,
        source_table_id,
        dest_table_id,
        source_dataset_id=bq_default_dataset(),
        dest_dataset_id=bq_default_dataset(),
        source_project_id=bq_default_project(),
        dest_project_id=bq_default_project()
    ):


        if not self.dataset_exists(dest_dataset_id):
            self.create_dataset(dest_dataset_id)

        if self.table_exists(table_id=dest_table_id, dataset_id=dest_dataset_id):
            self.delete_table(table_id=dest_table_id, dataset_id=dest_dataset_id)

        if self.table_exists(table_id=source_table_id, dataset_id=source_dataset_id, project_id=source_project_id):

            attributes = self.get_table_attributes(
                table_id=source_table_id,
                dataset_id=source_dataset_id,
                project_id=source_project_id
            )
            table_def = "{}.{}.{}".format(dest_project_id, dest_dataset_id, dest_table_id)
            table = bigquery.Table(table_def)

            for attribute_key, attribute_value in attributes.items():
                if not attribute_value is None:
                    setattr(table, attribute_key, attribute_value)

            self.client.create_table(table) # Make an API request.
            logging.info(
                "Created table {}.{}.{}".format(table.project, table.dataset_id, table.table_id)
            )



    def copy_table(self, source_table_id, dest_table_id,
                   source_dataset_id=bq_default_dataset(), dest_dataset_id=bq_default_dataset(),
                   source_project_id=bq_default_project(),dest_project_id = bq_default_project(), write_disposition='WRITE_TRUNCATE'):
        """ Copy a BigQuery table.

        Parameters:
            source_project_id (string): Source BigQuery Project ID.
            source_dataset_id (string): Source BigQuery Dataset ID.
            source_table_id (string): Source BigQuery table ID.
            dest_project_id (string): Source BigQuery Project ID.
            dest_dataset_id (string): Destination BigQuery Dataset ID.
            dest_table_id (string): Destination BigQuery table ID.
            write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.

        Limitations:
            Destination project is the default, i.e. specified in `local.env`
            Destination dataset must reside in the same location as source (US, EU, etc.)

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the source
        """
        source_table_ref = self.get_table_ref(source_dataset_id, source_table_id,project_id=source_project_id)
        dest_table_ref = self.get_table_ref(dest_dataset_id, dest_table_id,project_id=dest_project_id)

        if not self.dataset_exists(dest_dataset_id):
            self.create_dataset(dest_dataset_id)

        job_config = bigquery.CopyJobConfig()
        job_config.write_disposition = set_write_disposition(write_disposition)
        job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED
        job = self.client.copy_table(
            source_table_ref,
            dest_table_ref,
            project=source_project_id,
            job_config=job_config
        )
        job.result()
        logging.info(
            'Table %s:%s.%s copied to %s:%s.%s',
            source_project_id,
            source_dataset_id,
            source_table_id,
            dest_project_id,
            dest_dataset_id,
            dest_table_id
        )

    def list_datasets(self, project_id=bq_default_project(), include_all=False):
        """ Lists datasets in project

        Parameters:
            project_id (string): BigQuery project ID
            include_all (boolean): True if results include hidden datasets. Defaults to False.

        Returns:
            A sorted list of dataset ids

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the project
        """
        try:
            datasets = list(self.client.list_datasets(project=project_id, include_all=include_all))
            dataset_ids = [dataset.dataset_id for dataset in datasets]
            return dataset_ids
        except NotFound as error:
            logging.error(error)

    def list_tables(self, dataset_id=bq_default_dataset(), project_id=bq_default_project()):
        """ Lists tables in dataset

        Parameters:
            project_id (string): BigQuery project ID
            dataset_id (string): BigQuery dataset ID

        Returns:
            A sorted list of table ids

        Permissions:
            The email address specified in `access_token.json` must have read permissions for the project
        """
        try:
            dataset_ref = self.get_dataset_ref(dataset_id, project_id = project_id)
            tables = self.client.list_tables(dataset_ref)
            table_ids = [table.table_id for table in tables]
            return table_ids
        except NotFound as error:
            logging.error(error)

    def count_rows(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count rows in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.

        Returns:
            Number of rows in a BigQuery table.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)
        return table.num_rows

    def count_columns(self, table_id, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count columns in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.

        Returns:
            Number of columns in a BigQuery table.
        """
        table_ref = self.get_table_ref(dataset_id, table_id,project_id = project_id)
        table = self.client.get_table(table_ref)
        return len(table.schema)

    def count_duplicates(self, table_id, primary_key: list, dataset_id=bq_default_dataset(),project_id=bq_default_project()):
        """ Count duplicate rows in primary key

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.
            primary_key: a list of one or more column names, e.g. ['col1', 'col2']

        Returns:
            non-negative integer
        """
        if not primary_key:
            logging.warning("No primary key supplied, skipping duplicates count.")
            return 0
        data = self.execute_sql(
            f"""
            SELECT
                COALESCE(SUM(dup_count), 0) AS dup_total
            FROM (
                SELECT
                    {', '.join(primary_key)},
                    (COUNT(*) - 1) AS dup_count
                FROM
                    `{dataset_id}.{table_id}`
                GROUP BY
                    {', '.join(primary_key)}
            )
            """,
            project_id = project_id
        )
        return data['dup_total'].values[0]

    def assert_unique(self, table_id, primary_key: list, dataset_id=bq_default_dataset(),project_id=bq_default_project(), ignore_error=False, **kwargs):
        """ Assert uniqueness of primary key in table

        Parameters:
            project_id (string): BigQuery project ID.
            dataset_id (string): BigQuery dataset ID.
            table_id (string): BigQuery table ID.
            primary_key: a list of one or more column names, e.g. ['col1', 'col2']
            ignore_error: boolean flag to prevent error being raised, useful for debugging

        Returns:
            Nothing if there are no duplicate rows in primary key
            Raise and log AssertionError if there are duplicate rows and ignore_error=False (default)
            Log a warning if there are dupicate rows and ignore_error=True (debugging)
        """
        if self.count_duplicates(table_id, primary_key, dataset_id, project_id=project_id) != 0:
            msg = "Table %s:%s.%s is not unique on %s" % (
                project_id,
                dataset_id,
                table_id,
                ', '.join(primary_key)
            )
            if ignore_error:
                logging.warning(msg)
            else:
                raise AssertionError(msg)

    def assert_acceptance(self, sql, cte, output_table_name='expected_output', **kwargs):
        """
        Compares the target cte defined in the mock file with the output of the sql.

        Notes:
        - order of rows and columns doesn't matter
        - only handles 1 layer of nesting
        """
        try:
            sql_extract_output_table = "WITH {} ( SELECT * FROM `{}` )".format(cte, output_table_name)

            df_result = self.execute_sql(sql_extract_output_table)
        except:
            raise AssertionError("Bad SQL or CTE")
        composite_sql = "WITH {} ( {} )".format(cte, sql)
        df = self.execute_sql(composite_sql)

        # flatten and sort - ignore column and row order
        df = flatten_df_columns(df)
        df = df[df.columns.sort_values()]
        df = df.sort_values(list(df.columns)).reset_index(drop=True)
        df_result = flatten_df_columns(df_result)
        df_result = df_result[df_result.columns.sort_values()]
        df_result = df_result.sort_values(list(df_result.columns)).reset_index(drop=True)
        assert_frame_equal(df, df_result)

Module containing BigQueryExecutor for Python Read SQL file and apply arguments/keyword arguments.      If dataset_prefix is a kwarg then this function is being used in a dry-run test.     In this case, the dataset_prefix is added to the dataset name and project names     are removed from the SQL.      Args:         file (string): path to the SQL file from PROJECT_ROOT environment variable.         *kwargs can be passed if some parameters are to be passed.      Returns:         a SQL formated with **kwargs if applicable.      Example:         With SQL as:             "select .. {param2} .. {param1} .. {paramN}"         *kwargs as:             param1=value1             param2=value2             paranN=valueN         The functions returns:             "select .. value2 .. value1 .. valueN" BigQuery handler      Parameters:         project_id (sql_file): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.      Required:         GOOGLE_APPLICATION_CREDENTIALS (env variable).      Attributes:         client (Client object)         credentials (Credentials object)         project_id (string): BigQuery Project. Defaults to BIGQUERY_PROJECT environment variable.      Returns:         a BigQueryExecutor object. Resets client and credentials. Sets BigQuery client. Returns BigQuery DatasetReference object.          Parameters:             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             bigquery.dataset.DatasetReference object Returns BigQuery Table reference object.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             Returns BigQuery Table reference object. Checks if a BigQuery dataset exists.          Parameters:             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             True is the dataset exists, False otherwise Deletes a BigQuery dataset.          Parameters:             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             delete_contents (bool): If True, delete all the tables in the dataset.             If False and the dataset contains tables, the request will fail. Default is False. Creates a BigQuery dataset if the dataset does not exists. Otherwise pass.          Parameters:             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID. Checks if a BigQuery table exists.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.           Returns:             True is the table exists, False otherwise. Delete a BigQuery table.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID. Initiate a BigQuery table. If the table already exists, compares the schema_path and apply a patch if there is a schema change.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string: Path to the BigQuery table schema from the PROJECT_ROOT environement variable.             partition (bool): Specify whether the BigQuery table is partioned. Default to False.             partition_field (string): Specify partition_field if table partioned. Defaults to "_PARTITIONTIME".             clustering (list): List of clustering fields. Defaults to None. create a bigquery table from a sql query Delete partitions from table in backfill dataset between from_partition and to_partition          Args:             backfill_dataset_id (string): dataset where the backfill tables are created             table_id (string): table in backfill dataset that the partitions will be deleted from             partition_field (string): the time-unit column that the table is partitioned by             from_partition (date): date of the first partition to be deleted             to_partition (date): date of the latest partition to be deleted          Returns:             None DELETE                  `{backfill_dataset_id}.{table_id}`                WHERE                  CAST({partition_field} AS DATE) BETWEEN                    CAST('{from_partition}' AS DATE) AND CAST('{to_partition}' AS DATE) # wait for job to complete Run a full backfill for a time-unit column partitioned table          Args:             dataset_id (string): dataset of the table to be backfilled             backfill_dataset_id (string): dataset where the backfill tables are created             table_id (string): name of the table to be backfilled             schema_path (string): file path to the schema of the table to be backfilled             sql (string): query that's used in the pipeline to create this table                 Note: it has to contain the clause:                 WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)             partition_field (string): the time-unit column that the table is partitioned by          Returns:             None Run a partial backfill for a time-unit column partitioned table between the dates         from_partition and to_partition          Args:             dataset_id (string): dataset of the table to be backfilled             backfill_dataset_id (string): dataset where the backfill tables are created             table_id (string): name of the table to be backfilled             schema_path (string): file path to the schema of the table to be backfilled             sql (string): query that's used in the pipeline to create this table                 Note: it has to contain the clause:                 WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)             partition_field (string): the time-unit column that the table is partitioned by             from_partition (date): date of the first partition to be backfilled             to_partition (date): date of the latest partition to be backfilled          Returns:             None Run a full or partial backfill for a time-unit column partitioned table          Args:             backfill (string): specifies full or partial backfill (if running a backfill)             dataset_id (string): dataset of the table to be backfilled             table_id (string): name of the table to be backfilled             schema_path (string): file path to the schema of the table to be backfilled             sql (string): query that's used in the pipeline to create this table                 Note: it has to contain the clause:                 WHERE <partition_field> BETWEEN CAST(<from_partition> AS DATE) AND CAST(<to_partition> AS DATE)             partition_field (string): the time-unit column that the table is partitioned by          Returns:             None New and legacy ways of creating partition tables:          Legacy:         Partition to be generated are either passed through partition_dates or automatically generated using existing partitions.         To filter on a specific partition, the filter DATE(_PARTITIONTIME) = {partition_date} can be used in your sql query.          New:         pipeline_partition_date is required for a normal pipeline run and is a single date.         To run a backfill, the BACKFILL env var needs to be set to 'full' or 'partial' and an optional date range          can be set for a partial backfill by setting the FROM_PARTITION and TO_PARTITION env vars. # legacy approach # new approach, no backfill # change date format from YYYY-MM-DD to YYYYMMDD for partition_id # new approach, run backfill Apply partition_date to the SQL query.          Parameters:             sql (string): the SQL query             partition_date (string): the partiton date          Returns:             Formatted SQL query. Validates the partition_dates parameter. Checks whether it is a list and the elements are in the right format.          Parameters:             partition_dates (list of string): list of partition_dates             partition_date (string): the partiton date          Raises:             BigQueryExecutorError if the partition_dates parameter in not a list.             ValueError if incorrect date format. Validates the date and sets the partition_name associated.          Parameters:             table (string): table name             date (string): date with format '%Y%m%d'          Returns:             Partition name          Raises:             BigQueryExecutorError if the partition_dates parameter in not a list.             ValueError if incorrect date format. Returns the partitions dates required based on start_date, end_date and existing_dates.          Parameters:             start_date (string): First partition required.             end_date (string): Last partition required.             existing_dates (list of string): List of existing partition.          Returns:             List of partition dates required. Gets existing partitions from BigQuery as dataframe.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             DataFrame of existing partitions SELECT                     FORMAT_DATE('%Y%m%d', DATE(_PARTITIONTIME)) AS partition_id                   FROM                     `{dataset_id}.{table_id}`                   GROUP BY                     1 Gets existing partitions.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             List of existing partition dates. # checks that res has number of rows > 0 # if not, no existing partitions Gets table schema object          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.           Returns:             Table schema. Gets table partitioning_type          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             partitioning_type is 'DAY' if the table is partitioned, None otherwise. Gets table clustering_fields          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.          Returns:             list of clustering fields is the table has cluster fields, None otherwise. Identify new fields in based on a schema file.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to the schema to compare to.          Returns:             List of new fields. Append fields to a BigQuery table.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             field (string): Schema field object. # API request # creates a copy of the schema # API request Apply a patch to a BigQuery Table if required.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to schema file to compare to.          Returns:             Lenght of new schema Perform a schema update. Used to update descriptions.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to schema file to compare to.          Raises:             BadRequest if the update fails. # API request # API request Performs a table update to fill in description for the table.          Parameters:             table_id (string): BigQuery table ID.             description (string): The descriptive text to describe the content of the table.             project_id (string): BigQuery project ID.             dataset_id (string): BigQuery dataset ID.          Raises:             BadRequest if the update fails. # API request # API request Executes a SQL query and loads it as a DataFrame.          Parameters:             sql (string): SQL Query.             project_id (string): BigQuery Project ID.             dialect (string): BigQuery dialect. Defaults to standard.          Returns:             Query result as a DataFrame. Executes a SQL file and loads it as a DataFrame.          Parameters:             file (string): Path to SQL file.             project_id (string): BigQuery Project ID.             dialect (string): BigQuery dialect. Defaults to standard.          **kwargs can be passed if the SQL file contains arguments formatted with {}.          Forbidden kwargs:             partition_date          Returns:             Query result as a DataFrame. Loads DataFrame to BigQuery table.          Parameters:             df (pd.DataFrame): Pandas DataFrame             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to schema file.             write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY. Creates BigQuery Table with live connection to Google Sheets          Args:             table_id (str): BigQuery table ID             dataset_id (str): BigQuery dataset ID             project_id (str): BigQuery project ID             schema_path (str): Path to schema file, if not set then BQ will auto-detect             googlesheet_uri (str): Google Sheet URI             googlesheet_key (str): Google Sheet Key, an alternate option instead of URI             sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0             header (bool): Defaults to True Loads Google Sheets data into a normal BigQuery Table          Args:             table_id (str): BigQuery table ID             dataset_id (str): BigQuery dataset ID             project_id (str): BigQuery project ID             schema_path (str): Path to schema file, if not set then BQ will auto-detect             googlesheet_uri (str): Google Sheet URI             googlesheet_key (str): Google Sheet Key, an alternate option instead of URI             sheet_name (str): GS Sheet Name, defaults to first worksheet, index 0             description (str): The descriptive text to describe the content of the table             header (bool): Defaults to True             write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY Loads JSON file to BigQuery table.          Parameters:             file (string): Path to JSON file.             project_id (string): BigQuery Project ID.             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to schema file.             write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY. Loads JSON data to BigQuery table.          Parameters:             json (string): JSON data.             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             schema_path (string): Path to schema file.             write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY. Loads Google Cloud Storage CSV file into a BigQuery table.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             gcs_bucket (string): Google Cloud Storage Bucket.             gcs_path (string): Google Cloud Storage Path.             location (string): Defaults to 'US'.             schema NBD             header (bool): Defaults to True.             write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY. # Waits for job to complete. Load data from another db into a BigQuery table         Note:             Sqlalchemy's built-in drivers should be supported             Additional drivers can be added using requirements.txt             Postgres(pg8000) and MySQL(pymysql) are included in tests          Args:             dataset_id (str): BigQuery dataset ID.             table_id (str): BigQuery table ID.             project_id (str): BigQuery project ID.             schema_path (str): Path to schema file, if not set then BQ will auto-detect when creating a new table.             write_disposition (str): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.             source (str): Defines connection using env vars with this prefix, e.g. "ERP", "MAGENTO", etc.             source_url (str): Defines connection using SQLAlchemy database URL format, e.g.             sql (str): SQL query (or table name).             file (str): path to the SQL file relative to PROJECT_ROOT env var.             **kwargs: an be applied to pass parameters into SQL file # By default, convert_dtypes will attempt to convert a Series # (or each Series in a DataFrame) to dtypes that support pd.NA # This prevents, for example, int64 being force to float in columns that contain nulls Insert rows into a table via the streaming API.         Requires the table to exists in BigQuery.          Arguments:             - project_id (string): BigQuery project ID             - dataset_id (string): BigQuery dataset ID             - table_id (string): BigQuery table ID             - rows (one of: list of tuples, list of dictonaries): Row data to be inserted.             If a list of tuples is given, each tuple should contain data for each schema             field on the current table and in the same order as the schema fields.             If a list of dictionaries is given, the keys must include all required fields             in the schema. Keys which do not correspond to a field in the schema are ignored. Extract BigQuery table into Google Cloud Storage.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID.             gcs_bucket (string): Google Cloud Storage Bucket.             gcs_path (string): Google Cloud Storage Path.             location (string): Defaults to 'US'.             shard (bool) # Waits for job to complete. Delete all records from table but preserve structure e.g. schema, partitioning, clustering, description, labels, etc.          Parameters:             table_id (string): BigQuery table ID.             dataset_id (string): BigQuery dataset ID.             project_id (string): BigQuery project ID. # Make an API request. Copy a BigQuery table.          Parameters:             source_project_id (string): Source BigQuery Project ID.             source_dataset_id (string): Source BigQuery Dataset ID.             source_table_id (string): Source BigQuery table ID.             dest_project_id (string): Source BigQuery Project ID.             dest_dataset_id (string): Destination BigQuery Dataset ID.             dest_table_id (string): Destination BigQuery table ID.             write_disposition (string): Write disposition. Can be one of WRITE_TRUNCATE, WRITE_APPEND or WRITE_EMPTY.          Limitations:             Destination project is the default, i.e. specified in `local.env`             Destination dataset must reside in the same location as source (US, EU, etc.)          Permissions:             The email address specified in `access_token.json` must have read permissions for the source Lists datasets in project          Parameters:             project_id (string): BigQuery project ID             include_all (boolean): True if results include hidden datasets. Defaults to False.          Returns:             A sorted list of dataset ids          Permissions:             The email address specified in `access_token.json` must have read permissions for the project Lists tables in dataset          Parameters:             project_id (string): BigQuery project ID             dataset_id (string): BigQuery dataset ID          Returns:             A sorted list of table ids          Permissions:             The email address specified in `access_token.json` must have read permissions for the project Count rows in table          Parameters:             project_id (string): BigQuery project ID.             dataset_id (string): BigQuery dataset ID.             table_id (string): BigQuery table ID.          Returns:             Number of rows in a BigQuery table. Count columns in table          Parameters:             project_id (string): BigQuery project ID.             dataset_id (string): BigQuery dataset ID.             table_id (string): BigQuery table ID.          Returns:             Number of columns in a BigQuery table. Count duplicate rows in primary key          Parameters:             project_id (string): BigQuery project ID.             dataset_id (string): BigQuery dataset ID.             table_id (string): BigQuery table ID.             primary_key: a list of one or more column names, e.g. ['col1', 'col2']          Returns:             non-negative integer SELECT                 COALESCE(SUM(dup_count), 0) AS dup_total             FROM (                 SELECT                     {', '.join(primary_key)},                     (COUNT(*) - 1) AS dup_count                 FROM                     `{dataset_id}.{table_id}`                 GROUP BY                     {', '.join(primary_key)}             ) Assert uniqueness of primary key in table          Parameters:             project_id (string): BigQuery project ID.             dataset_id (string): BigQuery dataset ID.             table_id (string): BigQuery table ID.             primary_key: a list of one or more column names, e.g. ['col1', 'col2']             ignore_error: boolean flag to prevent error being raised, useful for debugging          Returns:             Nothing if there are no duplicate rows in primary key             Raise and log AssertionError if there are duplicate rows and ignore_error=False (default)             Log a warning if there are dupicate rows and ignore_error=True (debugging) Compares the target cte defined in the mock file with the output of the sql.          Notes:         - order of rows and columns doesn't matter         - only handles 1 layer of nesting # flatten and sort - ignore column and row order