misc.python.materialize.zippy.kafka_actions
1# Copyright Materialize, Inc. and contributors. All rights reserved. 2# 3# Use of this software is governed by the Business Source License 4# included in the LICENSE file at the root of this repository. 5# 6# As of the Change Date specified in that file, in accordance with 7# the Business Source License, use of this software will be governed 8# by the Apache License, Version 2.0. 9 10import random 11import string 12import threading 13from textwrap import dedent 14 15from materialize.mzcompose.composition import Composition 16from materialize.zippy.framework import ( 17 Action, 18 ActionFactory, 19 Capabilities, 20 Capability, 21 State, 22) 23from materialize.zippy.kafka_capabilities import Envelope, KafkaRunning, TopicExists 24from materialize.zippy.mz_capabilities import MzIsRunning 25 26SCHEMA = """ 27$ set keyschema={ 28 "type" : "record", 29 "name" : "test", 30 "fields" : [ 31 {"name":"key", "type":"long"} 32 ] 33 } 34 35$ set schema={ 36 "type" : "record", 37 "name" : "test", 38 "fields" : [ 39 {"name":"f1", "type":"long"}, 40 {"name":"pad", "type":"string"} 41 ] 42 } 43""" 44 45 46class KafkaStart(Action): 47 """Start a Kafka instance.""" 48 49 def provides(self) -> list[Capability]: 50 return [KafkaRunning()] 51 52 def run(self, c: Composition, state: State) -> None: 53 c.up("redpanda") 54 55 56class CreateTopicParameterized(ActionFactory): 57 """Creates a Kafka topic and decides on the envelope that will be used.""" 58 59 @classmethod 60 def requires(cls) -> set[type[Capability]]: 61 return {MzIsRunning, KafkaRunning} 62 63 def __init__( 64 self, 65 max_topics: int = 10, 66 envelopes_with_weights: dict[Envelope, int] = { 67 Envelope.NONE: 25, 68 Envelope.UPSERT: 75, 69 }, 70 ) -> None: 71 self.max_topics = max_topics 72 self.envelopes_with_weights = envelopes_with_weights 73 74 def new(self, capabilities: Capabilities) -> list[Action]: 75 new_topic_name = capabilities.get_free_capability_name( 76 TopicExists, self.max_topics 77 ) 78 79 if new_topic_name: 80 return [ 81 CreateTopic( 82 capabilities=capabilities, 83 topic=TopicExists( 84 name=new_topic_name, 85 envelope=random.choices( 86 list(self.envelopes_with_weights.keys()), 87 weights=list(self.envelopes_with_weights.values()), 88 )[0], 89 partitions=random.randint(1, 10), 90 ), 91 ) 92 ] 93 else: 94 return [] 95 96 97class CreateTopic(Action): 98 def __init__(self, capabilities: Capabilities, topic: TopicExists) -> None: 99 self.topic = topic 100 super().__init__(capabilities) 101 102 def provides(self) -> list[Capability]: 103 return [self.topic] 104 105 def run(self, c: Composition, state: State) -> None: 106 c.testdrive( 107 SCHEMA + dedent(f""" 108 $ kafka-create-topic topic={self.topic.name} partitions={self.topic.partitions} 109 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} repeat=1 110 {{"key": 0}} {{"f1": 0, "pad": ""}} 111 """), 112 mz_service=state.mz_service, 113 ) 114 115 116class Ingest(Action): 117 """Ingests data (inserts, updates or deletions) into a Kafka topic.""" 118 119 @classmethod 120 def requires(cls) -> set[type[Capability]]: 121 return {MzIsRunning, KafkaRunning, TopicExists} 122 123 def __init__(self, capabilities: Capabilities) -> None: 124 self.topic = random.choice(capabilities.get(TopicExists)) 125 self.delta = random.randint(1, 10000) 126 # Heavy-tailed pad sizes: mostly up to 10 bytes, some up to 100 bytes 127 # and outliers up to 256 bytes. int(paretovariate(0.6)) approximates a 128 # zipf(1.6) sample while staying reproducible under random.seed(). 129 self.pad = min(int(random.paretovariate(0.6)), 256) * random.choice( 130 string.ascii_letters 131 ) 132 super().__init__(capabilities) 133 134 def __str__(self) -> str: 135 return f"{Action.__str__(self)} {self.topic.name}" 136 137 138class KafkaInsert(Ingest): 139 """Inserts data into a Kafka topic.""" 140 141 def parallel(self) -> bool: 142 return False 143 144 def run(self, c: Composition, state: State) -> None: 145 prev_max = self.topic.watermarks.max 146 self.topic.watermarks.max = prev_max + self.delta 147 assert self.topic.watermarks.max >= 0 148 assert self.topic.watermarks.min >= 0 149 150 testdrive_str = SCHEMA + dedent(f""" 151 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} start-iteration={prev_max + 1} repeat={self.delta} 152 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad" : "{self.pad}"}} 153 """) 154 155 if self.parallel(): 156 mz_service = state.mz_service 157 158 def ingest() -> None: 159 try: 160 c.testdrive(testdrive_str, mz_service=mz_service) 161 except Exception as e: 162 state.background_errors.append(e) 163 164 thread = threading.Thread(target=ingest) 165 thread.start() 166 state.background_threads.append(thread) 167 else: 168 c.testdrive(testdrive_str, mz_service=state.mz_service) 169 170 171class KafkaInsertParallel(KafkaInsert): 172 """Inserts data into a Kafka topic using background threads.""" 173 174 @classmethod 175 def require_explicit_mention(cls) -> bool: 176 return True 177 178 def parallel(self) -> bool: 179 return True 180 181 182class KafkaUpsertFromHead(Ingest): 183 """Updates records from the head in-place by modifying their pad""" 184 185 def run(self, c: Composition, state: State) -> None: 186 if self.topic.envelope is Envelope.NONE: 187 return 188 189 head = self.topic.watermarks.max 190 start = max(head - self.delta, self.topic.watermarks.min) 191 actual_delta = head - start 192 193 if actual_delta > 0: 194 c.testdrive( 195 SCHEMA + dedent(f""" 196 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={start} repeat={actual_delta} 197 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 198 """), 199 mz_service=state.mz_service, 200 ) 201 202 203class KafkaDeleteFromHead(Ingest): 204 """Deletes the largest values previously inserted.""" 205 206 def run(self, c: Composition, state: State) -> None: 207 if self.topic.envelope is Envelope.NONE: 208 return 209 210 prev_max = self.topic.watermarks.max 211 self.topic.watermarks.max = max( 212 prev_max - self.delta, self.topic.watermarks.min 213 ) 214 assert self.topic.watermarks.max >= 0 215 assert self.topic.watermarks.min >= 0 216 217 actual_delta = prev_max - self.topic.watermarks.max 218 219 if actual_delta > 0: 220 c.testdrive( 221 SCHEMA + dedent(f""" 222 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={self.topic.watermarks.max + 1} repeat={actual_delta} 223 {{"key": ${{kafka-ingest.iteration}}}} 224 """), 225 mz_service=state.mz_service, 226 ) 227 228 229class KafkaUpsertFromTail(Ingest): 230 """Updates records from the tail in-place by modifying their pad""" 231 232 def run(self, c: Composition, state: State) -> None: 233 if self.topic.envelope is Envelope.NONE: 234 return 235 236 tail = self.topic.watermarks.min 237 end = min(tail + self.delta, self.topic.watermarks.max) 238 actual_delta = end - tail 239 240 if actual_delta > 0: 241 c.testdrive( 242 SCHEMA + dedent(f""" 243 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={tail} repeat={actual_delta} 244 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 245 """), 246 mz_service=state.mz_service, 247 ) 248 249 250class KafkaDeleteFromTail(Ingest): 251 """Deletes the smallest values previously inserted.""" 252 253 def run(self, c: Composition, state: State) -> None: 254 if self.topic.envelope is Envelope.NONE: 255 return 256 257 prev_min = self.topic.watermarks.min 258 self.topic.watermarks.min = min( 259 prev_min + self.delta, self.topic.watermarks.max 260 ) 261 assert self.topic.watermarks.max >= 0 262 assert self.topic.watermarks.min >= 0 263 actual_delta = self.topic.watermarks.min - prev_min 264 265 if actual_delta > 0: 266 c.testdrive( 267 SCHEMA + dedent(f""" 268 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={prev_min} repeat={actual_delta} 269 {{"key": ${{kafka-ingest.iteration}}}} 270 """), 271 mz_service=state.mz_service, 272 )
47class KafkaStart(Action): 48 """Start a Kafka instance.""" 49 50 def provides(self) -> list[Capability]: 51 return [KafkaRunning()] 52 53 def run(self, c: Composition, state: State) -> None: 54 c.up("redpanda")
Start a Kafka instance.
57class CreateTopicParameterized(ActionFactory): 58 """Creates a Kafka topic and decides on the envelope that will be used.""" 59 60 @classmethod 61 def requires(cls) -> set[type[Capability]]: 62 return {MzIsRunning, KafkaRunning} 63 64 def __init__( 65 self, 66 max_topics: int = 10, 67 envelopes_with_weights: dict[Envelope, int] = { 68 Envelope.NONE: 25, 69 Envelope.UPSERT: 75, 70 }, 71 ) -> None: 72 self.max_topics = max_topics 73 self.envelopes_with_weights = envelopes_with_weights 74 75 def new(self, capabilities: Capabilities) -> list[Action]: 76 new_topic_name = capabilities.get_free_capability_name( 77 TopicExists, self.max_topics 78 ) 79 80 if new_topic_name: 81 return [ 82 CreateTopic( 83 capabilities=capabilities, 84 topic=TopicExists( 85 name=new_topic_name, 86 envelope=random.choices( 87 list(self.envelopes_with_weights.keys()), 88 weights=list(self.envelopes_with_weights.values()), 89 )[0], 90 partitions=random.randint(1, 10), 91 ), 92 ) 93 ] 94 else: 95 return []
Creates a Kafka topic and decides on the envelope that will be used.
60 @classmethod 61 def requires(cls) -> set[type[Capability]]: 62 return {MzIsRunning, KafkaRunning}
Compute the capability classes that this Action Factory requires.
75 def new(self, capabilities: Capabilities) -> list[Action]: 76 new_topic_name = capabilities.get_free_capability_name( 77 TopicExists, self.max_topics 78 ) 79 80 if new_topic_name: 81 return [ 82 CreateTopic( 83 capabilities=capabilities, 84 topic=TopicExists( 85 name=new_topic_name, 86 envelope=random.choices( 87 list(self.envelopes_with_weights.keys()), 88 weights=list(self.envelopes_with_weights.values()), 89 )[0], 90 partitions=random.randint(1, 10), 91 ), 92 ) 93 ] 94 else: 95 return []
98class CreateTopic(Action): 99 def __init__(self, capabilities: Capabilities, topic: TopicExists) -> None: 100 self.topic = topic 101 super().__init__(capabilities) 102 103 def provides(self) -> list[Capability]: 104 return [self.topic] 105 106 def run(self, c: Composition, state: State) -> None: 107 c.testdrive( 108 SCHEMA + dedent(f""" 109 $ kafka-create-topic topic={self.topic.name} partitions={self.topic.partitions} 110 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} repeat=1 111 {{"key": 0}} {{"f1": 0, "pad": ""}} 112 """), 113 mz_service=state.mz_service, 114 )
Base class for an action that a Zippy test can take.
99 def __init__(self, capabilities: Capabilities, topic: TopicExists) -> None: 100 self.topic = topic 101 super().__init__(capabilities)
Construct a new action, possibly conditioning on the available capabilities.
Compute the capabilities that this action will make available.
106 def run(self, c: Composition, state: State) -> None: 107 c.testdrive( 108 SCHEMA + dedent(f""" 109 $ kafka-create-topic topic={self.topic.name} partitions={self.topic.partitions} 110 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} repeat=1 111 {{"key": 0}} {{"f1": 0, "pad": ""}} 112 """), 113 mz_service=state.mz_service, 114 )
Run this action on the provided composition.
117class Ingest(Action): 118 """Ingests data (inserts, updates or deletions) into a Kafka topic.""" 119 120 @classmethod 121 def requires(cls) -> set[type[Capability]]: 122 return {MzIsRunning, KafkaRunning, TopicExists} 123 124 def __init__(self, capabilities: Capabilities) -> None: 125 self.topic = random.choice(capabilities.get(TopicExists)) 126 self.delta = random.randint(1, 10000) 127 # Heavy-tailed pad sizes: mostly up to 10 bytes, some up to 100 bytes 128 # and outliers up to 256 bytes. int(paretovariate(0.6)) approximates a 129 # zipf(1.6) sample while staying reproducible under random.seed(). 130 self.pad = min(int(random.paretovariate(0.6)), 256) * random.choice( 131 string.ascii_letters 132 ) 133 super().__init__(capabilities) 134 135 def __str__(self) -> str: 136 return f"{Action.__str__(self)} {self.topic.name}"
Ingests data (inserts, updates or deletions) into a Kafka topic.
124 def __init__(self, capabilities: Capabilities) -> None: 125 self.topic = random.choice(capabilities.get(TopicExists)) 126 self.delta = random.randint(1, 10000) 127 # Heavy-tailed pad sizes: mostly up to 10 bytes, some up to 100 bytes 128 # and outliers up to 256 bytes. int(paretovariate(0.6)) approximates a 129 # zipf(1.6) sample while staying reproducible under random.seed(). 130 self.pad = min(int(random.paretovariate(0.6)), 256) * random.choice( 131 string.ascii_letters 132 ) 133 super().__init__(capabilities)
Construct a new action, possibly conditioning on the available capabilities.
139class KafkaInsert(Ingest): 140 """Inserts data into a Kafka topic.""" 141 142 def parallel(self) -> bool: 143 return False 144 145 def run(self, c: Composition, state: State) -> None: 146 prev_max = self.topic.watermarks.max 147 self.topic.watermarks.max = prev_max + self.delta 148 assert self.topic.watermarks.max >= 0 149 assert self.topic.watermarks.min >= 0 150 151 testdrive_str = SCHEMA + dedent(f""" 152 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} start-iteration={prev_max + 1} repeat={self.delta} 153 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad" : "{self.pad}"}} 154 """) 155 156 if self.parallel(): 157 mz_service = state.mz_service 158 159 def ingest() -> None: 160 try: 161 c.testdrive(testdrive_str, mz_service=mz_service) 162 except Exception as e: 163 state.background_errors.append(e) 164 165 thread = threading.Thread(target=ingest) 166 thread.start() 167 state.background_threads.append(thread) 168 else: 169 c.testdrive(testdrive_str, mz_service=state.mz_service)
Inserts data into a Kafka topic.
145 def run(self, c: Composition, state: State) -> None: 146 prev_max = self.topic.watermarks.max 147 self.topic.watermarks.max = prev_max + self.delta 148 assert self.topic.watermarks.max >= 0 149 assert self.topic.watermarks.min >= 0 150 151 testdrive_str = SCHEMA + dedent(f""" 152 $ kafka-ingest format=avro key-format=avro topic={self.topic.name} schema=${{schema}} key-schema=${{keyschema}} start-iteration={prev_max + 1} repeat={self.delta} 153 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad" : "{self.pad}"}} 154 """) 155 156 if self.parallel(): 157 mz_service = state.mz_service 158 159 def ingest() -> None: 160 try: 161 c.testdrive(testdrive_str, mz_service=mz_service) 162 except Exception as e: 163 state.background_errors.append(e) 164 165 thread = threading.Thread(target=ingest) 166 thread.start() 167 state.background_threads.append(thread) 168 else: 169 c.testdrive(testdrive_str, mz_service=state.mz_service)
Run this action on the provided composition.
172class KafkaInsertParallel(KafkaInsert): 173 """Inserts data into a Kafka topic using background threads.""" 174 175 @classmethod 176 def require_explicit_mention(cls) -> bool: 177 return True 178 179 def parallel(self) -> bool: 180 return True
Inserts data into a Kafka topic using background threads.
183class KafkaUpsertFromHead(Ingest): 184 """Updates records from the head in-place by modifying their pad""" 185 186 def run(self, c: Composition, state: State) -> None: 187 if self.topic.envelope is Envelope.NONE: 188 return 189 190 head = self.topic.watermarks.max 191 start = max(head - self.delta, self.topic.watermarks.min) 192 actual_delta = head - start 193 194 if actual_delta > 0: 195 c.testdrive( 196 SCHEMA + dedent(f""" 197 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={start} repeat={actual_delta} 198 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 199 """), 200 mz_service=state.mz_service, 201 )
Updates records from the head in-place by modifying their pad
186 def run(self, c: Composition, state: State) -> None: 187 if self.topic.envelope is Envelope.NONE: 188 return 189 190 head = self.topic.watermarks.max 191 start = max(head - self.delta, self.topic.watermarks.min) 192 actual_delta = head - start 193 194 if actual_delta > 0: 195 c.testdrive( 196 SCHEMA + dedent(f""" 197 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={start} repeat={actual_delta} 198 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 199 """), 200 mz_service=state.mz_service, 201 )
Run this action on the provided composition.
204class KafkaDeleteFromHead(Ingest): 205 """Deletes the largest values previously inserted.""" 206 207 def run(self, c: Composition, state: State) -> None: 208 if self.topic.envelope is Envelope.NONE: 209 return 210 211 prev_max = self.topic.watermarks.max 212 self.topic.watermarks.max = max( 213 prev_max - self.delta, self.topic.watermarks.min 214 ) 215 assert self.topic.watermarks.max >= 0 216 assert self.topic.watermarks.min >= 0 217 218 actual_delta = prev_max - self.topic.watermarks.max 219 220 if actual_delta > 0: 221 c.testdrive( 222 SCHEMA + dedent(f""" 223 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={self.topic.watermarks.max + 1} repeat={actual_delta} 224 {{"key": ${{kafka-ingest.iteration}}}} 225 """), 226 mz_service=state.mz_service, 227 )
Deletes the largest values previously inserted.
207 def run(self, c: Composition, state: State) -> None: 208 if self.topic.envelope is Envelope.NONE: 209 return 210 211 prev_max = self.topic.watermarks.max 212 self.topic.watermarks.max = max( 213 prev_max - self.delta, self.topic.watermarks.min 214 ) 215 assert self.topic.watermarks.max >= 0 216 assert self.topic.watermarks.min >= 0 217 218 actual_delta = prev_max - self.topic.watermarks.max 219 220 if actual_delta > 0: 221 c.testdrive( 222 SCHEMA + dedent(f""" 223 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={self.topic.watermarks.max + 1} repeat={actual_delta} 224 {{"key": ${{kafka-ingest.iteration}}}} 225 """), 226 mz_service=state.mz_service, 227 )
Run this action on the provided composition.
230class KafkaUpsertFromTail(Ingest): 231 """Updates records from the tail in-place by modifying their pad""" 232 233 def run(self, c: Composition, state: State) -> None: 234 if self.topic.envelope is Envelope.NONE: 235 return 236 237 tail = self.topic.watermarks.min 238 end = min(tail + self.delta, self.topic.watermarks.max) 239 actual_delta = end - tail 240 241 if actual_delta > 0: 242 c.testdrive( 243 SCHEMA + dedent(f""" 244 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={tail} repeat={actual_delta} 245 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 246 """), 247 mz_service=state.mz_service, 248 )
Updates records from the tail in-place by modifying their pad
233 def run(self, c: Composition, state: State) -> None: 234 if self.topic.envelope is Envelope.NONE: 235 return 236 237 tail = self.topic.watermarks.min 238 end = min(tail + self.delta, self.topic.watermarks.max) 239 actual_delta = end - tail 240 241 if actual_delta > 0: 242 c.testdrive( 243 SCHEMA + dedent(f""" 244 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={tail} repeat={actual_delta} 245 {{"key": ${{kafka-ingest.iteration}}}} {{"f1": ${{kafka-ingest.iteration}}, "pad": "{self.pad}"}} 246 """), 247 mz_service=state.mz_service, 248 )
Run this action on the provided composition.
251class KafkaDeleteFromTail(Ingest): 252 """Deletes the smallest values previously inserted.""" 253 254 def run(self, c: Composition, state: State) -> None: 255 if self.topic.envelope is Envelope.NONE: 256 return 257 258 prev_min = self.topic.watermarks.min 259 self.topic.watermarks.min = min( 260 prev_min + self.delta, self.topic.watermarks.max 261 ) 262 assert self.topic.watermarks.max >= 0 263 assert self.topic.watermarks.min >= 0 264 actual_delta = self.topic.watermarks.min - prev_min 265 266 if actual_delta > 0: 267 c.testdrive( 268 SCHEMA + dedent(f""" 269 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={prev_min} repeat={actual_delta} 270 {{"key": ${{kafka-ingest.iteration}}}} 271 """), 272 mz_service=state.mz_service, 273 )
Deletes the smallest values previously inserted.
254 def run(self, c: Composition, state: State) -> None: 255 if self.topic.envelope is Envelope.NONE: 256 return 257 258 prev_min = self.topic.watermarks.min 259 self.topic.watermarks.min = min( 260 prev_min + self.delta, self.topic.watermarks.max 261 ) 262 assert self.topic.watermarks.max >= 0 263 assert self.topic.watermarks.min >= 0 264 actual_delta = self.topic.watermarks.min - prev_min 265 266 if actual_delta > 0: 267 c.testdrive( 268 SCHEMA + dedent(f""" 269 $ kafka-ingest format=avro topic={self.topic.name} key-format=avro key-schema=${{keyschema}} schema=${{schema}} start-iteration={prev_min} repeat={actual_delta} 270 {{"key": ${{kafka-ingest.iteration}}}} 271 """), 272 mz_service=state.mz_service, 273 )
Run this action on the provided composition.