base.py

import json
import os
import random
import shutil
import tempfile
import time
import traceback
import unittest
from contextlib import contextmanager
from datetime import datetime
from functools import wraps

import pytest
import yaml
from unittest.mock import patch

import dbt.main as dbt
import dbt.flags as flags
from dbt.adapters.factory import get_adapter, reset_adapters, register_adapter
from dbt.clients.jinja import template_cache
from dbt.config import RuntimeConfig
from dbt.context import providers
from dbt.logger import GLOBAL_LOGGER as logger, log_manager


INITIAL_ROOT = os.getcwd()


def normalize(path):
    """On windows, neither is enough on its own:

    >>> normcase('C:\\documents/ALL CAPS/subdir\\..')
    'c:\\documents\\all caps\\subdir\\..'
    >>> normpath('C:\\documents/ALL CAPS/subdir\\..')
    'C:\\documents\\ALL CAPS'
    >>> normpath(normcase('C:\\documents/ALL CAPS/subdir\\..'))
    'c:\\documents\\all caps'
    """
    return os.path.normcase(os.path.normpath(path))


class Normalized:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f'Normalized({self.value!r})'

    def __str__(self):
        return f'Normalized({self.value!s})'

    def __eq__(self, other):
        return normalize(self.value) == normalize(other)


class FakeArgs:
    def __init__(self):
        self.threads = 1
        self.data = False
        self.schema = True
        self.full_refresh = False
        self.models = None
        self.exclude = None
        self.single_threaded = False


class TestArgs:
    def __init__(self, kwargs):
        self.which = 'run'
        self.single_threaded = False
        self.profiles_dir = None
        self.__dict__.update(kwargs)


def _profile_from_test_name(test_name):
    adapter_names = ('postgres', 'snowflake', 'redshift', 'bigquery', 'presto')
    adapters_in_name = sum(x in test_name for x in adapter_names)
    if adapters_in_name != 1:
        raise ValueError(
            'test names must have exactly 1 profile choice embedded, {} has {}'
            .format(test_name, adapters_in_name)
        )

    for adapter_name in adapter_names:
        if adapter_name in test_name:
            return adapter_name

    raise ValueError(
        'could not find adapter name in test name {}'.format(test_name)
    )


def _pytest_test_name():
    return os.environ['PYTEST_CURRENT_TEST'].split()[0]


def _pytest_get_test_root():
    test_path = _pytest_test_name().split('::')[0]
    relative_to = INITIAL_ROOT
    head = os.path.relpath(test_path, relative_to)

    path_parts = []
    while head:
        head, tail = os.path.split(head)
        path_parts.append(tail)
    path_parts.reverse()
    # dbt tests are all of the form 'test/integration/XXX_suite_name'
    target = os.path.join(*path_parts[:3])
    return os.path.join(relative_to, target)


def _really_makedirs(path):
    while not os.path.exists(path):
        try:
            os.makedirs(path)
        except EnvironmentError:
            raise


