run_me_first.py

# -*- coding: utf-8 -*-
#
# Copyright 2022 Confluent Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Create topics and ksqlDB Streams

import sys
import time
import logging

from confluent_kafka.admin import NewTopic

from utils import (
    ksqldb,
    log_ini,
    log_exception,
    get_script_name,
    validate_cli_args,
    get_system_config,
    get_topic_partitions,
    set_producer_consumer,
)


####################
# Global variables #
####################
SCRIPT = get_script_name(__file__)
log_ini(SCRIPT, to_disk=False)

# Validate command arguments
kafka_config_file, sys_config_file = validate_cli_args(SCRIPT)

# Get system config file
SYS_CONFIG = get_system_config(sys_config_file)

# Set producer/consumer objects
KAFKA_CONFIG, _, _, ADMIN_CLIENT = set_producer_consumer(
    kafka_config_file,
    disable_producer=True,
    disable_consumer=True,
)

TOPIC_ORDERED = SYS_CONFIG["kafka-topics"]["pizza_ordered"]
STREAM_ORDERED = TOPIC_ORDERED.replace("-", "_").upper()

TOPIC_PENDING = SYS_CONFIG["kafka-topics"]["pizza_pending"]
STREAM_PENDING = TOPIC_PENDING.replace("-", "_").upper()

TOPIC_ASSEMBLED = SYS_CONFIG["kafka-topics"]["pizza_assembled"]
STREAM_ASSEMBLED = TOPIC_ASSEMBLED.replace("-", "_").upper()

TOPIC_BAKED = SYS_CONFIG["kafka-topics"]["pizza_baked"]
STREAM_BAKED = TOPIC_BAKED.replace("-", "_").upper()

TOPIC_DELIVERED = SYS_CONFIG["kafka-topics"]["pizza_delivered"]
STREAM_DELIVERED = TOPIC_DELIVERED.replace("-", "_").upper()

TOPIC_STATUS = SYS_CONFIG["kafka-topics"]["pizza_status"]
STREAM_STATUS = TOPIC_STATUS.replace("-", "_").upper()

# TOPIC_STATUS_SESSION = SYS_CONFIG["kafka-topics"]["pizza_status_session"]
# TABLE_STATUS_SESSION = TOPIC_STATUS_SESSION.replace("-", "_").upper()

# Persistent queries to push all statuses generated by each stream to the status stream
PERSISTENT_QUERIES = {
    f"PERSISTENT_QUERY_{stream}": f"""INSERT INTO {STREAM_STATUS} SELECT order_id, status, timestamp FROM {stream} EMIT CHANGES;"""
    for stream in [
        STREAM_ORDERED,
        STREAM_PENDING,
        STREAM_ASSEMBLED,
        STREAM_BAKED,
        STREAM_DELIVERED,
    ]
}

KSQL_STATEMENTS = {
    STREAM_ORDERED: f"""CREATE STREAM IF NOT EXISTS {STREAM_ORDERED} (
        order_id VARCHAR KEY,
        status INT,
        timestamp BIGINT,
        order STRUCT<
            extra_toppings ARRAY<STRING>,
            username STRING,
            customer_id STRING,
            sauce STRING,
            cheese STRING,
            main_topping STRING
        >
    ) WITH (
        KAFKA_TOPIC = '{TOPIC_ORDERED}',
        VALUE_FORMAT = 'JSON',
        TIMESTAMP = 'timestamp'
    );""",
    STREAM_ASSEMBLED: f"""CREATE STREAM IF NOT EXISTS {STREAM_ASSEMBLED} (
        order_id VARCHAR KEY,
        status INT,
        baking_time INT,
        timestamp BIGINT
    ) WITH (
        KAFKA_TOPIC = '{TOPIC_ASSEMBLED}',
        VALUE_FORMAT = 'JSON',
        TIMESTAMP = 'timestamp'
    );""",
    STREAM_BAKED: f"""CREATE STREAM IF NOT EXISTS {STREAM_BAKED} (
        order_id VARCHAR KEY,
        status INT,
        timestamp BIGINT
    ) WITH (
        KAFKA_TOPIC = '{TOPIC_BAKED}',
        VALUE_FORMAT = 'JSON',
        TIMESTAMP = 'timestamp'
    );""",
    STREAM_DELIVERED: f"""CREATE STREAM IF NOT EXISTS {STREAM_DELIVERED} (
        order_id VARCHAR KEY,
        status INT,
        timestamp BIGINT
    ) WITH (
        KAFKA_TOPIC = '{TOPIC_DELIVERED}',
        VALUE_FORMAT = 'JSON',
        TIMESTAMP = 'timestamp'
    );""",
    STREAM_PENDING: f"""CREATE STREAM IF NOT EXISTS {STREAM_PENDING} (
        order_id VARCHAR KEY,
        status INT,
        timestamp BIGINT
    ) WITH (
        KAFKA_TOPIC = '{TOPIC_PENDING}',
        VALUE_FORMAT = 'JSON',
        TIMESTAMP = 'timestamp'
    );""",
    STREAM_STATUS: f"""CREATE STREAM IF NOT EXISTS {STREAM_STATUS} (
        order_id VARCHAR KEY,
        status INT,
        timestamp BIGINT
    ) WITH (
        KAFKA_TOPIC='{TOPIC_STATUS}',
        VALUE_FORMAT='JSON',
        TIMESTAMP='timestamp'
    );""",
    **PERSISTENT_QUERIES,
}

if __name__ == "__main__":
    # Create topics
    num_partitions = int(SYS_CONFIG["kafka-topic-config"]["num_partitions"])
    replication_factor = int(SYS_CONFIG["kafka-topic-config"]["replication_factor"])
    for alias, topic in SYS_CONFIG["kafka-topics"].items():
        logging.info(
            f"Creating topic {alias} as '{topic}' with {num_partitions} partition(s) and replication factor {replication_factor}..."
        )
        try:
            partitions = get_topic_partitions(
                ADMIN_CLIENT,
                topic,
                default_partition_number=0,
            )
            if partitions == 0:
                result = (
                    ADMIN_CLIENT.create_topics(
                        [
                            NewTopic(
                                topic,
                                num_partitions,
                                replication_factor,
                            ),
                        ]
                    )
                    .get(topic)
                    .result()
                )
            else:
                logging.info("Topic already exists!")
                continue
        except Exception:
            log_exception(
                f"Error when creating topic!",
                sys.exc_info(),
            )
        else:
            logging.info("Done!")

    # Create ksqlDB Streams/Tables/Persistent queries
    for ksql_name, statement in KSQL_STATEMENTS.items():
        logging.info(f"Creating ksqlDB: {ksql_name}...")
        ksqldb(
            KAFKA_CONFIG["ksqldb"]["endpoint"],
            statement,
            username=KAFKA_CONFIG["ksqldb"].get("username"),
            password=KAFKA_CONFIG["ksqldb"].get("password"),
        )
        time.sleep(3)