"""===============================================================================

        FILE: /Users/nailbiter/Documents/datawise/alex_python_toolbox/alex_python_toolbox/bigquery_caching.py

       USAGE: (not intended to be directly executed)

 DESCRIPTION:

     OPTIONS: ---
REQUIREMENTS: ---
        BUGS: ---
       NOTES: ---
      AUTHOR: <NAME> (<EMAIL>)
ORGANIZATION: Datawise Inc.
     VERSION: ---
     CREATED: 2020-11-02T18:39:07.242011
    REVISION: ---

TODO:
    2. determine `sources` automatically from sql when `sources==None`
    4. make it usable with ThreadPool, ProcessPool
    5. make robust w.r.t. non-trivial data-types in table (e.g. arrays etc.); at least ensure clean failure
    6. add CLI-tool (commands: show --what tables|sqls, remove-table)
    9. 0a5cafd7878d3d0a -- make prefix more distinct than `t_` -- breaking change
    10. add option `save-to-disk` and check it is parallel
    11. FIXME: add insted of update to _SQL_SOURCE_TABLE_NAME
    14. (breaking_change)when using `sources`: add source names in computation, not only modification times; also remove duplicates
    15. add `purge` functionality to cache;
    17. optimize
    20. try faster hash computation methods
    21. persistent `get_total_bytes_processed` across runs
    22. make `destination` optional in `query_to_table` (in that case, encode label in table name)
    23. `get_total_bytes_processed_by_sql`: output zero if already in cache
    24(done). output how many time spent on sql
==============================================================================="""
from google.cloud import bigquery
import hashlib
import sqlite3
import pandas as pd
import logging
from os import makedirs
from os.path import join
from datetime import datetime
import sqlparse
import google
import json
import time
import itertools
import time


