# Copyright Contributors to the Amundsen project.
# SPDX-License-Identifier: Apache-2.0

import base64
import tempfile
import unittest
from typing import Any

from mock import Mock, patch
from pyhocon import ConfigFactory

from databuilder import Scoped
from databuilder.extractor.bigquery_usage_extractor import BigQueryTableUsageExtractor, TableColumnUsageTuple

CORRECT_DATA = {
    "entries": [{
        "protoPayload": {
            "@type": "type.googleapis.com/google.cloud.audit.AuditLog",
            "status": {},
            "authenticationInfo": {
                "principalEmail": "your-user-here@test.com"
            },
            "serviceName": "bigquery.googleapis.com",
            "methodName": "jobservice.jobcompleted",
            "resourceName": "projects/your-project-here/jobs/bquxjob_758c08d1_16a96889839",
            "serviceData": {
                "@type": "type.googleapis.com/google.cloud.bigquery.logging.v1.AuditData",
                "jobCompletedEvent": {
                    "eventName": "query_job_completed",
                    "job": {
                        "jobName": {
                            "projectId": "your-project-here",
                            "jobId": "bquxjob_758c08d1_16a96889839",
                            "location": "US"
                        },
                        "jobConfiguration": {
                            "query": {
                                "query": "select descript from "
                                         "`bigquery-public-data.austin_incidents.incidents_2008`\n",
                                "destinationTable": {
                                    "projectId": "your-project-here",
                                    "datasetId": "_07147a061ddfd6dcaf246cfc5e858a0ccefa7080",
                                    "tableId": "anon1dd83635c62357091e55a5f76fb62d7deebcfa4c"
                                },
                                "createDisposition": "CREATE_IF_NEEDED",
                                "writeDisposition": "WRITE_TRUNCATE",
                                "defaultDataset": {},
                                "queryPriority": "QUERY_INTERACTIVE",
                                "statementType": "SELECT"
                            }
                        },
                        "jobStatus": {
                            "state": "DONE",
                            "error": {}
                        },
                        "jobStatistics": {
                            "createTime": "2019-05-08T08:22:56.349Z",
                            "startTime": "2019-05-08T08:22:56.660Z",
                            "endTime": "2019-05-08T08:23:00.049Z",
                            "totalProcessedBytes": "3637807",
                            "totalBilledBytes": "10485760",
                            "billingTier": 1,
                            "totalSlotMs": "452",
                            "referencedTables": [
                                {
                                    "projectId": "bigquery-public-data",
                                    "datasetId": "austin_incidents",
                                    "tableId": "incidents_2008"
                                }
                            ],
                            "totalTablesProcessed": 1,
                            "queryOutputRowCount": "179524"
                        }
                    }
                }
            }
        },
        "insertId": "-jyqvjse6lwjz",
        "resource": {
            "type": "bigquery_resource",
            "labels": {
                "project_id": "your-project-here"
            }
        },
        "timestamp": "2019-05-08T08:23:00.061Z",
        "severity": "INFO",
        "logName": "projects/your-project-here/logs/cloudaudit.googleapis.com%2Fdata_access",
        "receiveTimestamp": "2019-05-08T08:23:00.310709609Z"
    }]
}  # noqa

FAILURE = {
    "entries": [{
        "protoPayload": {
            "authenticationInfo": {
                "principalEmail": "your-user-here@test.com"
            },
            "methodName": "jobservice.jobcompleted",
            "serviceData": {
                "jobCompletedEvent": {
                    "job": {
                        "jobStatus": {
                            "state": "DONE",
                            "error": {
                                "code": 11,
                                "message": "Some descriptive error message"
                            }
                        },
                        "jobStatistics": {
                            "createTime": "2019-05-08T08:22:56.349Z",
                            "startTime": "2019-05-08T08:22:56.660Z",
                            "endTime": "2019-05-08T08:23:00.049Z",
                            "totalProcessedBytes": "3637807",
                            "totalBilledBytes": "10485760",
                            "referencedTables": [
                                {
                                    "projectId": "bigquery-public-data",
                                    "datasetId": "austin_incidents",
                                    "tableId": "incidents_2008"
                                }
                            ]
                        }
                    }
                }
            },
        },
    }]
}  # noqa

