mage-ai 0.9.45__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/shared/models.py

@@ -0,0 +1,254 @@
+import typing
+from dataclasses import dataclass, make_dataclass
+from enum import Enum
+from functools import reduce
+from typing import Dict, List
+
+import inflection
+
+from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.parsers import encode_complex
+
+
+class BaseClass:
+    attribute_aliases = {}
+    disable_attribute_snake_case = False
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @classmethod
+    def all_annotations(self) -> Dict:
+        annotations = {}
+
+        for key, value in self.__dataclass_fields__.items():
+            annotations[key] = value.type
+
+        for parent_class in self.__bases__:
+            if issubclass(parent_class, BaseDataClass):
+                for key, value in parent_class.__dataclass_fields__.items():
+                    if key not in annotations:
+                        annotations[key] = value.type
+
+        return annotations
+
+    @classmethod
+    def load(self, **kwargs):
+        annotations = self.all_annotations()
+
+        props_init = self.load_to_dict(**kwargs)
+
+        props = {}
+        props_not_set = {}
+        if props_init:
+            for key, value in props_init.items():
+                if self.attribute_aliases and key in self.attribute_aliases:
+                    key = self.attribute_aliases[key]
+
+                annotation = annotations.get(key)
+                if annotation:
+                    props[key] = self.convert_value(value, annotation)
+                else:
+                    props_not_set[key] = value
+
+        model = self(**props)
+
+        for key, value in props_not_set.items():
+            try:
+                if not callable(getattr(model, key)):
+                    model.set_value(key, value)
+            except AttributeError as err:
+                print(f'[WARNING] {self.__name__}.load: {err}')
+                raise err
+
+        return model
+
+    def set_value(self, key, value):
+        value = self.convert_value(value)
+
+        try:
+            setattr(self, key, value)
+        except AttributeError:
+            pass
+
+    @classmethod
+    def convert_value(
+        self,
+        value,
+        annotation=None,
+        convert_enum: bool = False,
+        ignore_empty: bool = False,
+    ):
+        is_list = isinstance(value, list)
+        if is_list:
+            if ignore_empty and len(value) == 0:
+                return None
+
+            return [self.convert_value(
+                v,
+                convert_enum=convert_enum,
+                ignore_empty=ignore_empty,
+            ) for v in value]
+
+        if not annotation:
+            annotation = type(value)
+
+        is_dict_class = isinstance(value, dict)
+
+        def _build_dict(acc, kv, cls=self):
+            key, value = kv
+            acc[key] = cls.convert_value(
+                value,
+                convert_enum=convert_enum,
+                ignore_empty=ignore_empty,
+            )
+            return acc
+
+        is_typing_class = False
+        if hasattr(typing, '_BaseGenericAlias'):
+            is_typing_class = issubclass(annotation.__class__, typing._BaseGenericAlias)
+        elif hasattr(typing, '_GenericAlias'):
+            is_typing_class = issubclass(annotation.__class__, typing._GenericAlias)
+
+        if is_typing_class:
+            if is_dict_class:
+                if ignore_empty and len(value) == 0:
+                    return None
+
+                return reduce(_build_dict, value.items(), {})
+            else:
+                return value
+
+        is_data_class = issubclass(annotation, BaseDataClass)
+        if is_data_class:
+            if is_dict_class:
+                if ignore_empty and len(value) == 0:
+                    return None
+
+                return annotation.load(**value)
+            elif isinstance(value, BaseDataClass):
+                return value.to_dict(
+                    convert_enum=convert_enum,
+                    ignore_empty=ignore_empty,
+                )
+
+        is_enum_class = issubclass(annotation, Enum)
+        is_enum = isinstance(value, Enum)
+        if is_enum_class and not is_enum:
+            try:
+                return annotation(value)
+            except ValueError:
+                pass
+
+        if convert_enum and is_enum:
+            return value.value
+
+        return value
+
+    @classmethod
+    def load_to_dict(self, **kwargs) -> Dict:
+        data = {}
+        for key, value in kwargs.items():
+            if not self.disable_attribute_snake_case:
+                key = inflection.underscore(key)
+            data[key] = value
+        return data
+
+    def serialize_attribute_class(self, attribute_name: str, attribute_class, **kwargs):
+        try:
+            value = getattr(self, attribute_name)
+            if value and isinstance(value, dict):
+                setattr(self, attribute_name, attribute_class.load(**merge_dict(value, kwargs)))
+        except AttributeError as err:
+            print(f'[WARNING] {self.__class__.__name__}.serialize_attribute_class: {err}')
+
+    def serialize_attribute_classes(self, attribute_name: str, attribute_class, **kwargs):
+        try:
+            value = getattr(self, attribute_name)
+            if value and isinstance(value, list):
+                arr = []
+                for model in value:
+                    if isinstance(model, dict):
+                        arr.append(attribute_class.load(**merge_dict(model, kwargs)))
+                    else:
+                        arr.append(model)
+                setattr(self, attribute_name, arr)
+        except AttributeError as err:
+            print(f'[WARNING] {self.__class__.__name__}.serialize_attribute_classes: {err}')
+
+    def serialize_attribute_enum(self, attribute_name: str, enum_class):
+        try:
+            value = getattr(self, attribute_name)
+            if value and isinstance(value, str):
+                setattr(self, attribute_name, enum_class(value))
+        except AttributeError as err:
+            print(f'[WARNING] {self.__class__.__name__}.serialize_attribute_enum: {err}')
+
+    def serialize_attribute_enums(self, attribute_name: str, enum_class):
+        try:
+            values = getattr(self, attribute_name)
+            if values and isinstance(values, list):
+                arr = []
+                for value in values:
+                    if isinstance(value, str):
+                        arr.append(enum_class(value))
+                    else:
+                        arr.append(value)
+                setattr(self, attribute_name, arr)
+        except AttributeError as err:
+            print(f'[WARNING] {self.__class__.__name__}.serialize_attribute_enums: {err}')
+
+    def to_dict(self, convert_enum: bool = False, ignore_empty: bool = False, **kwargs) -> Dict:
+        data = {}
+
+        for key, annotation in self.all_annotations().items():
+            value = None
+            try:
+                value = getattr(self, key)
+            except AttributeError as err:
+                print(f'[WARNING] {self.__class__.__name__}.to_dict: {err}')
+
+            value = self.convert_value(
+                value,
+                annotation,
+                convert_enum=convert_enum,
+                ignore_empty=ignore_empty,
+            )
+            if not ignore_empty or value is not None:
+                if self.attribute_aliases and key in self.attribute_aliases:
+                    key = self.attribute_aliases[key]
+
+                if ignore_empty and \
+                        (isinstance(value, list) or isinstance(value, dict)) and \
+                        len(value) == 0:
+
+                    continue
+
+                data[key] = encode_complex(value)
+
+        return data
+
+    def update_attributes(self, **kwargs):
+        for key, value in kwargs.items():
+            self.set_value(key, value)
+
+
+@dataclass
+class BaseDataClass(BaseClass):
+    @classmethod
+    def dynamic_fields(self, *args, **kwargs) -> List:
+        return None
+
+    def __new__(cls, *args, **kwargs):
+        fields = cls.dynamic_fields()
+        if fields:
+            cls.__class__ = make_dataclass(
+                cls.__name__,
+                fields=fields,
+                bases=(cls,),
+            )
+
+        obj = object.__new__(cls)
+        BaseClass.__init__(obj, *args, **kwargs)
+
+        return obj

