from google.cloud import bigquery
import pandas as pd

# Initialize BigQuery client
client = bigquery.Client()

# Query to get all datasets
schemata_query = """
SELECT schema_name FROM region-us.INFORMATION_SCHEMA.SCHEMATA;
"""

# Query template to get table metadata for each dataset
dataset_tables_query = """
SELECT
    table_id,
    dataset_id,
    DATE(TIMESTAMP_MILLIS(last_modified_time)) AS last_modified_date,
    size_bytes / 1024 / 1024 / 1024 as size_gb,
    size_bytes / 1024 / 1024 / 1024 /1024 as size_tb
FROM `project-name.{dataset}.__TABLES__`;
"""

# Create DataFrame to store results
dwh_analysis_df = pd.DataFrame(columns=[
    'table_id', 'dataset_id', 'last_modified_date', 'size_gb', 'size_tb'
])

def fetch_schemata(schemata_query, client):
    query_job = client.query(schemata_query)
    return query_job.result()

def schemata_data_to_list(schemata_data):
    datasets = []
    for row in schemata_data:
        datasets.append(row[0])
    return datasets

def query_table_metadata(dataset_list, client, dataset_tables_query, analysis_df):
    for dataset in dataset_list:
        tables_query = dataset_tables_query.format(dataset=dataset)
        bq_df = client.query(tables_query).to_dataframe()
        print(f'Loading {dataset}')
        analysis_df = pd.concat([analysis_df, bq_df])
    return analysis_df.to_csv('output.csv', sep=',', index=False, encoding='utf-8')

# Main execution
schemata_data = fetch_schemata(schemata_query, client)
dataset_list = schemata_data_to_list(schemata_data)
query_table_metadata(dataset_list, client, dataset_tables_query, dwh_analysis_df)

CREATE OR REPLACE TABLE `project-name.analysis_dataset.table_activity` AS(
    SELECT
        dataset,
        COUNT(DISTINCT job_id) AS jobs_in_last_3_months,
        ARRAY_AGG(DISTINCT user_email) AS users_emails
    FROM (
        SELECT
            user_email,
            rt.dataset_id AS dataset,
            job_id
        FROM
            `project-name.region-us.INFORMATION_SCHEMA.JOBS_BY_ORGANIZATION`,
            UNNEST(referenced_tables) rt
        WHERE
            rt.project_id = 'project-name'
            AND creation_time >= TIMESTAMP(DATE_SUB(CURRENT_DATE(), INTERVAL 3 MONTH))
            AND rt.dataset_id NOT LIKE '%_script%'
    )
    GROUP BY 1
)

SELECT
    'project-name' AS project_id,
    COALESCE(ta.dataset, dwh.dataset_id) AS dataset_id,
    ta.jobs_in_last_3_months,
    CASE
        WHEN dataset_id IN ('data_monitoring','prod_analytics','ml','dimensions')
        THEN TRUE
        ELSE FALSE
    END AS should_have_activity,
    dataset_size_tb,
    STRUCT(
        dwh.tables AS tables,
        ta.users_emails AS users_emails
    ) AS dataset_struct
FROM
    dwh_analysis dwh
LEFT JOIN
    table_activity ta
ON
    ta.dataset = dwh.dataset_id

from google.cloud import bigquery
from google.api_core.exceptions import BadRequest, NotFound

# Initialize BigQuery client
client = bigquery.Client()

# Query to identify tables for deprecation
query = """
SELECT
    project_id,
    dataset_id,
    table_id
FROM
    `project-name.analysis_dataset.warehouse_activity_analysis`
WHERE
    should_have_activity IS FALSE
    AND table_id NOT LIKE '%_to_be_archived%'
"""

def check_table_type(table_dict):
    """Check if the given table ID corresponds to a TABLE type."""
    table_ref = client.dataset(table_dict['dataset_id']).table(table_dict['table_id'])
    table = client.get_table(table_ref)
    return table.table_type

def add_archive_suffix(table_dict, suffix='_to_be_archived'):
    """Add suffix to table name to identify it for archiving."""
    if check_table_type(table_dict) == "TABLE":
        table_id = f"{table_dict['project_id']}.{table_dict['dataset_id']}.{table_dict['table_id']}"

        # Create SQL to rename the table
        alter_table_query = f"ALTER TABLE IF EXISTS {table_id} RENAME TO {table_dict['table_id']}{suffix}"

        try:
            query_job = client.query(alter_table_query)
            query_job.result()  # Wait for the query to complete
            print(f"Successfully renamed {table_id}")
            return True
        except BadRequest as e:
            print(f"Cannot rename table {table_id}: {e}")
            return False
        except NotFound as e:
            print(f"Table {table_id} not found: {e}")
            return False
    else:
        print(f"{table_dict['table_id']} is not a table, skipping")
        return False

# Main execution
datasets_df = client.query(query).to_dataframe()
table_list = datasets_df.to_dict(orient='records')

for table in table_list:
    add_archive_suffix(table)

import time
from google.cloud import bigquery
from google.api_core.exceptions import NotFound

# Initialize BigQuery client
client = bigquery.Client()

# Query to get datasets with marked tables
fetch_datasets_query = """
SELECT
    dataset_id
FROM
    `project-name.analysis_dataset.warehouse_activity_analysis`,
    UNNEST(dataset_struct.tables) as dst
WHERE
    should_have_activity IS FALSE
    AND table_id LIKE '%_to_be_archived%'
GROUP BY 1
"""

def dataset_exists(client, dataset_id):
    """Check if dataset exists."""
    try:
        dataset_ref = client.dataset(dataset_id)
        dataset = client.get_dataset(dataset_ref)
        tables = list(client.list_tables(dataset))
        return True
    except Exception:
        return False

def delete_dataset(client, dataset_id):
    """Delete dataset and all its contents."""
    dataset_id = 'project-name' + '.' + dataset_id
    client.delete_dataset(
        dataset_id,
        delete_contents=True,
        not_found_ok=False
    )
    print(f"Deleted dataset {dataset_id}")

# Main execution
dataset_ids_list = client.query(fetch_datasets_query).to_dataframe()['dataset_id'].tolist()

for dataset in dataset_ids_list:
    if dataset_exists(client, dataset):
        # For safety, require manual confirmation before deletion
        user_input = input(f'Do you want to delete the dataset {dataset}? (Y/n): ')
        if user_input.lower() == "y" or not user_input.strip():
            print(f'Deleting {dataset}')
            delete_dataset(client, dataset)
        else:
            print(f'Skipping {dataset}')

print('All marked datasets have been processed')