# An empty dict will be ignored, but putting in nextPageToken causes the test
# to loop infinitely, so we need a bogus key/value to ensure that we will try
# to read entries
NO_ENTRIES = {'key': 'value'}  # noqa

KEYFILE_DATA = """
ewogICJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsCiAgInByb2plY3RfaWQiOiAieW91ci1wcm9q
ZWN0LWhlcmUiLAogICJwcml2YXRlX2tleV9pZCI6ICJiMDQ0N2U1ODEyYTg5ZTAyOTgxYjRkMWE1
YjE1N2NlNzZkOWJlZTc3IiwKICAicHJpdmF0ZV9rZXkiOiAiLS0tLS1CRUdJTiBQUklWQVRFIEtF
WS0tLS0tXG5NSUlFdkFJQkFEQU5CZ2txaGtpRzl3MEJBUUVGQUFTQ0JLWXdnZ1NpQWdFQUFvSUJB
UUM1UzBYRWtHY2NuOEsxXG5ZbHhRbXlhRWFZK2grYnRacHRVWjJiK2J1cTluNExKU3I3eTdPQWll
ZjBWazIyQnc1TFRsUXRQSUtNVkh6MzJMXG5Ld0lJYmY5Wkwzamd5UC9hNHIveHVhMVdzNFF2YVkz
TGoxRG1ITm40L3hQNXdDY0VscHIxV2RXL05VZ1RQV1A2XG5LZnVDdHhyQTJxbHJNazhyYklXVTRm
WTAzQmFqdzNHT0p4VDBvbXlCVmdGSzJTdGRFUVVYMm9YQVdSNXJyR21qXG5qWTNzb3lNU0NwSWtT
b0h4b1BrVEM0VzZ2a3dJRlk4SUkwbmhsWUZHc3FiZjdkbTBLVEZmVVh5SUFTOHd6RCtlXG54UFVQ
V3k0UXA5cTVyNTVPRmlxdWt3TGNZei9BQXFpYTU3KzhURmhiWXcwUXNsZ2xSaWFLWkVhQyt4M0pD
OEhuXG5KajY2WE5mTEFnTUJBQUVDZ2dFQVMyNFlGYi9QS2ZqamM2RjZBUnBYNExsMFRqVHlqcmw2
c001UzBSdDdRbWRYXG5VSS9YM2NNZXh4NzZhZWRnYURUQ2F6MzhKdFJxRXlTbGI5enZNKzFMY013
QmdraHcxM05OUGlNZkxGZGg3VWNrXG5BUVR6b3VtRjFuWklkSGhEcWZ1QlUzWGhyTGdOQWtBUWpn
cy9KdVJSVU1iekJ2OXcrVFZ4WDcxbzAvWHdoWE5kXG5kSWlWdE1TbnFWQ0J2cEp3ZXBoR3FxNGQ3
VEIzb2F3UUg1QkFGeHk5NGpoT0dwaVFWYW8yQmtPdEVyVVBQYjkrXG5vRzByZTM3WHVtQzZRWENv
VSs4Zm4vcE1YVWVOUitXSm5tY1lndVZqWDl6QzJ3MU13cmVmOFVKa1Q4SHJxZ09KXG5sWnNFcVJr
aHBYUFVzdmt2dWxQTWQ3TitJdlFvYTh0N3ZaZFkrR1lMdVFLQmdRRHd2enY0alhVUStIU1RaVm1p
XG5hQmNMVGRMRE5WNlpuT25aTEhxbDZaQmloTUhZNi9qS2xDN1hqWGJaQ2NqS05MMkE1am9mQ0d5
bHFhNFRrZnArXG5rYmJKQ29KS2tFY1pSWGQ3NEdXb0J1V2d3enY2WWFkcDNxS2x0RndhM1FjMkJ3
SlNlazkrTzd6OGs2d0dvclZJXG5OK3ZNMVd3OWJPa1VaaXh4T2g2V2ZKSTl6UUtCZ1FERkNLQXZ2
b3FUQnErMnovazhLYy9lTHVRdThPWWNXVm9GXG55eXprOTN2QnBXcEVPT1hybnNsUFFtQldUdTN5
UWpRN08zd2t1c0g3VUtJQTg0MDVHbDlwbmJvTmlaSVdBRlpvXG4vVWlVVm5aa3pvZER5Tk9PUjBm
UW5zM1BaeE5peklSSjh2Mm93a2d3MExFYWEwaWUyNU92bFJmQ2pmYlVZL0EzXG5wbU9SVkdFVDl3
S0JnR0Zab3lHRjZoRzd0a0FvR28vT3NZclRwR2RsZkdSM2pDUlNsU0hrQ1l1ZERWbnZTY0o1XG5H
MXYwaTF1R1ZsaFY3VTlqU1p0azU3SXhvLytyNXZRcGJoVnJsM1laVTNiSG5XSk5RaTRvNDlBWFFu
aWo1bk9zXG5JRzhMT0xkd0swdFFtRUxMekx0SjRzanIyZ013NWtkV3ZaWXRzMEEvZXh6Um1DVU5F
SE5mMmk3OUFvR0FESVpkXG4yR3NlVi9aRzJUSWpQOFhRcHVrSUxFdTM5UGxoRlpreXcyTlFCS0ZG
UGd6MzRLQjVYNFp5cFVuaktsRTNETVRkXG5RV0IxMEVueDRtbVpBcFpBbG5BbVVaSDdMVmJjSjFS
aWRydUFUeXdwd1E5VkUyaElrbVJsNU5kQ2pqYzkrWTF1XG52bm1MS1Q4NjR0a0xCcjRpaHpqTkI5
c0tZN251blRzQWZVNkYxVVVDZ1lBMmdlMFdiVEVwRlBuN05YYjZ4citiXG5QK1RFVEVWZzhRS0Z1
OUtHVk03NXI5dmhYblNicmphbGVCSzJFQzBLK2F2d2hHTTd3eXRqM0FrTjRac2NKNWltXG5VZTBw
Z3pVSE1RSVI1OWlGVmt5WVVjZnZMSERZU0xmeW9QVU5RWWduVXBKYlZOczZtWFRqQ3o2UERrb0tX
ZzcyXG4rS3p4RWhubWJzY0NiSFRpQ08wNEtBPT1cbi0tLS0tRU5EIFBSSVZBVEUgS0VZLS0tLS1c
biIsCiAgImNsaWVudF9lbWFpbCI6ICJ0ZXN0LTE2MkB5b3VyLXByb2plY3QtaGVyZS5pYW0uZ3Nl
cnZpY2VhY2NvdW50LmNvbSIsCiAgImNsaWVudF9pZCI6ICIxMDg2NTMzMjY0MzE1NDU2ODg3MTAi
LAogICJhdXRoX3VyaSI6ICJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20vby9vYXV0aDIvYXV0
aCIsCiAgInRva2VuX3VyaSI6ICJodHRwczovL29hdXRoMi5nb29nbGVhcGlzLmNvbS90b2tlbiIs
CiAgImF1dGhfcHJvdmlkZXJfeDUwOV9jZXJ0X3VybCI6ICJodHRwczovL3d3dy5nb29nbGVhcGlz
LmNvbS9vYXV0aDIvdjEvY2VydHMiLAogICJjbGllbnRfeDUwOV9jZXJ0X3VybCI6ICJodHRwczov
L3d3dy5nb29nbGVhcGlzLmNvbS9yb2JvdC92MS9tZXRhZGF0YS94NTA5L3Rlc3QtMTYyJTQweW91
ci1wcm9qZWN0LWhlcmUuaWFtLmdzZXJ2aWNlYWNjb3VudC5jb20iCn0KCgo=
"""