mage_ai/streaming/constants.py

@@ -5,11 +5,13 @@ DEFAULT_TIMEOUT_MS = 500
 
 
 class SourceType(str, Enum):
+    ACTIVEMQ = 'activemq'
     AMAZON_SQS = 'amazon_sqs'
     AZURE_EVENT_HUB = 'azure_event_hub'
     GOOGLE_CLOUD_PUBSUB = 'google_cloud_pubsub'
     INFLUXDB = 'influxdb'
     KAFKA = 'kafka'
+    NATS = 'nats'
     KINESIS = 'kinesis'
     RABBITMQ = 'rabbitmq'
     MONGODB = 'mongodb'
@@ -35,6 +37,7 @@ class SinkType(str, Enum):
     OPENSEARCH = 'opensearch'
     ORACLEDB = 'oracledb'
     POSTGRES = 'postgres'
+    RABBITMQ = 'rabbitmq'
     REDSHIFT = 'redshift'
     SNOWFLAKE = 'snowflake'
     TRINO = 'trino'

mage_ai/streaming/sinks/postgres.py

@@ -44,6 +44,8 @@ class PostgresSink(BaseSink):
         self.batch_write([message])
 
     def batch_write(self, messages: List[Dict]):
+        if not messages:
+            return
         self._print(
             f'Batch ingest {len(messages)} records, time={time.time()}. Sample: {messages[0]}')
 

