extract_spotify_data.py

"""
Extract Coldplay data from Spotify API using DLT
"""
import os
import base64
import requests
from datetime import datetime
from typing import Dict, List, Any, Iterator

import dlt
from dlt.sources.rest_api import rest_api_source
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Spotify API configuration
SPOTIFY_CLIENT_ID = os.getenv("SPOTIFY_CLIENT_ID")
SPOTIFY_CLIENT_SECRET = os.getenv("SPOTIFY_CLIENT_SECRET")
SPOTIFY_API_BASE_URL = "https://api.spotify.com/v1"

# Coldplay's Spotify artist ID
COLDPLAY_ARTIST_ID = "4gzpq5DPGxSnKTe4SA8HAU"


def get_spotify_access_token() -> str:
    """Get Spotify access token using Client Credentials flow"""
    auth_url = "https://accounts.spotify.com/api/token"
    
    # Encode client credentials
    client_credentials = f"{SPOTIFY_CLIENT_ID}:{SPOTIFY_CLIENT_SECRET}"
    client_credentials_b64 = base64.b64encode(client_credentials.encode()).decode()
    
    headers = {
        "Authorization": f"Basic {client_credentials_b64}",
        "Content-Type": "application/x-www-form-urlencoded"
    }
    
    data = {
        "grant_type": "client_credentials"
    }
    
    response = requests.post(auth_url, headers=headers, data=data)
    response.raise_for_status()
    
    return response.json()["access_token"]


@dlt.resource(name="artists", write_disposition="replace", primary_key="id")
def get_artist_info() -> Iterator[Dict[str, Any]]:
    """Fetch Coldplay artist information"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(
        f"{SPOTIFY_API_BASE_URL}/artists/{COLDPLAY_ARTIST_ID}",
        headers=headers
    )
    response.raise_for_status()
    
    artist_data = response.json()
    artist_data["extracted_at"] = datetime.now().isoformat()
    yield artist_data


@dlt.resource(name="albums", write_disposition="replace", primary_key="id")
def get_artist_albums() -> Iterator[Dict[str, Any]]:
    """Fetch all Coldplay albums"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    offset = 0
    limit = 50
    
    while True:
        response = requests.get(
            f"{SPOTIFY_API_BASE_URL}/artists/{COLDPLAY_ARTIST_ID}/albums",
            headers=headers,
            params={
                "include_groups": "album,single,compilation",
                "limit": limit,
                "offset": offset,
                "market": "US"
            }
        )
        response.raise_for_status()
        
        data = response.json()
        
        for album in data["items"]:
            album["extracted_at"] = datetime.now().isoformat()
            yield album
        
        if not data.get("next"):
            break
        offset += limit


@dlt.resource(name="tracks", write_disposition="replace", primary_key="id")
def get_all_tracks() -> Iterator[Dict[str, Any]]:
    """Fetch all tracks from all Coldplay albums"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    # First get all albums
    albums = list(get_artist_albums())
    
    for album in albums:
        album_id = album["id"]
        
        # Get tracks for each album
        offset = 0
        limit = 50
        
        while True:
            response = requests.get(
                f"{SPOTIFY_API_BASE_URL}/albums/{album_id}/tracks",
                headers=headers,
                params={
                    "limit": limit,
                    "offset": offset,
                    "market": "US"
                }
            )
            response.raise_for_status()
            
            data = response.json()
            
            for track in data["items"]:
                track["album_id"] = album_id
                track["album_name"] = album["name"]
                track["album_release_date"] = album["release_date"]
                track["extracted_at"] = datetime.now().isoformat()
                yield track
            
            if not data.get("next"):
                break
            offset += limit


@dlt.resource(name="audio_features", write_disposition="replace", primary_key="id")
def get_audio_features() -> Iterator[Dict[str, Any]]:
    """Fetch audio features for all tracks"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    # Get all track IDs
    tracks = list(get_all_tracks())
    track_ids = [track["id"] for track in tracks]
    
    # Spotify allows up to 100 tracks per request
    batch_size = 100
    
    for i in range(0, len(track_ids), batch_size):
        batch_ids = track_ids[i:i + batch_size]
        
        response = requests.get(
            f"{SPOTIFY_API_BASE_URL}/audio-features",
            headers=headers,
            params={"ids": ",".join(batch_ids)}
        )
        response.raise_for_status()
        
        data = response.json()
        
        for features in data["audio_features"]:
            if features:  # Some tracks might not have audio features
                features["extracted_at"] = datetime.now().isoformat()
                yield features


@dlt.resource(name="top_tracks", write_disposition="replace", primary_key="id")
def get_top_tracks() -> Iterator[Dict[str, Any]]:
    """Fetch Coldplay's top tracks"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(
        f"{SPOTIFY_API_BASE_URL}/artists/{COLDPLAY_ARTIST_ID}/top-tracks",
        headers=headers,
        params={"market": "US"}
    )
    response.raise_for_status()
    
    data = response.json()
    
    for idx, track in enumerate(data["tracks"]):
        track["rank"] = idx + 1
        track["extracted_at"] = datetime.now().isoformat()
        yield track


@dlt.resource(name="related_artists", write_disposition="replace", primary_key="id")
def get_related_artists() -> Iterator[Dict[str, Any]]:
    """Fetch artists related to Coldplay"""
    token = get_spotify_access_token()
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(
        f"{SPOTIFY_API_BASE_URL}/artists/{COLDPLAY_ARTIST_ID}/related-artists",
        headers=headers
    )
    response.raise_for_status()
    
    data = response.json()
    
    for artist in data["artists"]:
        artist["related_to_artist_id"] = COLDPLAY_ARTIST_ID
        artist["extracted_at"] = datetime.now().isoformat()
        yield artist


def run_pipeline():
    """Run the DLT pipeline to extract Coldplay data"""
    # Create pipeline
    pipeline = dlt.pipeline(
        pipeline_name="spotify_coldplay",
        destination="duckdb",
        dataset_name="coldplay_data",
        dev_mode=True  # Use dev_mode instead of full_refresh
    )
    
    # Run the pipeline with available resources 
    # Note: audio_features and related_artists are restricted for new apps as of Nov 2024
    info = pipeline.run([
        get_artist_info(),
        get_artist_albums(),
        get_all_tracks(),
        get_top_tracks()
    ])
    
    print(info)
    print(f"\nPipeline completed successfully!")
    print(f"Data loaded to: {pipeline.pipeline_name}.duckdb")
    
    # Print some statistics
    with pipeline.sql_client() as client:
        # Count records in each table
        tables = ["artists", "albums", "tracks", "top_tracks"]
        
        print("\nData Statistics:")
        for table in tables:
            result = client.execute_sql(f"SELECT COUNT(*) as count FROM {table}")
            count = list(result)[0][0]
            print(f"  {table}: {count} records")


if __name__ == "__main__":
    run_pipeline()