class MockLoggingClient():
    def __init__(self, data: Any) -> None:
        self.data = data
        self.a = Mock()
        self.a.execute.return_value = self.data
        self.b = Mock()
        self.b.list.return_value = self.a

    def entries(self) -> Any:
        return self.b


# Patch fallback auth method to avoid actually calling google API
@patch('google.auth.default', lambda scopes: ['dummy', 'dummy'])
class TestBigqueryUsageExtractor(unittest.TestCase):

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_basic_extraction(self, mock_build: Any) -> None:
        """
        Test Extraction using mock class
        """
        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
        }
        conf = ConfigFactory.from_dict(config_dict)

        mock_build.return_value = MockLoggingClient(CORRECT_DATA)
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))
        result = extractor.extract()
        assert result is not None
        self.assertIsInstance(result, tuple)

        (key, value) = result
        self.assertIsInstance(key, TableColumnUsageTuple)
        self.assertIsInstance(value, int)

        self.assertEqual(key.database, 'bigquery')
        self.assertEqual(key.cluster, 'bigquery-public-data')
        self.assertEqual(key.schema, 'austin_incidents')
        self.assertEqual(key.table, 'incidents_2008')
        self.assertEqual(key.email, 'your-user-here@test.com')
        self.assertEqual(value, 1)

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_no_entries(self, mock_build: Any) -> None:
        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
        }
        conf = ConfigFactory.from_dict(config_dict)

        mock_build.return_value = MockLoggingClient(NO_ENTRIES)
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))
        result = extractor.extract()
        self.assertIsNone(result)

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_key_path(self, mock_build: Any) -> None:
        """
        Test key_path can be used
        """

        with tempfile.NamedTemporaryFile() as keyfile:
            # There are many github scanners looking for API / cloud keys, so in order not to get a
            # false positive triggering everywhere, I base64 encoded the key.
            # This is written to a tempfile as part of this test and then used.
            keyfile.write(base64.b64decode(KEYFILE_DATA))
            keyfile.flush()
            config_dict = {
                f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
                f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.KEY_PATH_KEY}': keyfile.name,
            }
            conf = ConfigFactory.from_dict(config_dict)

            mock_build.return_value = MockLoggingClient(CORRECT_DATA)
            extractor = BigQueryTableUsageExtractor()
            extractor.init(Scoped.get_scoped_conf(conf=conf,
                                                  scope=extractor.get_scope()))

            args, kwargs = mock_build.call_args
            creds = kwargs['http'].credentials
            self.assertEqual(creds.project_id, 'your-project-here')
            self.assertEqual(creds.service_account_email, 'test-162@your-project-here.iam.gserviceaccount.com')

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_timestamp_pagesize_settings(self, mock_build: Any) -> None:
        """
        Test timestamp and pagesize can be set
        """
        TIMESTAMP = '2019-01-01T00:00:00.00Z'
        PAGESIZE = 215

        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.TIMESTAMP_KEY}': TIMESTAMP,
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PAGE_SIZE_KEY}': PAGESIZE,
        }
        conf = ConfigFactory.from_dict(config_dict)

        client = MockLoggingClient(CORRECT_DATA)
        mock_build.return_value = client
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))

        args, kwargs = client.b.list.call_args
        body = kwargs['body']

        self.assertEqual(body['pageSize'], PAGESIZE)
        self.assertEqual(TIMESTAMP in body['filter'], True)

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_failed_jobs_should_not_be_counted(self, mock_build: Any) -> None:
        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
        }
        conf = ConfigFactory.from_dict(config_dict)

        client = MockLoggingClient(FAILURE)
        mock_build.return_value = client
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))

        result = extractor.extract()
        self.assertIsNone(result)

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_email_filter_not_counted(self, mock_build: Any) -> None:
        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.EMAIL_PATTERN}': 'emailFilter',
        }
        conf = ConfigFactory.from_dict(config_dict)

        mock_build.return_value = MockLoggingClient(CORRECT_DATA)
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))
        result = extractor.extract()
        self.assertIsNone(result)

    @patch('databuilder.extractor.base_bigquery_extractor.build')
    def test_email_filter_counted(self, mock_build: Any) -> None:
        config_dict = {
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}': 'your-project-here',
            f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.EMAIL_PATTERN}': '.*@test.com.*',
        }
        conf = ConfigFactory.from_dict(config_dict)

        mock_build.return_value = MockLoggingClient(CORRECT_DATA)
        extractor = BigQueryTableUsageExtractor()
        extractor.init(Scoped.get_scoped_conf(conf=conf,
                                              scope=extractor.get_scope()))
        result = extractor.extract()
        assert result is not None
        self.assertIsInstance(result, tuple)

        (key, value) = result
        self.assertIsInstance(key, TableColumnUsageTuple)
        self.assertIsInstance(value, int)

        self.assertEqual(key.database, 'bigquery')
        self.assertEqual(key.cluster, 'bigquery-public-data')
        self.assertEqual(key.schema, 'austin_incidents')
        self.assertEqual(key.table, 'incidents_2008')
        self.assertEqual(key.email, 'your-user-here@test.com')
        self.assertEqual(value, 1)