mage_ai/streaming/sinks/rabbitmq.py

@@ -0,0 +1,76 @@
+import json
+from dataclasses import dataclass
+from typing import Dict, List
+
+import pika
+
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sinks.base import BaseSink
+
+
+@dataclass
+class RabbitMQConfig(BaseConfig):
+    connection_host: str
+    connection_port: int
+    queue_name: str
+    username: str = 'guest'
+    password: str = 'guest'
+    amqp_url_virtual_host: str = r'%2f'
+
+
+class RabbitMQSink(BaseSink):
+    config_class = RabbitMQConfig
+
+    def init_client(self):
+        self._print('Start initializing producer.')
+        # Initialize RabbitMQ producer
+        queue_name = self.config.queue_name
+        username = self.config.username
+        password = self.config.password
+        connection_host = self.config.connection_host
+        connection_port = self.config.connection_port
+        vt_host = self.config.amqp_url_virtual_host
+
+        self._print(f'Starting to initialize producer for queue {queue_name}')
+
+        try:
+            generated_url = f"amqp://{username}:{password}@" \
+                            f"{connection_host}:{connection_port}/{vt_host}"
+
+            self._print(f'Trying to connect on {generated_url}')
+            self.connection = pika.BlockingConnection(
+                pika.URLParameters(
+                    generated_url
+                )
+            )
+
+            self.main_channel = self.connection.channel()
+            self._print('Finish initializing producer.')
+        except Exception as e:
+            self._print('Connection Error! Please check RabbitMQ connection')
+            raise e
+
+    def write(self, message: Dict):
+        pass
+
+    def batch_write(self, messages: List[Dict]):
+        if not messages:
+            return
+        self._print(
+            f'Batch ingest {len(messages)} messages. Sample: {messages[0]}'
+        )
+        for message in messages:
+            if isinstance(message, dict):
+                data = message.get('data', message)
+                metadata = message.get('metadata', None)
+            else:
+                data = message
+                metadata = None
+            message_properties = pika.BasicProperties(headers=metadata)
+            self.main_channel.basic_publish(
+                exchange='',
+                routing_key=self.config.queue_name,
+                body=json.dumps(data).encode('utf-8'),
+                properties=message_properties,
+                mandatory=True,
+            )

mage_ai/streaming/sinks/sink_factory.py

@@ -53,6 +53,10 @@ class SinkFactory:
             from mage_ai.streaming.sinks.postgres import PostgresSink
 
             return PostgresSink(config, **kwargs)
+        elif connector_type == SinkType.RABBITMQ:
+            from mage_ai.streaming.sinks.rabbitmq import RabbitMQSink
+
+            return RabbitMQSink(config, **kwargs)
         elif connector_type in GENERIC_IO_SINK_TYPES:
             from mage_ai.streaming.sinks.generic_io import GenericIOSink
 

mage_ai/streaming/sources/activemq.py