class BigqueryCache:
    _SQL_SOURCE_TABLE_NAME = "bigquery_cache_sql_sources"

    def __init__(self, bigquery_client=None, cache_db="bigquery_cache.db", cache_table_settings=None):
        """
        cache_table_settings: {dataset: str [, prefix: str, assume_sync: boolean]}
        """
        # FIXME: change cache_table_settings into `namedtuple`
        if bigquery_client is None:
            bigquery_client = bigquery.Client()
        self._cache_db = cache_db
        self._client = bigquery_client
        self._logger = logging.getLogger(self.__class__.__name__)
        self._total_bytes_processed = 0
        self._wait_after_table_creation_sec = 0
        self._status = None

        if cache_table_settings is not None:
            self._cache_table_settings = {
                **{"prefix": "", "assume_sync": False},
                **cache_table_settings,
            }
        else:
            self._cache_table_settings = None
        if self._cache_table_settings is not None and self._cache_table_settings["assume_sync"]:
            _fetched = {}
            # FIXME: potentially, this may cause troubles if we have colliding names in tables/models
            # note that this won't happen if tables/models in dataset are created exclusively through this package
            for k in "table,model".split(","):
                _fetched[f"{k}s"] = [
                    getattr(t, f"{k}_id")
                    for t
                    in getattr(bigquery_client, f"list_{k}s")(dataset=self._cache_table_settings["dataset"])
                    if getattr(t, f"{k}_id").startswith(self._cache_table_settings["prefix"])
                ]
            self._cache_table_settings["_cached_tables"] = {
                f"{self._cache_table_settings['dataset']}.{x}"
                for x
                in itertools.chain(*_fetched.values())
            }

    def get_prefix(self):
        return self._cache_table_settings["dataset"]+"."+self._cache_table_settings["prefix"]

    def set_wait_after_table_creation_sec(self, wait_after_table_creation_sec):
        assert wait_after_table_creation_sec >= 0
        self._wait_after_table_creation_sec = wait_after_table_creation_sec

    @staticmethod
    def _list_tables(client, cache, project, dataset):
        if (project, dataset) not in cache:
            cache[(project, dataset)] = [
                t.table_id for t in client.list_tables(dataset=f"{project}.{dataset}")]
        return cache[(project, dataset)]

    def purge(self, table_hash):
        conn = self._get_conn()
        c = conn.cursor()
        tn = f"t_{table_hash}"
        c.execute(f"drop table if exists {tn};")
        c.execute("vacuum")
        self._logger.warning(f"dropped {tn}")
        conn.commit()
        conn.close()

    def show(self, what):
        conn = self._get_conn()
        if what == "tables":
            res = pd.read_sql_query(
                "SELECT name FROM sqlite_master WHERE type='table'", conn)
            sqls_df = self._get_sql_sources_df(conn)  # FIXME: use `show`

            res = pd.DataFrame([
                {
                    "hash": name[len("t_"):],
                    "row_count":pd.read_sql_query(f"select count(*) as row_count from {name}", conn).to_dict(orient="records")[0]["row_count"],
                }
                for name in res["name"]
                if name.startswith("t_")
            ])
            conn.close()

            res = res.set_index("hash")
            sqls_df = sqls_df.rename(columns={"sql_hash": "hash"})
            sqls_df["is_cache_hit"] = sqls_df["is_cache_hit"].apply(
                lambda f: f == 1.0)

            sqls_df_created = sqls_df[[
                is_cache_hit for is_cache_hit in sqls_df.is_cache_hit]]
            sqls_df_created = pd.DataFrame([
                {"hash": hash_, "creation_date": max(slice_["date"])}
                for hash_, slice_
                in sqls_df_created.groupby("hash")
            ])
            sqls_df_created = sqls_df_created.set_index("hash")

            sqls_df_modified = pd.DataFrame([
                {"hash": hash_, "last_access_date": max(slice_["date"])}
                for hash_, slice_
                in sqls_df.groupby("hash")
            ])
            sqls_df_modified = sqls_df_modified.set_index("hash")

            res = res.join(sqls_df_created).join(
                sqls_df_modified).reset_index()
            return res
        elif what == "sqls":
            res = self._get_sql_sources_df(conn)
            conn.close()
            return res
        else:
            conn.close()
            raise NotImplementedError

    def _get_conn(self):
        return sqlite3.connect(self._cache_db)

    def remove_table(i):
        pass

    _SQL_SOURCES_DF_HEADERS = {
        "sql_hash": 0, "sql_source": 0, "sources": 0, "date": 0, "is_cache_hit": 0, "destination": 0, "total_bytes_processed": 0, }

    def _get_empty_sql_sources_df(self):
        return pd.DataFrame(
            {k: [] for k in BigqueryCache._SQL_SOURCES_DF_HEADERS})

    def _get_sql_sources_df(self, conn):
        try:
            sql_sources_df = pd.read_sql_query(
                f"select * from {BigqueryCache._SQL_SOURCE_TABLE_NAME}", conn)
        except pd.io.sql.DatabaseError:
            self._logger.warning(
                f"could not fetch `{BigqueryCache._SQL_SOURCE_TABLE_NAME}` => returning empty table")
            sql_sources_df = self._get_empty_sql_sources_df()

        if set(sql_sources_df) != set(BigqueryCache._SQL_SOURCES_DF_HEADERS):
            self._logger.warning(
                f"{set(sql_sources_df)} => {set(BigqueryCache._SQL_SOURCES_DF_HEADERS)}")
            sql_sources_df = sql_sources_df.drop(columns=[
                                                 fn for fn in sql_sources_df if fn not in BigqueryCache._SQL_SOURCES_DF_HEADERS])
            for fn, v in BigqueryCache._SQL_SOURCES_DF_HEADERS.items():
                if fn not in sql_sources_df:
                    sql_sources_df[fn] = v
            self._logger.warning(f"preparing to rewrite")
            # self._logger.warning(sql_sources_df)
            sql_sources_df.to_sql(
                BigqueryCache._SQL_SOURCE_TABLE_NAME, conn, if_exists="replace", index=None)

        sql_sources_df = sql_sources_df.loc[:,
                                            BigqueryCache._SQL_SOURCES_DF_HEADERS]
        return sql_sources_df

    __TOKEN_BLACKLIST = "Token.Text.Whitespace.Newline Token.Comment.Single Token.Text.Whitespace".split(
        " ")

    def __update_sql_hash(self, m, sql):
        for stat in sqlparse.parse(sql):
            for token in stat.tokens:
                for token_ in token.flatten():
                    ttype = str(token_.ttype)
                    if ttype not in BigqueryCache.__TOKEN_BLACKLIST:
                        # self._logger.info(f"ttype: {ttype}")
                        m.update(str(token_).encode())

    @staticmethod
    def _expand_sources(sources, bq_client, logger=None):
        if logger is None:
            logger = logging.getLogger("bigquery_cache._expand_sources")

        _list_tables_cache = {}
        _sources = []
        ASTERISK = "*"
        for source in sources:
            logger.info(f"source: {source}")
            if source.endswith(ASTERISK):
                p, d, t = source.split(".")
                lt = BigqueryCache._list_tables(
                    bq_client, _list_tables_cache, p, d)
                prefix = t[:-len(ASTERISK)]
                _sources.extend(
                    [f"{p}.{d}.{tn}" for tn in lt if tn.startswith(prefix)])
            else:
                _sources.append(source)
        _sources = sorted(_sources)
        logger.info(f"_sources: {_sources}")
        return _sources

    def _get_sql_hash(self, sql, sources, additional_data={}):
        m = hashlib.sha256()
        self.__update_sql_hash(m, sql)

        _sources = BigqueryCache._expand_sources(
            sources, self._client, logger=self._logger)

        for source in _sources:
            m.update(str(self._client.get_table(source).modified).encode())
        m.update(json.dumps(additional_data, sort_keys=True).encode())
        return m.hexdigest(), _sources

    _BIGQUERY_CACHE_TABLE_DIGEST_KEY = "<KEY>"

    # https://cloud.google.com/bigquery/docs/labels-intro
    _MAX_LENGTH_OF_BQ_LABEL_VALUE = 63
    # https://cloud.google.com/bigquery/docs/tables#table_naming
    _MAX_LENGTH_OF_BQ_TABLE_NAME = 1024

    def _get_table_digest(self, table):
        try:
            table_ = self._client.get_table(table)
            if BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY not in table_.labels:
                return None
            else:
                return table_.labels[BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY][:BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        except google.api_core.exceptions.NotFound:
            return None

    def _set_table_digest(self, table, hexdigest):
        table_ = self._client.get_table(table)
        table_.labels[BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY] = hexdigest[:
                                                                                  BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        self._client.update_table(table_, ["labels"])

    @staticmethod
    def _is_hexdigest_same_as_remote_label_hexdigest(hexdigest, remote_hexdigest):
        hexdigest = hexdigest[:BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        return hexdigest == remote_hexdigest

    def _table_exists(self, table_name):
        res = True
        try:
            self._client.get_table(table_name)
        except google.api_core.exceptions.NotFound:
            res = False
        return res

    def df_to_table(self, df, force_overwrite=False, wait_after_creation_sec=None):
        if wait_after_creation_sec is None:
            wait_after_creation_sec = self._wait_after_table_creation_sec
        assert wait_after_creation_sec >= 0
        m = hashlib.sha256()
        for x in list(pd.util.hash_pandas_object(df)):
            m.update(str(x).encode())
        m.update(json.dumps(sorted(list(df))).encode())
        hash_ = m.hexdigest()
        tn = f"{self._cache_table_settings['prefix']}{hash_}"
        tn = tn[:BigqueryCache._MAX_LENGTH_OF_BQ_TABLE_NAME]
        destination = f"{self._cache_table_settings['dataset']}.{tn}"
        te = self._table_exists(destination)
        if not te or force_overwrite:
            self._logger.warning(f"creating table {destination}")
            self._client.delete_table(destination, not_found_ok=True)
            self._client.load_table_from_dataframe(
                dataframe=df, destination=destination)
            if wait_after_creation_sec > 0:
                time.sleep(wait_after_creation_sec)
        else:
            self._logger.warning(f"skipping table {destination}")
        return destination

    def query_to_table(self, sql, *args, destination=None, force_overwrite=False, dry_run=False, sources=[], partitioning_sql="", preamble_sql="", wait_after_creation_sec=None, entity_to_create="table", **kwargs):
        if wait_after_creation_sec is None:
            wait_after_creation_sec = self._wait_after_table_creation_sec
        assert wait_after_creation_sec >= 0
        assert destination is not None or self._cache_table_settings is not None, "destination is None and cache_table_settings was not set in constructor"

        hexdigest, expanded_sources = self._get_sql_hash(
            sql, sources, {"partitioning_sql": partitioning_sql, "preamble_sql": preamble_sql})
        if destination is None:
            hash_injection_method = "table_name"
            tn = f"{self._cache_table_settings['prefix']}{hexdigest}"
            tn = tn[:BigqueryCache._MAX_LENGTH_OF_BQ_TABLE_NAME]
            destination = f"{self._cache_table_settings['dataset']}.{tn}"
            if self._cache_table_settings["assume_sync"]:
                is_should_work = destination not in self._cache_table_settings["_cached_tables"]
                self._logger.info(
                    f"""{destination}, {self._cache_table_settings['_cached_tables']} => {is_should_work}""")
            else:
                try:
                    table_ = self._client.get_table(destination)
                    is_should_work = False
                except google.api_core.exceptions.NotFound:
                    is_should_work = True
            hash_injection_method = "table_name"
        else:
            remote_hexdigest = self._get_table_digest(destination)
            hash_injection_method = "label"
            is_should_work = BigqueryCache._is_hexdigest_same_as_remote_label_hexdigest(
                hexdigest, remote_hexdigest)

        is_cache_hit = True
        total_bytes_processed = 0
        _status = {}
        if force_overwrite or is_should_work:
            if dry_run:
                self._logger.warning(
                    f"no cached result for {hexdigest}, {destination}")
            else:
                tic = time.time()
                if len(partitioning_sql.strip()) > 0 or entity_to_create != "table":
                    # FIXME: do this without writing everything in SQL
                    sql = f"""
                    {preamble_sql}
                    create or replace {entity_to_create} `{destination}`
                    {partitioning_sql} as (
                        {sql}
                    )
                    """
                    self._client.query(sql, *args, **kwargs).result()
                else:
                    self._client.query(sql, *args, job_config=bigquery.job.QueryJobConfig(
                        destination=destination, write_disposition="WRITE_TRUNCATE"), **kwargs).result()
                toc = time.time()
                _status = {"tictoc": toc-tic, "is_cached": False}
                if wait_after_creation_sec > 0:
                    time.sleep(wait_after_creation_sec)
                total_bytes_processed = self.get_total_bytes_processed_by_sql(
                    sql)
                self._total_bytes_processed += total_bytes_processed
                if hash_injection_method == "label":
                    self._set_table_digest(destination, hexdigest)
                elif hash_injection_method == "table_name" and self._cache_table_settings["assume_sync"]:
                    self._cache_table_settings["_cached_tables"].add(
                        destination)
                is_cache_hit = False
        else:
            _status = {"is_cached": True}
            self._logger.warning(
                f"using cached result for {hexdigest}, {destination}")

        if not dry_run:
            self._status = _status
            self._append_to_sqls(sql, expanded_sources, hexdigest,
                                 is_cache_hit, total_bytes_processed, destination=destination)

        return destination

    def get_status(self):
        return {**self._status}

    def _append_to_sqls(self, sql, expanded_sources, sql_hash, is_cache_hit, total_bytes_processed, destination=None):
        sql_sources_df = self._get_empty_sql_sources_df()
        new_record = {
            "sql_hash": sql_hash,
            "sql_source": sql,
            "sources": json.dumps(expanded_sources),
            "date": datetime.now().isoformat(),
            "is_cache_hit": is_cache_hit,
            "destination": destination,
            "total_bytes_processed": total_bytes_processed
        }
        assert(set(new_record.keys()) == set(
            BigqueryCache._SQL_SOURCES_DF_HEADERS))
        sql_sources_df = sql_sources_df.append(new_record, ignore_index=True)
        self._logger.info(
            f"save {len(sql_sources_df)} records to {BigqueryCache._SQL_SOURCE_TABLE_NAME}")

        conn = self._get_conn()
        sql_sources_df.to_sql(
            BigqueryCache._SQL_SOURCE_TABLE_NAME, conn, if_exists="append", index=None)
        conn.close()

    def query_to_dataframe(self, sql, *args, force_overwrite=False, sources=[], to_dataframe_args=[], to_dataframe_kwargs={}, **kwargs):
        """
        `force_overwrite` means "force cache overwrite" (maybe change name later)
        """

        hexdigest, expanded_sources = self._get_sql_hash(sql, sources)
        sql_hash = hexdigest

        conn = self._get_conn()
        self._logger.info(f"sql_hash: {sql_hash}")
        is_cache_hit = True
        sql_cache_table_name = f"t_{sql_hash}"  # FIXME: 0a5cafd7878d3d0a
        total_bytes_processed = 0
        try:
            df = pd.read_sql_query(
                f"select * from {sql_cache_table_name}", conn)
            if force_overwrite:
                raise pd.io.sql.DatabaseError()
            self._logger.warning(f"using cached result for {hexdigest}")
        except pd.io.sql.DatabaseError:
            self._logger.info("no cached result")
            df = self._client.query(
                sql, *args, **kwargs).to_dataframe(*to_dataframe_args, **to_dataframe_kwargs)
            total_bytes_processed = self.get_total_bytes_processed_by_sql(sql)
            self._total_bytes_processed += total_bytes_processed
            df.to_sql(
                sql_cache_table_name, conn, if_exists="replace", index=None)
            is_cache_hit = False
        finally:
            conn.close()

        self._append_to_sqls(sql, expanded_sources, sql_hash,
                             is_cache_hit, total_bytes_processed)

        return df

    def get_total_bytes_processed_by_sql(self, sql):
        client = self._client
        return client.query(sql, job_config=bigquery.job.QueryJobConfig(dry_run=True, use_query_cache=False)).total_bytes_processed

    def get_total_bytes_processed(self):
        return self._total_bytes_processed

"""===============================================================================

        FILE: /Users/nailbiter/Documents/datawise/alex_python_toolbox/alex_python_toolbox/bigquery_caching.py

       USAGE: (not intended to be directly executed)

 DESCRIPTION:

     OPTIONS: ---
REQUIREMENTS: ---
        BUGS: ---
       NOTES: ---
      AUTHOR: <NAME> (<EMAIL>)
ORGANIZATION: Datawise Inc.
     VERSION: ---
     CREATED: 2020-11-02T18:39:07.242011
    REVISION: ---

TODO:
    2. determine `sources` automatically from sql when `sources==None`
    4. make it usable with ThreadPool, ProcessPool
    5. make robust w.r.t. non-trivial data-types in table (e.g. arrays etc.); at least ensure clean failure
    6. add CLI-tool (commands: show --what tables|sqls, remove-table)
    9. 0a5cafd7878d3d0a -- make prefix more distinct than `t_` -- breaking change
    10. add option `save-to-disk` and check it is parallel
    11. FIXME: add insted of update to _SQL_SOURCE_TABLE_NAME
    14. (breaking_change)when using `sources`: add source names in computation, not only modification times; also remove duplicates
    15. add `purge` functionality to cache;
    17. optimize
    20. try faster hash computation methods
    21. persistent `get_total_bytes_processed` across runs
    22. make `destination` optional in `query_to_table` (in that case, encode label in table name)
    23. `get_total_bytes_processed_by_sql`: output zero if already in cache
    24(done). output how many time spent on sql
==============================================================================="""
from google.cloud import bigquery
import hashlib
import sqlite3
import pandas as pd
import logging
from os import makedirs
from os.path import join
from datetime import datetime
import sqlparse
import google
import json
import time
import itertools
import time


class BigqueryCache:
    _SQL_SOURCE_TABLE_NAME = "bigquery_cache_sql_sources"

    def __init__(self, bigquery_client=None, cache_db="bigquery_cache.db", cache_table_settings=None):
        """
        cache_table_settings: {dataset: str [, prefix: str, assume_sync: boolean]}
        """
        # FIXME: change cache_table_settings into `namedtuple`
        if bigquery_client is None:
            bigquery_client = bigquery.Client()
        self._cache_db = cache_db
        self._client = bigquery_client
        self._logger = logging.getLogger(self.__class__.__name__)
        self._total_bytes_processed = 0
        self._wait_after_table_creation_sec = 0
        self._status = None

        if cache_table_settings is not None:
            self._cache_table_settings = {
                **{"prefix": "", "assume_sync": False},
                **cache_table_settings,
            }
        else:
            self._cache_table_settings = None
        if self._cache_table_settings is not None and self._cache_table_settings["assume_sync"]:
            _fetched = {}
            # FIXME: potentially, this may cause troubles if we have colliding names in tables/models
            # note that this won't happen if tables/models in dataset are created exclusively through this package
            for k in "table,model".split(","):
                _fetched[f"{k}s"] = [
                    getattr(t, f"{k}_id")
                    for t
                    in getattr(bigquery_client, f"list_{k}s")(dataset=self._cache_table_settings["dataset"])
                    if getattr(t, f"{k}_id").startswith(self._cache_table_settings["prefix"])
                ]
            self._cache_table_settings["_cached_tables"] = {
                f"{self._cache_table_settings['dataset']}.{x}"
                for x
                in itertools.chain(*_fetched.values())
            }

    def get_prefix(self):
        return self._cache_table_settings["dataset"]+"."+self._cache_table_settings["prefix"]

    def set_wait_after_table_creation_sec(self, wait_after_table_creation_sec):
        assert wait_after_table_creation_sec >= 0
        self._wait_after_table_creation_sec = wait_after_table_creation_sec

    @staticmethod
    def _list_tables(client, cache, project, dataset):
        if (project, dataset) not in cache:
            cache[(project, dataset)] = [
                t.table_id for t in client.list_tables(dataset=f"{project}.{dataset}")]
        return cache[(project, dataset)]

    def purge(self, table_hash):
        conn = self._get_conn()
        c = conn.cursor()
        tn = f"t_{table_hash}"
        c.execute(f"drop table if exists {tn};")
        c.execute("vacuum")
        self._logger.warning(f"dropped {tn}")
        conn.commit()
        conn.close()

    def show(self, what):
        conn = self._get_conn()
        if what == "tables":
            res = pd.read_sql_query(
                "SELECT name FROM sqlite_master WHERE type='table'", conn)
            sqls_df = self._get_sql_sources_df(conn)  # FIXME: use `show`

            res = pd.DataFrame([
                {
                    "hash": name[len("t_"):],
                    "row_count":pd.read_sql_query(f"select count(*) as row_count from {name}", conn).to_dict(orient="records")[0]["row_count"],
                }
                for name in res["name"]
                if name.startswith("t_")
            ])
            conn.close()

            res = res.set_index("hash")
            sqls_df = sqls_df.rename(columns={"sql_hash": "hash"})
            sqls_df["is_cache_hit"] = sqls_df["is_cache_hit"].apply(
                lambda f: f == 1.0)

            sqls_df_created = sqls_df[[
                is_cache_hit for is_cache_hit in sqls_df.is_cache_hit]]
            sqls_df_created = pd.DataFrame([
                {"hash": hash_, "creation_date": max(slice_["date"])}
                for hash_, slice_
                in sqls_df_created.groupby("hash")
            ])
            sqls_df_created = sqls_df_created.set_index("hash")

            sqls_df_modified = pd.DataFrame([
                {"hash": hash_, "last_access_date": max(slice_["date"])}
                for hash_, slice_
                in sqls_df.groupby("hash")
            ])
            sqls_df_modified = sqls_df_modified.set_index("hash")

            res = res.join(sqls_df_created).join(
                sqls_df_modified).reset_index()
            return res
        elif what == "sqls":
            res = self._get_sql_sources_df(conn)
            conn.close()
            return res
        else:
            conn.close()
            raise NotImplementedError

    def _get_conn(self):
        return sqlite3.connect(self._cache_db)

    def remove_table(i):
        pass

    _SQL_SOURCES_DF_HEADERS = {
        "sql_hash": 0, "sql_source": 0, "sources": 0, "date": 0, "is_cache_hit": 0, "destination": 0, "total_bytes_processed": 0, }

    def _get_empty_sql_sources_df(self):
        return pd.DataFrame(
            {k: [] for k in BigqueryCache._SQL_SOURCES_DF_HEADERS})

    def _get_sql_sources_df(self, conn):
        try:
            sql_sources_df = pd.read_sql_query(
                f"select * from {BigqueryCache._SQL_SOURCE_TABLE_NAME}", conn)
        except pd.io.sql.DatabaseError:
            self._logger.warning(
                f"could not fetch `{BigqueryCache._SQL_SOURCE_TABLE_NAME}` => returning empty table")
            sql_sources_df = self._get_empty_sql_sources_df()

        if set(sql_sources_df) != set(BigqueryCache._SQL_SOURCES_DF_HEADERS):
            self._logger.warning(
                f"{set(sql_sources_df)} => {set(BigqueryCache._SQL_SOURCES_DF_HEADERS)}")
            sql_sources_df = sql_sources_df.drop(columns=[
                                                 fn for fn in sql_sources_df if fn not in BigqueryCache._SQL_SOURCES_DF_HEADERS])
            for fn, v in BigqueryCache._SQL_SOURCES_DF_HEADERS.items():
                if fn not in sql_sources_df:
                    sql_sources_df[fn] = v
            self._logger.warning(f"preparing to rewrite")
            # self._logger.warning(sql_sources_df)
            sql_sources_df.to_sql(
                BigqueryCache._SQL_SOURCE_TABLE_NAME, conn, if_exists="replace", index=None)

        sql_sources_df = sql_sources_df.loc[:,
                                            BigqueryCache._SQL_SOURCES_DF_HEADERS]
        return sql_sources_df

    __TOKEN_BLACKLIST = "Token.Text.Whitespace.Newline Token.Comment.Single Token.Text.Whitespace".split(
        " ")

    def __update_sql_hash(self, m, sql):
        for stat in sqlparse.parse(sql):
            for token in stat.tokens:
                for token_ in token.flatten():
                    ttype = str(token_.ttype)
                    if ttype not in BigqueryCache.__TOKEN_BLACKLIST:
                        # self._logger.info(f"ttype: {ttype}")
                        m.update(str(token_).encode())

    @staticmethod
    def _expand_sources(sources, bq_client, logger=None):
        if logger is None:
            logger = logging.getLogger("bigquery_cache._expand_sources")

        _list_tables_cache = {}
        _sources = []
        ASTERISK = "*"
        for source in sources:
            logger.info(f"source: {source}")
            if source.endswith(ASTERISK):
                p, d, t = source.split(".")
                lt = BigqueryCache._list_tables(
                    bq_client, _list_tables_cache, p, d)
                prefix = t[:-len(ASTERISK)]
                _sources.extend(
                    [f"{p}.{d}.{tn}" for tn in lt if tn.startswith(prefix)])
            else:
                _sources.append(source)
        _sources = sorted(_sources)
        logger.info(f"_sources: {_sources}")
        return _sources

    def _get_sql_hash(self, sql, sources, additional_data={}):
        m = hashlib.sha256()
        self.__update_sql_hash(m, sql)

        _sources = BigqueryCache._expand_sources(
            sources, self._client, logger=self._logger)

        for source in _sources:
            m.update(str(self._client.get_table(source).modified).encode())
        m.update(json.dumps(additional_data, sort_keys=True).encode())
        return m.hexdigest(), _sources

    _BIGQUERY_CACHE_TABLE_DIGEST_KEY = "<KEY>"

    # https://cloud.google.com/bigquery/docs/labels-intro
    _MAX_LENGTH_OF_BQ_LABEL_VALUE = 63
    # https://cloud.google.com/bigquery/docs/tables#table_naming
    _MAX_LENGTH_OF_BQ_TABLE_NAME = 1024

    def _get_table_digest(self, table):
        try:
            table_ = self._client.get_table(table)
            if BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY not in table_.labels:
                return None
            else:
                return table_.labels[BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY][:BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        except google.api_core.exceptions.NotFound:
            return None

    def _set_table_digest(self, table, hexdigest):
        table_ = self._client.get_table(table)
        table_.labels[BigqueryCache._BIGQUERY_CACHE_TABLE_DIGEST_KEY] = hexdigest[:
                                                                                  BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        self._client.update_table(table_, ["labels"])

    @staticmethod
    def _is_hexdigest_same_as_remote_label_hexdigest(hexdigest, remote_hexdigest):
        hexdigest = hexdigest[:BigqueryCache._MAX_LENGTH_OF_BQ_LABEL_VALUE]
        return hexdigest == remote_hexdigest

    def _table_exists(self, table_name):
        res = True
        try:
            self._client.get_table(table_name)
        except google.api_core.exceptions.NotFound:
            res = False
        return res

    def df_to_table(self, df, force_overwrite=False, wait_after_creation_sec=None):
        if wait_after_creation_sec is None:
            wait_after_creation_sec = self._wait_after_table_creation_sec
        assert wait_after_creation_sec >= 0
        m = hashlib.sha256()
        for x in list(pd.util.hash_pandas_object(df)):
            m.update(str(x).encode())
        m.update(json.dumps(sorted(list(df))).encode())
        hash_ = m.hexdigest()
        tn = f"{self._cache_table_settings['prefix']}{hash_}"
        tn = tn[:BigqueryCache._MAX_LENGTH_OF_BQ_TABLE_NAME]
        destination = f"{self._cache_table_settings['dataset']}.{tn}"
        te = self._table_exists(destination)
        if not te or force_overwrite:
            self._logger.warning(f"creating table {destination}")
            self._client.delete_table(destination, not_found_ok=True)
            self._client.load_table_from_dataframe(
                dataframe=df, destination=destination)
            if wait_after_creation_sec > 0:
                time.sleep(wait_after_creation_sec)
        else:
            self._logger.warning(f"skipping table {destination}")
        return destination

    def query_to_table(self, sql, *args, destination=None, force_overwrite=False, dry_run=False, sources=[], partitioning_sql="", preamble_sql="", wait_after_creation_sec=None, entity_to_create="table", **kwargs):
        if wait_after_creation_sec is None:
            wait_after_creation_sec = self._wait_after_table_creation_sec
        assert wait_after_creation_sec >= 0
        assert destination is not None or self._cache_table_settings is not None, "destination is None and cache_table_settings was not set in constructor"

        hexdigest, expanded_sources = self._get_sql_hash(
            sql, sources, {"partitioning_sql": partitioning_sql, "preamble_sql": preamble_sql})
        if destination is None:
            hash_injection_method = "table_name"
            tn = f"{self._cache_table_settings['prefix']}{hexdigest}"
            tn = tn[:BigqueryCache._MAX_LENGTH_OF_BQ_TABLE_NAME]
            destination = f"{self._cache_table_settings['dataset']}.{tn}"
            if self._cache_table_settings["assume_sync"]:
                is_should_work = destination not in self._cache_table_settings["_cached_tables"]
                self._logger.info(
                    f"""{destination}, {self._cache_table_settings['_cached_tables']} => {is_should_work}""")
            else:
                try:
                    table_ = self._client.get_table(destination)
                    is_should_work = False
                except google.api_core.exceptions.NotFound:
                    is_should_work = True
            hash_injection_method = "table_name"
        else:
            remote_hexdigest = self._get_table_digest(destination)
            hash_injection_method = "label"
            is_should_work = BigqueryCache._is_hexdigest_same_as_remote_label_hexdigest(
                hexdigest, remote_hexdigest)

        is_cache_hit = True
        total_bytes_processed = 0
        _status = {}
        if force_overwrite or is_should_work:
            if dry_run:
                self._logger.warning(
                    f"no cached result for {hexdigest}, {destination}")
            else:
                tic = time.time()
                if len(partitioning_sql.strip()) > 0 or entity_to_create != "table":
                    # FIXME: do this without writing everything in SQL
                    sql = f"""
                    {preamble_sql}
                    create or replace {entity_to_create} `{destination}`
                    {partitioning_sql} as (
                        {sql}
                    )
                    """
                    self._client.query(sql, *args, **kwargs).result()
                else:
                    self._client.query(sql, *args, job_config=bigquery.job.QueryJobConfig(
                        destination=destination, write_disposition="WRITE_TRUNCATE"), **kwargs).result()
                toc = time.time()
                _status = {"tictoc": toc-tic, "is_cached": False}
                if wait_after_creation_sec > 0:
                    time.sleep(wait_after_creation_sec)
                total_bytes_processed = self.get_total_bytes_processed_by_sql(
                    sql)
                self._total_bytes_processed += total_bytes_processed
                if hash_injection_method == "label":
                    self._set_table_digest(destination, hexdigest)
                elif hash_injection_method == "table_name" and self._cache_table_settings["assume_sync"]:
                    self._cache_table_settings["_cached_tables"].add(
                        destination)
                is_cache_hit = False
        else:
            _status = {"is_cached": True}
            self._logger.warning(
                f"using cached result for {hexdigest}, {destination}")

        if not dry_run:
            self._status = _status
            self._append_to_sqls(sql, expanded_sources, hexdigest,
                                 is_cache_hit, total_bytes_processed, destination=destination)

        return destination

    def get_status(self):
        return {**self._status}

    def _append_to_sqls(self, sql, expanded_sources, sql_hash, is_cache_hit, total_bytes_processed, destination=None):
        sql_sources_df = self._get_empty_sql_sources_df()
        new_record = {
            "sql_hash": sql_hash,
            "sql_source": sql,
            "sources": json.dumps(expanded_sources),
            "date": datetime.now().isoformat(),
            "is_cache_hit": is_cache_hit,
            "destination": destination,
            "total_bytes_processed": total_bytes_processed
        }
        assert(set(new_record.keys()) == set(
            BigqueryCache._SQL_SOURCES_DF_HEADERS))
        sql_sources_df = sql_sources_df.append(new_record, ignore_index=True)
        self._logger.info(
            f"save {len(sql_sources_df)} records to {BigqueryCache._SQL_SOURCE_TABLE_NAME}")

        conn = self._get_conn()
        sql_sources_df.to_sql(
            BigqueryCache._SQL_SOURCE_TABLE_NAME, conn, if_exists="append", index=None)
        conn.close()

    def query_to_dataframe(self, sql, *args, force_overwrite=False, sources=[], to_dataframe_args=[], to_dataframe_kwargs={}, **kwargs):
        """
        `force_overwrite` means "force cache overwrite" (maybe change name later)
        """

        hexdigest, expanded_sources = self._get_sql_hash(sql, sources)
        sql_hash = hexdigest

        conn = self._get_conn()
        self._logger.info(f"sql_hash: {sql_hash}")
        is_cache_hit = True
        sql_cache_table_name = f"t_{sql_hash}"  # FIXME: 0a5cafd7878d3d0a
        total_bytes_processed = 0
        try:
            df = pd.read_sql_query(
                f"select * from {sql_cache_table_name}", conn)
            if force_overwrite:
                raise pd.io.sql.DatabaseError()
            self._logger.warning(f"using cached result for {hexdigest}")
        except pd.io.sql.DatabaseError:
            self._logger.info("no cached result")
            df = self._client.query(
                sql, *args, **kwargs).to_dataframe(*to_dataframe_args, **to_dataframe_kwargs)
            total_bytes_processed = self.get_total_bytes_processed_by_sql(sql)
            self._total_bytes_processed += total_bytes_processed
            df.to_sql(
                sql_cache_table_name, conn, if_exists="replace", index=None)
            is_cache_hit = False
        finally:
            conn.close()

        self._append_to_sqls(sql, expanded_sources, sql_hash,
                             is_cache_hit, total_bytes_processed)

        return df

    def get_total_bytes_processed_by_sql(self, sql):
        client = self._client
        return client.query(sql, job_config=bigquery.job.QueryJobConfig(dry_run=True, use_query_cache=False)).total_bytes_processed

    def get_total_bytes_processed(self):
        return self._total_bytes_processed

===============================================================================          FILE: /Users/nailbiter/Documents/datawise/alex_python_toolbox/alex_python_toolbox/bigquery_caching.py         USAGE: (not intended to be directly executed)   DESCRIPTION:       OPTIONS: --- REQUIREMENTS: ---         BUGS: ---        NOTES: ---       AUTHOR: <NAME> (<EMAIL>) ORGANIZATION: Datawise Inc.      VERSION: ---      CREATED: 2020-11-02T18:39:07.242011     REVISION: ---  TODO:     2. determine `sources` automatically from sql when `sources==None`     4. make it usable with ThreadPool, ProcessPool     5. make robust w.r.t. non-trivial data-types in table (e.g. arrays etc.); at least ensure clean failure     6. add CLI-tool (commands: show --what tables|sqls, remove-table)     9. 0a5cafd7878d3d0a -- make prefix more distinct than `t_` -- breaking change     10. add option `save-to-disk` and check it is parallel     11. FIXME: add insted of update to _SQL_SOURCE_TABLE_NAME     14. (breaking_change)when using `sources`: add source names in computation, not only modification times; also remove duplicates     15. add `purge` functionality to cache;     17. optimize     20. try faster hash computation methods     21. persistent `get_total_bytes_processed` across runs     22. make `destination` optional in `query_to_table` (in that case, encode label in table name)     23. `get_total_bytes_processed_by_sql`: output zero if already in cache     24(done). output how many time spent on sql =============================================================================== cache_table_settings: {dataset: str [, prefix: str, assume_sync: boolean]} # FIXME: change cache_table_settings into `namedtuple` # FIXME: potentially, this may cause troubles if we have colliding names in tables/models # note that this won't happen if tables/models in dataset are created exclusively through this package # FIXME: use `show` # self._logger.warning(sql_sources_df) # self._logger.info(f"ttype: {ttype}") # https://cloud.google.com/bigquery/docs/labels-intro # https://cloud.google.com/bigquery/docs/tables#table_naming {destination}, {self._cache_table_settings['_cached_tables']} => {is_should_work} # FIXME: do this without writing everything in SQL {preamble_sql}                     create or replace {entity_to_create} `{destination}`                     {partitioning_sql} as (                         {sql}                     ) `force_overwrite` means "force cache overwrite" (maybe change name later) # FIXME: 0a5cafd7878d3d0a