class DBTIntegrationTest(unittest.TestCase):
    CREATE_SCHEMA_STATEMENT = 'CREATE SCHEMA {}'
    DROP_SCHEMA_STATEMENT = 'DROP SCHEMA IF EXISTS {} CASCADE'

    prefix = "test{}{:04}".format(int(time.time()), random.randint(0, 9999))
    setup_alternate_db = False

    @property
    def database_host(self):
        if os.name == 'nt':
            return 'localhost'
        return 'database'

    def postgres_profile(self):
        return {
            'config': {
                'send_anonymous_usage_stats': False
            },
            'test': {
                'outputs': {
                    'default2': {
                        'type': 'postgres',
                        'threads': 4,
                        'host': self.database_host,
                        'port': 5432,
                        'user': 'root',
                        'pass': 'password',
                        'dbname': 'dbt',
                        'schema': self.unique_schema()
                    },
                    'noaccess': {
                        'type': 'postgres',
                        'threads': 4,
                        'host': self.database_host,
                        'port': 5432,
                        'user': 'noaccess',
                        'pass': 'password',
                        'dbname': 'dbt',
                        'schema': self.unique_schema()
                    }
                },
                'target': 'default2'
            }
        }

    def redshift_profile(self):
        return {
            'config': {
                'send_anonymous_usage_stats': False
            },
            'test': {
                'outputs': {
                    'default2': {
                        'type': 'redshift',
                        'threads': 1,
                        'host': os.getenv('REDSHIFT_TEST_HOST'),
                        'port': os.getenv('REDSHIFT_TEST_PORT'),
                        'user': os.getenv('REDSHIFT_TEST_USER'),
                        'pass': os.getenv('REDSHIFT_TEST_PASS'),
                        'dbname': os.getenv('REDSHIFT_TEST_DBNAME'),
                        'schema': self.unique_schema()
                    }
                },
                'target': 'default2'
            }
        }

    def snowflake_profile(self):
        return {
            'config': {
                'send_anonymous_usage_stats': False
            },
            'test': {
                'outputs': {
                    'default2': {
                        'type': 'snowflake',
                        'threads': 4,
                        'account': os.getenv('SNOWFLAKE_TEST_ACCOUNT'),
                        'user': os.getenv('SNOWFLAKE_TEST_USER'),
                        'password': os.getenv('SNOWFLAKE_TEST_PASSWORD'),
                        'database': os.getenv('SNOWFLAKE_TEST_DATABASE'),
                        'schema': self.unique_schema(),
                        'warehouse': os.getenv('SNOWFLAKE_TEST_WAREHOUSE'),
                    },
                    'noaccess': {
                        'type': 'snowflake',
                        'threads': 4,
                        'account': os.getenv('SNOWFLAKE_TEST_ACCOUNT'),
                        'user': 'noaccess',
                        'password': 'password',
                        'database': os.getenv('SNOWFLAKE_TEST_DATABASE'),
                        'schema': self.unique_schema(),
                        'warehouse': os.getenv('SNOWFLAKE_TEST_WAREHOUSE'),
                    }
                },
                'target': 'default2'
            }
        }

    def bigquery_profile(self):
        credentials_json_str = os.getenv('BIGQUERY_SERVICE_ACCOUNT_JSON').replace("'", '')
        credentials = json.loads(credentials_json_str)
        project_id = credentials.get('project_id')

        return {
            'config': {
                'send_anonymous_usage_stats': False
            },
            'test': {
                'outputs': {
                    'default2': {
                        'type': 'bigquery',
                        'method': 'service-account-json',
                        'threads': 1,
                        'project': project_id,
                        'keyfile_json': credentials,
                        'schema': self.unique_schema(),
                    },
                },
                'target': 'default2'
            }
        }

    def presto_profile(self):
        return {
            'config': {
                'send_anonymous_usage_stats': False
            },
            'test': {
                'outputs': {
                    'default2': {
                        'type': 'presto',
                        'method': 'none',
                        'threads': 1,
                        'schema': self.unique_schema(),
                        'database': 'hive',
                        'host': 'presto',
                        'port': 8080,
                    },
                },
                'target': 'default2'
            }
        }

    @property
    def packages_config(self):
        return None

    def unique_schema(self):
        schema = self.schema

        to_return = "{}_{}".format(self.prefix, schema)

        if self.adapter_type == 'snowflake':
            return to_return.upper()

        return to_return.lower()

    @property
    def default_database(self):
        database = self.config.credentials.database
        if self.adapter_type == 'snowflake':
            return database.upper()
        return database

    @property
    def alternative_database(self):
        if self.adapter_type == 'bigquery':
            return os.environ['BIGQUERY_TEST_ALT_DATABASE']
        elif self.adapter_type == 'snowflake':
            return os.environ['SNOWFLAKE_TEST_ALT_DATABASE']
        return None

    def get_profile(self, adapter_type):
        if adapter_type == 'postgres':
            return self.postgres_profile()
        elif adapter_type == 'snowflake':
            return self.snowflake_profile()
        elif adapter_type == 'bigquery':
            return self.bigquery_profile()
        elif adapter_type == 'redshift':
            return self.redshift_profile()
        elif adapter_type == 'presto':
            return self.presto_profile()
        else:
            raise ValueError('invalid adapter type {}'.format(adapter_type))

    def _pick_profile(self):
        test_name = self.id().split('.')[-1]
        return _profile_from_test_name(test_name)

    def _symlink_test_folders(self):
        for entry in os.listdir(self.test_original_source_path):
            src = os.path.join(self.test_original_source_path, entry)
            tst = os.path.join(self.test_root_dir, entry)
            if os.path.isdir(src) or src.endswith('.sql'):
                # symlink all sql files and all directories.
                os.symlink(src, tst)
        os.symlink(self._logs_dir, os.path.join(self.test_root_dir, 'logs'))

    def setUp(self):
        self.dbt_core_install_root = os.path.dirname(dbt.__file__)
        log_manager.reset_handlers()
        self.initial_dir = INITIAL_ROOT
        os.chdir(self.initial_dir)
        # before we go anywhere, collect the initial path info
        self._logs_dir = os.path.join(self.initial_dir, 'logs', self.prefix)
        _really_makedirs(self._logs_dir)
        self.test_original_source_path = _pytest_get_test_root()
        self.test_root_dir = normalize(tempfile.mkdtemp(prefix='dbt-int-test-'))
        os.chdir(self.test_root_dir)
        try:
            self._symlink_test_folders()
        except Exception as exc:
            msg = '\n\t'.join((
                'Failed to symlink test folders!',
                'initial_dir={0.initial_dir}',
                'test_original_source_path={0.test_original_source_path}',
                'test_root_dir={0.test_root_dir}'
            )).format(self)
            logger.exception(msg)

            # if logging isn't set up, I still really want this message.
            print(msg)
            traceback.print_exc()

            raise

        self._created_schemas = set()
        flags.reset()
        template_cache.clear()

        self.use_profile(self._pick_profile())
        self.use_default_project()
        self.set_packages()
        self.load_config()

    def use_default_project(self, overrides=None):
        # create a dbt_project.yml
        base_project_config = {
            'name': 'test',
            'version': '1.0',
            'test-paths': [],
            'source-paths': [self.models],
            'profile': 'test',
        }

        project_config = {}
        project_config.update(base_project_config)
        project_config.update(self.project_config)
        project_config.update(overrides or {})

        with open("dbt_project.yml", 'w') as f:
            yaml.safe_dump(project_config, f, default_flow_style=True)

    def use_profile(self, adapter_type):
        self.adapter_type = adapter_type

        profile_config = {}
        default_profile_config = self.get_profile(adapter_type)

        profile_config.update(default_profile_config)
        profile_config.update(self.profile_config)

        if not os.path.exists(self.test_root_dir):
            os.makedirs(self.test_root_dir)

        profiles_path = os.path.join(self.test_root_dir, 'profiles.yml')
        with open(profiles_path, 'w') as f:
            yaml.safe_dump(profile_config, f, default_flow_style=True)
        self._profile_config = profile_config

    def set_packages(self):
        if self.packages_config is not None:
            with open('packages.yml', 'w') as f:
                yaml.safe_dump(self.packages_config, f, default_flow_style=True)

    def load_config(self):
        # we've written our profile and project. Now we want to instantiate a
        # fresh adapter for the tests.
        # it's important to use a different connection handle here so
        # we don't look into an incomplete transaction
        kwargs = {
            'profile': None,
            'profiles_dir': self.test_root_dir,
            'target': None,
        }

        config = RuntimeConfig.from_args(TestArgs(kwargs))

        register_adapter(config)
        adapter = get_adapter(config)
        adapter.cleanup_connections()
        self.adapter_type = adapter.type()
        self.adapter = adapter
        self.config = config

        self._drop_schemas()
        self._create_schemas()

    def quote_as_configured(self, value, quote_key):
        return self.adapter.quote_as_configured(value, quote_key)

    def tearDown(self):
        # get any current run adapter and clean up its connections before we
        # reset them. It'll probably be different from ours because
        # handle_and_check() calls reset_adapters().
        register_adapter(self.config)
        adapter = get_adapter(self.config)
        if adapter is not self.adapter:
            adapter.cleanup_connections()
        if not hasattr(self, 'adapter'):
            self.adapter = adapter

        self._drop_schemas()

        self.adapter.cleanup_connections()
        reset_adapters()
        os.chdir(INITIAL_ROOT)
        try:
            shutil.rmtree(self.test_root_dir)
        except EnvironmentError:
            logger.exception('Could not clean up after test - {} not removable'
                             .format(self.test_root_dir))

    def _get_schema_fqn(self, database, schema):
        schema_fqn = self.quote_as_configured(schema, 'schema')
        if self.adapter_type == 'snowflake':
            database = self.quote_as_configured(database, 'database')
            schema_fqn = '{}.{}'.format(database, schema_fqn)
        return schema_fqn

    def _create_schema_named(self, database, schema):
        if self.adapter_type == 'bigquery':
            self.adapter.create_schema(database, schema)
        else:
            schema_fqn = self._get_schema_fqn(database, schema)
            self.run_sql(self.CREATE_SCHEMA_STATEMENT.format(schema_fqn))
            self._created_schemas.add(schema_fqn)

    def _drop_schema_named(self, database, schema):
        if self.adapter_type == 'bigquery' or self.adapter_type == 'presto':
            self.adapter.drop_schema(database, schema)
        else:
            schema_fqn = self._get_schema_fqn(database, schema)
            self.run_sql(self.DROP_SCHEMA_STATEMENT.format(schema_fqn))

    def _create_schemas(self):
        schema = self.unique_schema()
        with self.adapter.connection_named('__test'):
            self._create_schema_named(self.default_database, schema)
            if self.setup_alternate_db and self.adapter_type == 'snowflake':
                self._create_schema_named(self.alternative_database, schema)

    def _drop_schemas_adapter(self):
        schema = self.unique_schema()
        if self.adapter_type == 'bigquery' or self.adapter_type == 'presto':
            self._drop_schema_named(self.default_database, schema)
            if self.setup_alternate_db and self.alternative_database:
                self._drop_schema_named(self.alternative_database, schema)

    def _drop_schemas_sql(self):
        schema = self.unique_schema()
        # we always want to drop these if necessary, we'll clear it soon.
        self._created_schemas.add(
            self._get_schema_fqn(self.default_database, schema)
        )
        # on postgres/redshift, this will make you sad
        drop_alternative = (
            self.setup_alternate_db and
            self.adapter_type not in {'postgres', 'redshift'} and
            self.alternative_database
        )
        if drop_alternative:
            self._created_schemas.add(
                self._get_schema_fqn(self.alternative_database, schema)
            )

        for schema_fqn in self._created_schemas:
            self.run_sql(self.DROP_SCHEMA_STATEMENT.format(schema_fqn))

        self._created_schemas.clear()

    def _drop_schemas(self):
        with self.adapter.connection_named('__test'):
            if self.adapter_type == 'bigquery' or self.adapter_type == 'presto':
                self._drop_schemas_adapter()
            else:
                self._drop_schemas_sql()

    @property
    def project_config(self):
        return {}

    @property
    def profile_config(self):
        return {}

    def run_dbt(self, args=None, expect_pass=True, strict=True, parser=True, profiles_dir=True):
        res, success = self.run_dbt_and_check(args=args, strict=strict, parser=parser, profiles_dir=profiles_dir)

        self.assertEqual(
            success, expect_pass,
            "dbt exit state did not match expected")

        return res

    def run_dbt_and_check(self, args=None, strict=True, parser=False, profiles_dir=True):
        log_manager.reset_handlers()
        if args is None:
            args = ["run"]

        final_args = []

        if strict:
            final_args.append('--strict')
        if parser:
            final_args.append('--test-new-parser')
        if os.getenv('DBT_TEST_SINGLE_THREADED') in ('y', 'Y', '1'):
            final_args.append('--single-threaded')

        final_args.extend(args)

        if profiles_dir:
            final_args.extend(['--profiles-dir', self.test_root_dir])
        final_args.append('--log-cache-events')

        logger.info("Invoking dbt with {}".format(final_args))
        return dbt.handle_and_check(final_args)

    def run_sql_file(self, path, kwargs=None):
        with open(path, 'r') as f:
            statements = f.read().split(";")
            for statement in statements:
                self.run_sql(statement, kwargs=kwargs)

    # horrible hack to support snowflake for right now
    def transform_sql(self, query, kwargs=None):
        to_return = query

        if self.adapter_type == 'snowflake':
            to_return = to_return.replace("BIGSERIAL", "BIGINT AUTOINCREMENT")

        base_kwargs = {
            'schema': self.unique_schema(),
            'database': self.adapter.quote(self.default_database),
        }
        if kwargs is None:
            kwargs = {}
        base_kwargs.update(kwargs)


        to_return = to_return.format(**base_kwargs)

        return to_return

    def run_sql_bigquery(self, sql, fetch):
        """Run an SQL query on a bigquery adapter. No cursors, transactions,
        etc. to worry about"""

        do_fetch = fetch != 'None'
        _, res = self.adapter.execute(sql, fetch=do_fetch)

        # convert dataframe to matrix-ish repr
        if fetch == 'one':
            return res[0]
        else:
            return list(res)

    def run_sql_presto(self, sql, fetch, conn):
        cursor = conn.handle.cursor()
        try:
            cursor.execute(sql)
            if fetch == 'one':
                return cursor.fetchall()[0]
            elif fetch == 'all':
                return cursor.fetchall()
            else:
                # we have to fetch.
                cursor.fetchall()
        except Exception as e:
            conn.handle.rollback()
            conn.transaction_open = False
            print(sql)
            print(e)
            raise
        else:
            conn.handle.commit()
            conn.transaction_open = False

    def run_sql_common(self, sql, fetch, conn):
        with conn.handle.cursor() as cursor:
            try:
                cursor.execute(sql)
                conn.handle.commit()
                if fetch == 'one':
                    return cursor.fetchone()
                elif fetch == 'all':
                    return cursor.fetchall()
                else:
                    return
            except BaseException as e:
                if conn.handle and not conn.handle.closed:
                    conn.handle.rollback()
                print(sql)
                print(e)
                raise
            finally:
                conn.transaction_open = False

    def run_sql(self, query, fetch='None', kwargs=None, connection_name=None):
        if connection_name is None:
            connection_name = '__test'

        if query.strip() == "":
            return

        sql = self.transform_sql(query, kwargs=kwargs)

        with self.get_connection(connection_name) as conn:
            logger.debug('test connection "{}" executing: {}'.format(conn.name, sql))
            if self.adapter_type == 'bigquery':
                return self.run_sql_bigquery(sql, fetch)
            elif self.adapter_type == 'presto':
                return self.run_sql_presto(sql, fetch, conn)
            else:
                return self.run_sql_common(sql, fetch, conn)

    def _ilike(self, target, value):
        # presto has this regex substitution monstrosity instead of 'ilike'
        if self.adapter_type == 'presto':
            return r"regexp_like({}, '(?i)\A{}\Z')".format(target, value)
        else:
            return "{} ilike '{}'".format(target, value)

    def get_many_table_columns_snowflake(self, tables, schema, database=None):
        tables = set(tables)
        if database is None:
            database = self.default_database
        sql = 'show columns in schema {database}.{schema}'.format(
            database=self.quote_as_configured(database, 'database'),
            schema=self.quote_as_configured(schema, 'schema')
        )
        # assumption: this will be much  faster than doing one query/table
        # because in tests, we'll want most of our tables most of the time.
        columns = self.run_sql(sql, fetch='all')
        results = []
        for column in columns:
            table_name, _, column_name, json_data_type = column[:4]
            character_maximum_length = None
            if table_name in tables:
                typeinfo = json.loads(json_data_type)
                data_type = typeinfo['type']
                if data_type == 'TEXT':
                    character_maximum_length = max(typeinfo['length'], 16777216)
                results.append((table_name, column_name, data_type, character_maximum_length))
        return results

    def get_many_table_columns_information_schema(self, tables, schema, database=None):
        if self.adapter_type == 'presto':
            columns = 'table_name, column_name, data_type'
        else:
            columns = 'table_name, column_name, data_type, character_maximum_length'

        sql = """
                select {columns}
                from {db_string}information_schema.columns
                where {schema_filter}
                  and ({table_filter})
                order by column_name asc"""

        db_string = ''
        if database:
            db_string = self.quote_as_configured(database, 'database') + '.'

        table_filters_s = " OR ".join(
            self._ilike('table_name', table.replace('"', ''))
            for table in tables
        )
        schema_filter = self._ilike('table_schema', schema)

        sql = sql.format(
                columns=columns,
                schema_filter=schema_filter,
                table_filter=table_filters_s,
                db_string=db_string)

        columns = self.run_sql(sql, fetch='all')
        return list(map(self.filter_many_columns, columns))

    def get_many_table_columns_bigquery(self, tables, schema, database=None):
        result = []
        for table in tables:
            relation = self._make_relation(table, schema, database)
            columns = self.adapter.get_columns_in_relation(relation)
            for col in columns:
                result.append((table, col.column, col.dtype, col.char_size))
        return result

    def get_many_table_columns(self, tables, schema, database=None):
        if self.adapter_type == 'snowflake':
            result = self.get_many_table_columns_snowflake(tables, schema, database)
        elif self.adapter_type == 'bigquery':
            result = self.get_many_table_columns_bigquery(tables, schema, database)
        else:
            result = self.get_many_table_columns_information_schema(tables, schema, database)
        result.sort(key=lambda x: '{}.{}'.format(x[0], x[1]))
        return result

    def filter_many_columns(self, column):
        if len(column) == 3:
            table_name, column_name, data_type = column
            char_size = None
        else:
            table_name, column_name, data_type, char_size = column
        # in snowflake, all varchar widths are created equal
        if self.adapter_type == 'snowflake':
            if char_size and char_size < 16777216:
                char_size = 16777216
        return (table_name, column_name, data_type, char_size)

    @contextmanager
    def get_connection(self, name=None):
        """Create a test connection context where all executed macros, etc will
        get self.adapter as the adapter.

        This allows tests to run normal adapter macros as if reset_adapters()
        were not called by handle_and_check (for asserts, etc)
        """
        if name is None:
            name = '__test'
        with patch.object(providers, 'get_adapter', return_value=self.adapter):
            with self.adapter.connection_named(name):
                conn = self.adapter.connections.get_thread_connection()
                yield conn

    def get_relation_columns(self, relation):
        with self.get_connection():
            columns = self.adapter.get_columns_in_relation(relation)

        return sorted(((c.name, c.dtype, c.char_size) for c in columns),
                      key=lambda x: x[0])

    def get_table_columns(self, table, schema=None, database=None):
        schema = self.unique_schema() if schema is None else schema
        database = self.default_database if database is None else database
        relation = self.adapter.Relation.create(
            database=database,
            schema=schema,
            identifier=table,
            type='table',
            quote_policy=self.config.quoting
        )
        return self.get_relation_columns(relation)

    def get_table_columns_as_dict(self, tables, schema=None):
        col_matrix = self.get_many_table_columns(tables, schema)
        res = {}
        for row in col_matrix:
            table_name = row[0]
            col_def = row[1:]
            if table_name not in res:
                res[table_name] = []
            res[table_name].append(col_def)
        return res

    def get_models_in_schema_snowflake(self, schema):
        sql = 'show objects in schema {}.{}'.format(
            self.quote_as_configured(self.default_database, 'database'),
            self.quote_as_configured(schema, 'schema')
        )
        results = {}
        for row in self.run_sql(sql, fetch='all'):
            # I sure hope these never change!
            name = row[1]
            kind = row[4]

            if kind == 'TABLE':
                kind = 'table'
            elif kind == 'VIEW':
                kind = 'view'

            results[name] = kind
        return results

    def get_models_in_schema(self, schema=None):
        schema = self.unique_schema() if schema is None else schema
        if self.adapter_type == 'snowflake':
            return self.get_models_in_schema_snowflake(schema)

        sql = """
                select table_name,
                        case when table_type = 'BASE TABLE' then 'table'
                             when table_type = 'VIEW' then 'view'
                             else table_type
                        end as materialization
                from information_schema.tables
                where {}
                order by table_name
                """

        sql = sql.format(self._ilike('table_schema', schema))
        result = self.run_sql(sql, fetch='all')

        return {model_name: materialization for (model_name, materialization) in result}

    def _assertTablesEqualSql(self, relation_a, relation_b, columns=None):
        if columns is None:
            columns = self.get_relation_columns(relation_a)

        columns_csv = ', '.join([self.adapter.quote(record[0]) for record in columns])

        if self.adapter_type == 'bigquery':
            except_operator = 'EXCEPT DISTINCT'
        else:
            except_operator = 'EXCEPT'

        sql = """
            with diff_count as (
                SELECT
                    1 as id,
                    COUNT(*) as num_missing FROM (
                        (SELECT {columns} FROM {relation_a} {except_op}
                         SELECT {columns} FROM {relation_b})
                         UNION ALL
                        (SELECT {columns} FROM {relation_b} {except_op}
                         SELECT {columns} FROM {relation_a})
                    ) as a
            ), table_a as (
                SELECT COUNT(*) as num_rows FROM {relation_a}
            ), table_b as (
                SELECT COUNT(*) as num_rows FROM {relation_b}
            ), row_count_diff as (
                select
                    1 as id,
                    table_a.num_rows - table_b.num_rows as difference
                from table_a, table_b
            )
            select
                row_count_diff.difference as row_count_difference,
                diff_count.num_missing as num_mismatched
            from row_count_diff
            join diff_count using (id)
            """.strip().format(
                columns=columns_csv,
                relation_a=str(relation_a),
                relation_b=str(relation_b),
                except_op=except_operator
            )

        return sql

    def assertTablesEqual(self, table_a, table_b,
                          table_a_schema=None, table_b_schema=None,
                          table_a_db=None, table_b_db=None):
        if table_a_schema is None:
            table_a_schema = self.unique_schema()

        if table_b_schema is None:
            table_b_schema = self.unique_schema()

        if table_a_db is None:
            table_a_db = self.default_database

        if table_b_db is None:
            table_b_db = self.default_database

        relation_a = self._make_relation(table_a, table_a_schema, table_a_db)
        relation_b = self._make_relation(table_b, table_b_schema, table_b_db)

        self._assertTableColumnsEqual(relation_a, relation_b)

        sql = self._assertTablesEqualSql(relation_a, relation_b)
        result = self.run_sql(sql, fetch='one')

        self.assertEqual(
            result[0],
            0,
            'row_count_difference nonzero: ' + sql
        )
        self.assertEqual(
            result[1],
            0,
            'num_mismatched nonzero: ' + sql
        )

    def _make_relation(self, identifier, schema=None, database=None):
        if schema is None:
            schema = self.unique_schema()
        if database is None:
            database = self.default_database
        return self.adapter.Relation.create(
            database=database,
            schema=schema,
            identifier=identifier,
            quote_policy=self.config.quoting
        )

    def get_many_relation_columns(self, relations):
        """Returns a dict of (datbase, schema) -> (dict of (table_name -> list of columns))
        """
        schema_fqns = {}
        for rel in relations:
            this_schema = schema_fqns.setdefault((rel.database, rel.schema), [])
            this_schema.append(rel.identifier)

        column_specs = {}
        for key, tables in schema_fqns.items():
            database, schema = key
            columns = self.get_many_table_columns(tables, schema, database=database)
            table_columns = {}
            for col in columns:
                table_columns.setdefault(col[0], []).append(col[1:])
            for rel_name, columns in table_columns.items():
                key = (database, schema, rel_name)
                column_specs[key] = columns

        return column_specs

    def assertManyRelationsEqual(self, relations, default_schema=None, default_database=None):
        if default_schema is None:
            default_schema = self.unique_schema()
        if default_database is None:
            default_database = self.default_database

        specs = []
        for relation in relations:
            if not isinstance(relation, (tuple, list)):
                relation = [relation]

            assert len(relation) <= 3

            if len(relation) == 3:
                relation = self._make_relation(*relation)
            elif len(relation) == 2:
                relation = self._make_relation(relation[0], relation[1], default_database)
            elif len(relation) == 1:
                relation = self._make_relation(relation[0], default_schema, default_database)
            else:
                raise ValueError('relation must be a sequence of 1, 2, or 3 values')

            specs.append(relation)

        with self.get_connection():
            column_specs = self.get_many_relation_columns(specs)

        # make sure everyone has equal column definitions
        first_columns = None
        for relation in specs:
            key = (relation.database, relation.schema, relation.identifier)
            columns = column_specs[key]
            if first_columns is None:
                first_columns = columns
            else:
                self.assertEqual(
                    first_columns, columns,
                    '{} did not match {}'.format(str(specs[0]), str(relation))
                )

        # make sure everyone has the same data. if we got here, everyone had
        # the same column specs!
        first_relation = None
        for relation in specs:
            if first_relation is None:
                first_relation = relation
            else:
                sql = self._assertTablesEqualSql(first_relation, relation,
                                                 columns=first_columns)
                result = self.run_sql(sql, fetch='one')

                self.assertEqual(
                    result[0],
                    0,
                    'row_count_difference nonzero: ' + sql
                )
                self.assertEqual(
                    result[1],
                    0,
                    'num_mismatched nonzero: ' + sql
                )

    def assertManyTablesEqual(self, *args):
        schema = self.unique_schema()

        all_tables = []
        for table_equivalencies in args:
            all_tables += list(table_equivalencies)

        all_cols = self.get_table_columns_as_dict(all_tables, schema)

        for table_equivalencies in args:
            first_table = table_equivalencies[0]
            first_relation = self._make_relation(first_table)

            # assert that all tables have the same columns
            base_result = all_cols[first_table]
            self.assertTrue(len(base_result) > 0)

            for other_table in table_equivalencies[1:]:
                other_result = all_cols[other_table]
                self.assertTrue(len(other_result) > 0)
                self.assertEqual(base_result, other_result)

                other_relation = self._make_relation(other_table)
                sql = self._assertTablesEqualSql(first_relation,
                                                 other_relation,
                                                 columns=base_result)
                result = self.run_sql(sql, fetch='one')

                self.assertEqual(
                    result[0],
                    0,
                    'row_count_difference nonzero: ' + sql
                )
                self.assertEqual(
                    result[1],
                    0,
                    'num_mismatched nonzero: ' + sql
                )


    def _assertTableRowCountsEqual(self, relation_a, relation_b):
        cmp_query = """
            with table_a as (

                select count(*) as num_rows from {}

            ), table_b as (

                select count(*) as num_rows from {}

            )

            select table_a.num_rows - table_b.num_rows as difference
            from table_a, table_b

        """.format(str(relation_a), str(relation_b))


        res = self.run_sql(cmp_query, fetch='one')

        self.assertEqual(int(res[0]), 0, "Row count of table {} doesn't match row count of table {}. ({} rows different)".format(
                relation_a.identifier,
                relation_b.identifier,
                res[0]
            )
        )

    def assertTableDoesNotExist(self, table, schema=None):
        columns = self.get_table_columns(table, schema)

        self.assertEqual(
            len(columns),
            0
        )

    def assertTableDoesExist(self, table, schema=None):
        columns = self.get_table_columns(table, schema)

        self.assertGreater(
            len(columns),
            0
        )

    def _assertTableColumnsEqual(self, relation_a, relation_b):
        table_a_result = self.get_relation_columns(relation_a)
        table_b_result = self.get_relation_columns(relation_b)

        text_types = {'text', 'character varying', 'character', 'varchar'}

        self.assertEqual(len(table_a_result), len(table_b_result))
        for a_column, b_column in zip(table_a_result, table_b_result):
            a_name, a_type, a_size = a_column
            b_name, b_type, b_size = b_column
            self.assertEqual(a_name, b_name,
                '{} vs {}: column "{}" != "{}"'.format(
                    relation_a, relation_b, a_name, b_name
                ))

            self.assertEqual(a_type, b_type,
                '{} vs {}: column "{}" has type "{}" != "{}"'.format(
                    relation_a, relation_b, a_name, a_type, b_type
                ))

            if self.adapter_type == 'presto' and None in (a_size, b_size):
                # None is compatible with any size
                continue

            self.assertEqual(a_size, b_size,
                '{} vs {}: column "{}" has size "{}" != "{}"'.format(
                    relation_a, relation_b, a_name, a_size, b_size
                ))

    def assertEquals(self, *args, **kwargs):
        # assertEquals is deprecated. This makes the warnings less chatty
        self.assertEqual(*args, **kwargs)

    def assertBetween(self, timestr, start, end=None):
        datefmt = '%Y-%m-%dT%H:%M:%S.%fZ'
        if end is None:
            end = datetime.utcnow()

        parsed = datetime.strptime(timestr, datefmt)

        self.assertLessEqual(start, parsed,
            'parsed date {} happened before {}'.format(
                parsed,
                start.strftime(datefmt))
        )
        self.assertGreaterEqual(end, parsed,
            'parsed date {} happened after {}'.format(
                parsed,
                end.strftime(datefmt))
        )