@@ -0,0 +1,89 @@
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Callable
+
+import stomp
+from stomp.exception import ConnectFailedException
+
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sources.base import BaseSource
+
+
+@dataclass
+class ActiveMQConfig(BaseConfig):
+    connection_host: str
+    connection_port: int
+    queue_name: str
+    configure_consume: bool = False
+    username: str = 'admin'
+    password: str = 'admin'
+
+
+def messageProcessingFunction(message, handler):
+    print('Recieved message: "%s"' % message)
+    handler([message])
+
+
+class ActiveMQMsgListener(stomp.ConnectionListener):
+    processMessage = None
+    conn = None
+
+    def __init__(self, processMessage, conn, handler):
+        self.processMessage = processMessage
+        self.conn = conn
+        self.handler = handler
+
+    def on_error(self, frame):
+        print('Received an error "%s"' % frame.body)
+
+    def on_message(self, frame):
+        self.processMessage(frame.body, self.handler)
+
+
+class ActiveMQSource(BaseSource):
+    config_class = ActiveMQConfig
+
+    def init_client(self):
+        queue_name = self.config.queue_name
+        username = self.config.username
+        password = self.config.password
+        connection_host = self.config.connection_host
+        connection_port = self.config.connection_port
+
+        self._print(f'Starting to initialize consumer for queue {queue_name}')
+
+        try:
+            conn = stomp.Connection11([(connection_host, connection_port)])
+            self._print('Connecting to broker')
+            conn.connect(username, password, wait=True)
+            self.connection = conn
+
+            self._print('Connected to broker')
+
+        except ConnectFailedException:
+            self._print('Connection Error! Please check broker connection')
+            raise ConnectFailedException
+        except Exception as e:
+            self._print(e)
+            raise e
+
+    def read(self, handler: Callable):
+        pass
+
+    def batch_read(self, handler: Callable):
+        self._print('Start consuming messages.')
+        listener = ActiveMQMsgListener(processMessage=messageProcessingFunction,
+                                       conn=self.connection,
+                                       handler=handler)
+        self.connection.set_listener('ActiceMQListerner', listener)
+        self.connection.subscribe(destination=f'/queue/{self.config.queue_name}',
+                                  id=uuid.uuid4(), headers={})
+        try:
+            while True:
+                time.sleep(10)
+        except KeyboardInterrupt:
+            self._print('Interrupted! Hence, exiting!')
+
+    def __print_message(self, method):
+        self._print(f'Received message {method}')

mage_ai/streaming/sources/nats_js.py

