import logging
import findspark
from cassandra.cluster import Cluster
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType



def setup_keyspace(session):
    session.execute("""
        CREATE KEYSPACE IF NOT EXISTS my_spark_data
        WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'};
    """)

    print("Keyspace set up successfully!")


def setup_table(session):
    session.execute("""
    CREATE TABLE IF NOT EXISTS my_spark_data.stream_users (
        id UUID PRIMARY KEY,
        first_name TEXT,
        last_name TEXT,
        gender TEXT,
        address TEXT,
        post_code TEXT,
        email TEXT,
        username TEXT,
        registered_date TEXT,
        phone TEXT,
        picture TEXT);
    """)

    print("Table created successfully!")


def add_data(session, **kwargs):
    print("inserting data...")

    user_id = kwargs.get('id')
    first_name = kwargs.get('first_name')
    last_name = kwargs.get('last_name')
    gender = kwargs.get('gender')
    address = kwargs.get('address')
    postcode = kwargs.get('post_code')
    email = kwargs.get('email')
    username = kwargs.get('username')
    dob = kwargs.get('dob')
    registered_date = kwargs.get('registered_date')
    phone = kwargs.get('phone')
    picture = kwargs.get('picture')

    try:
        session.execute("""
            INSERT INTO my_spark_data.stream_users(id, first_name, last_name, gender, address, 
                post_code, email, username, dob, registered_date, phone, picture)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """, (user_id, first_name, last_name, gender, address,
              postcode, email, username, dob, registered_date, phone, picture))
        logging.info(f"Data inserted for {first_name} {last_name}")

    except Exception as e:
        logging.error(f'could not insert data due to {e}')


def establish_spark_connection():
    s_conn = None

    try:
        s_conn = SparkSession.builder \
            .appName('StreamingProcessing') \
            .config('spark.jars.packages', "com.datastax.spark:spark-cassandra-connector_2.13:3.4.1,"
                                           "org.apache.spark:spark-sql-kafka-0-10_2.13:3.4.1") \
            .config('spark.cassandra.connection.host', 'localhost') \
            .getOrCreate()

        s_conn.sparkContext.setLogLevel("ERROR")
        logging.info("Spark connection established successfully!")
    except Exception as e:
        logging.error(f"Failed to create the spark session due to exception {e}")

    return s_conn


def link_to_kafka(spark_conn):
    spark_df = None
    try:
        spark_df = spark_conn.readStream \
            .format('kafka') \
            .option('kafka.bootstrap.servers', 'localhost:9092') \
            .option('subscribe', 'users_streamed') \
            .option('startingOffsets', 'earliest') \
            .load()
        logging.info("kafka dataframe linked successfully")
    except Exception as e:
        logging.warning(f"unable to link kafka dataframe because: {e}")

    return spark_df


def connect_cassandra():
    try:
        # connecting to the cassandra cluster
        cluster = Cluster(['localhost:9042'])

        cas_session = cluster.connect()

        return cas_session
    except Exception as e:
        logging.error(f"Could not establish cassandra connection due to {e}")
        return None


def prepare_selection_df_from_kafka(spark_df):
    schema = StructType([
        StructField("id", StringType(), False),
        StructField("first_name", StringType(), False),
        StructField("last_name", StringType(), False),
        StructField("gender", StringType(), False),
        StructField("address", StringType(), False),
        StructField("post_code", StringType(), False),
        StructField("email", StringType(), False),
        StructField("username", StringType(), False),
        StructField("registered_date", StringType(), False),
        StructField("phone", StringType(), False),
        StructField("picture", StringType(), False)
    ])

    sel = spark_df.selectExpr("CAST(value AS STRING)") \
        .select(from_json(col('value'), schema).alias('data')).select("data.*")
    print(sel)

    return sel


if __name__ == "__main__":
    # establish spark connection
    spark_conn = establish_spark_connection()

    if spark_conn is not None:
        # link to kafka with spark connection
        spark_df = link_to_kafka(spark_conn)
        selection_df = prepare_selection_df_from_kafka(spark_df)
        session = connect_cassandra()

        if session is not None:
            setup_keyspace(session)
            setup_table(session)

            logging.info("Start streaming...")

            streaming_query = (selection_df.writeStream.format("org.apache.spark.sql.cassandra")
                               .option('checkpointLocation', '/tmp/checkpoint')
                               .option('keyspace', 'my_spark_data')
                               .option('table', 'stream_users')
                               .start())

            streaming_query.awaitTermination()

I encountered an issue when running " spark-submit --master spark://localhost:7077 .\spark_stream.py"

A worker is created in the Spark UI but it doesn't fetch the data coming from Kafka.

I am using Airflow to trigger a DAG from an API to stream data into Kafka, process it in Spark, and save it into Cassandra.

The problem seems to be with communication between Spark and Kafka, even though each component works fine individually.