def use_profile(profile_name):
    """A decorator to declare a test method as using a particular profile.
    Handles both setting the nose attr and calling self.use_profile.

    Use like this:

    class TestSomething(DBIntegrationTest):
        @use_profile('postgres')
        def test_postgres_thing(self):
            self.assertEqual(self.adapter_type, 'postgres')

        @use_profile('snowflake')
        def test_snowflake_thing(self):
            self.assertEqual(self.adapter_type, 'snowflake')
    """
    def outer(wrapped):
        @getattr(pytest.mark, 'profile_'+profile_name)
        @wraps(wrapped)
        def func(self, *args, **kwargs):
            return wrapped(self, *args, **kwargs)
        # sanity check at import time
        assert _profile_from_test_name(wrapped.__name__) == profile_name
        return func
    return outer


class AnyFloat:
    """Any float. Use this in assertEqual() calls to assert that it is a float.
    """
    def __eq__(self, other):
        return isinstance(other, float)


class AnyString:
    """Any string. Use this in assertEqual() calls to assert that it is a string.
    """
    def __eq__(self, other):
        return isinstance(other, str)


class AnyStringWith:
    def __init__(self, contains=None):
        self.contains = contains

    def __eq__(self, other):
        if not isinstance(other, str):
            return False

        if self.contains is None:
            return True

        return self.contains in other

    def __repr__(self):
        return 'AnyStringWith<{!r}>'.format(self.contains)