@@ -0,0 +1,182 @@
+import asyncio
+import json
+import ssl
+import threading
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional
+
+import nats
+from nats.errors import NoServersError
+
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS
+from mage_ai.streaming.sources.base import BaseSource
+
+
+@dataclass
+class SSLConfig:
+    cafile: str = None
+    certfile: str = None
+    keyfile: str = None
+    check_hostname: bool = False
+
+
+@dataclass
+class NATSConfig(BaseConfig):
+    server_url: str
+    stream_name: str
+    subject: str = None
+    nkeys_seed_str: Optional[str] = None
+    use_tls: bool = False
+    ssl_config: Optional[SSLConfig] = None
+    consumer_name: Optional[str] = None
+    batch_size: int = DEFAULT_BATCH_SIZE
+    timeout: int = DEFAULT_TIMEOUT_MS / 1000  # Convert to seconds
+
+    @classmethod
+    def parse_config(self, config: Dict) -> Dict:
+        ssl_config = config.get('ssl_config')
+        if ssl_config and type(ssl_config) is dict:
+            config['ssl_config'] = SSLConfig(**ssl_config)
+        return config
+
+
+class NATSSource(BaseSource):
+    config_class = NATSConfig
+
+    def __init__(self, config, **kwargs):
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self.start_loop, daemon=True)
+        self.thread.start()
+        super().__init__(config)
+
+    def start_loop(self):
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_forever()
+
+    def stop_loop(self):
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join()
+
+    async def ainit_client(self):
+        try:
+            connect_opts = {
+                "servers": [self.config.server_url],
+                "error_cb": self.error_cb,
+                "reconnected_cb": self.reconnected_cb,
+                "disconnected_cb": self.disconnected_cb,
+                "closed_cb": self.closed_cb,
+            }
+
+            # Configure SSL context if use_tls is True
+            if self.config.use_tls and self.config.ssl_config:
+                ssl_ctx = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
+                if self.config.ssl_config.cafile:
+                    ssl_ctx.load_verify_locations(self.config.ssl_config.cafile)
+                if self.config.ssl_config.certfile and self.config.ssl_config.keyfile:
+                    ssl_ctx.load_cert_chain(
+                        certfile=self.config.ssl_config.certfile,
+                        keyfile=self.config.ssl_config.keyfile
+                    )
+                connect_opts["tls"] = ssl_ctx
+
+            # Use NKEY if provided
+            if self.config.nkeys_seed_str:
+                connect_opts["nkeys_seed_str"] = self.config.nkeys_seed_str
+
+            # Establish connection with the configured options
+            self.nc = await nats.connect(**connect_opts)
+            self.js = self.nc.jetstream()
+
+            # Default consumer_name to stream_name if not provided
+            consumer_name = self.config.consumer_name or self.config.stream_name
+
+            await self.js.add_stream(name=self.config.stream_name, subjects=[self.config.subject])
+            self.psub = await self.js.pull_subscribe(self.config.subject, consumer_name)
+        except NoServersError as e:
+            self._print(f'Caught NoServersError while connecting to NATS server: {e}')
+        except Exception as e:
+            self._print(f'Caught exception while connecting to NATS server: {e}')
+
+    # Define callback methods
+    async def disconnected_cb(self):
+        self._print('Got disconnected!')
+
+    async def reconnected_cb(self):
+        self._print(f'Got reconnected to {self.nc.connected_url.netloc}')
+
+    async def error_cb(self, e):
+        self._print(f'There was an error: {e}')
+
+    async def closed_cb(self):
+        self._print('Connection is closed')
+
+    async def aclose_client(self):
+        await self.nc.close()
+
+    def init_client(self):
+        future = asyncio.run_coroutine_threadsafe(self.ainit_client(), self.loop)
+        future.result()
+
+    def close_client(self):
+        asyncio.run_coroutine_threadsafe(self.aclose_client(), self.loop)
+
+    def batch_read(self, handler: Callable):
+        self.init_client()
+
+        try:
+            while True:
+                self._print("Fetching messages...")
+                message_tuples = self.fetch_messages()
+                self._print(f"Fetched {len(message_tuples)} messages")
+
+                if not message_tuples:
+                    self._print("No messages fetched, continuing to next iteration")
+                    continue
+
+                processed_messages = []  # List to store successfully processed messages
+
+                for decoded_message, msg in message_tuples:
+                    try:
+                        # Attempt to process the message
+                        processed_messages.append(decoded_message)
+                        asyncio.run_coroutine_threadsafe(msg.ack(), self.loop)
+                    except Exception as e:
+                        self._print(f"Error processing message: {e}")
+
+                # Once all messages in the batch have been attempted,
+                # pass the successfully processed messages to the handler
+                if processed_messages:
+                    handler(processed_messages)
+
+        finally:
+            self.close_client()
+            self.stop_loop()
+
+    def read(self, handler: Callable):
+        self.init_client()
+
+        try:
+            while True:
+                messages = self.fetch_messages()
+                if messages:
+                    for data, msg in messages:
+                        handler(data)  # Process each decoded message
+
+                        # Acknowledge the original message
+                        asyncio.run_coroutine_threadsafe(msg.ack(), self.loop)
+        finally:
+            self.close_client()
+            self.stop_loop()
+
+    async def afetch_messages(self):
+        try:
+            msgs = await self.psub.fetch(self.config.batch_size, timeout=self.config.timeout)
+            # Return a tuple of (decoded_message, message)
+            return [(json.loads(msg.data.decode()), msg) for msg in msgs]
+        except nats.errors.TimeoutError:
+            return []
+
+    def fetch_messages(self):
+        future = asyncio.run_coroutine_threadsafe(self.afetch_messages(), self.loop)
+        return future.result()

mage_ai/streaming/sources/source_factory.py

@@ -29,6 +29,10 @@ class SourceFactory:
             from mage_ai.streaming.sources.kafka import KafkaSource
 
             return KafkaSource(config, **kwargs)
+        elif connector_type == SourceType.NATS:
+            from mage_ai.streaming.sources.nats_js import NATSSource
+
+            return NATSSource(config, **kwargs)
         elif connector_type == SourceType.KINESIS:
             from mage_ai.streaming.sources.kinesis import KinesisSource
 
@@ -37,6 +41,10 @@ class SourceFactory:
             from mage_ai.streaming.sources.rabbitmq import RabbitMQSource
 
             return RabbitMQSource(config, **kwargs)
+        elif connector_type == SourceType.ACTIVEMQ:
+            from mage_ai.streaming.sources.activemq import ActiveMQSource
+
+            return ActiveMQSource(config, **kwargs)
         elif connector_type == SourceType.MONGODB:
             from mage_ai.streaming.sources.mongodb import MongoSource
             return MongoSource(config, **kwargs)